Um artigo publicado no site de Machine Learning Research da Apple examina quanta melhoria de recompensa o alinhamento de modelos de linguagem pode alcançar teoricamente sob um orçamento fixo de divergência. O trabalho, intitulado "Theoretical Limits of Language Model Alignment", é assinado por Lucas Monteiro Paes, Natalie Mackraz, Barry-John Theobald e Federico Danieli, com Paes e Mackraz creditados como contribuintes em igualdade de condições.
Os métodos de alinhamento buscam moldar as saídas do modelo para que reflitam as preferências humanas, preservando ao mesmo tempo as capacidades do modelo base. Os autores observam que duas famílias de técnicas dominam a prática: o aprendizado por reforço, que maximiza a recompensa esperada sujeita a uma restrição de divergência KL, e o alinhamento best-of-N, que seleciona a saída com maior recompensa entre N amostras extraídas de forma independente. Apesar do uso difundido de ambas, o teto fundamental de melhoria de recompensa sob um orçamento de KL ainda era pouco compreendido.
O artigo caracteriza os limites teóricos da informação do alinhamento regularizado por KL ao derivar o maior ganho de recompensa esperada possível para um dado orçamento de divergência KL. Seu primeiro resultado é uma expressão em forma fechada para a melhoria ótima, governada por um termo de divergência de Jeffreys, em vez da raiz quadrada da divergência KL usada em análises anteriores. Os pesquisadores também reformulam essa expressão como uma covariância sob o modelo base, o que gera um estimador prático capaz de prever os ganhos de alinhamento alcançáveis apenas a partir de amostras do modelo base.
A análise é então estendida ao cenário de recompensas proxy. Os autores mostram que a lacuna entre o alinhamento ideal e o proxy, comumente descrita como reward hacking, aumenta com a magnitude do erro de recompensa e quando o fator de penalização KL diminui. Eles também provam que o ensemble de recompensas mitiga o reward hacking, oferecendo uma justificativa teórica para uma técnica já empregada na prática.
No lado empírico, a equipe calculou a fronteira de Pareto entre KL e recompensa para duas tarefas de modelos de linguagem, segurança e sumarização. Verificou-se que a amostragem best-of-N se aproxima bastante do limite teórico, enquanto PPO e GRPO permanecem substancialmente abaixo do ótimo. Segundo os autores, os resultados ajudam a explicar diversos fenômenos observados empiricamente na literatura de alinhamento e sugerem que melhorias algorítmicas serão necessárias para alcançar um alinhamento ótimo sem altos custos de inferência.
O artigo foi aceito no Fine-Tuning in Modern Machine Learning (FITML) Workshop da NeurIPS 2024.