Resumo Técnico: Treinamento Conjunto de Predição de Múltiplos Tokens em Aprendizado por Reforço via Calibração Ótima de Coeficientes

Declaração do Problema

O Aprendizado por Reforço a partir de Recompensas Verificáveis (RLVR) tornou-se o paradigma padrão para aprimorar as capacidades de raciocínio de Modelos de Linguagem de Grande Escala (LLMs). Concurrentemente, a Predição de Múltiplos Tokens (MTP), um objetivo auxiliar treinado durante o pré-treinamento para prever múltiplos tokens futuros, demonstrou benefícios significativos na eficiência de amostragem e como um modelo de rascunho para decodificação especulativa.

Intuitivamente, a otimização conjunta de MTP e objetivos de RL durante o pós-treinamento deveria permitir que os modelos se beneficiassem da supervisão de múltiplos tokens enquanto adquiriam capacidades de raciocínio. No entanto, evidências empíricas mostram que o treinamento conjunto MTP-RL frequentemente leva a uma degradação significativa no desempenho do modelo principal. Consequentemente, as práticas atuais (por exemplo, em frameworks como veRL e slime) tipicamente desacoplam os gradientes de MTP do modelo principal durante o treinamento de RL. Esse desacoplamento força os parâmetros de MTP a se adaptarem apenas através de ajuste fino pós-hoc independente, sacrificando os ganhos potenciais da supervisão simultânea de múltiplos tokens. O problema central abordado é: Por que o treinamento conjunto MTP-RL falha e pode ser projetada uma estratégia principiante para permitir um treinamento conjunto eficaz?

Metodologia

Análise Teórica

Os autores abordam o problema a partir de uma perspectiva de otimização, analisando o efeito por etapa de MTP sobre o objetivo de RL $J(\theta)$. Assumindo uma função objetivo $L$-suave, eles derivam um limite inferior para a melhoria por etapa. Essa decomposição revela que o efeito de MTP é governado por dois termos concorrentes:

1. Termo de Correlação de Primeira Ordem: Determinado pelo acordo direcional (produto interno) entre o gradiente de RL ($g_{RL}$) e o gradiente de MTP ($g_{MTP}$).
2. Penalidade de Perturbação de Segunda Ordem: Um termo não positivo que reflete a perturbação introduzida pela magnitude do gradiente de MTP.

A melhoria no desempenho ocorre apenas se o termo de correlação superar o termo de penalidade.

Análise de Regimes Existentes

Utilizando essa decomposição, o artigo unifica e explica três regimes de treinamento existentes:

• Desacoplamento (Detach): $g_{MTP} = 0$ em relação ao modelo principal. $\Delta_{MTP} = 0$. Isso é seguro, mas renuncia a ganhos potenciais.
• Perda de Entropia Cruzada (CE Loss): O termo de correlação é próximo de zero em expectativa, pois o RL aumenta o peso de amostras de alta recompensa, enquanto o CE trata todas as amostras igualmente. Consequentemente, a penalidade negativa de segunda ordem domina, causando degradação.
• Perda de Política (Policy Loss): Utiliza a mesma formulação de objetivo de RL para MTP. Inicialmente, os gradientes estão alinhados e o termo de correlação domina, gerando ganhos. No entanto, à medida que o treinamento progride e o modelo entra em regiões planas da paisagem de perda, a correlação decai enquanto a norma do gradiente (e, portanto, a penalidade) persiste. Isso leva a uma curva de desempenho de "subida-então-descida", onde um coeficiente fixo não consegue acompanhar a transição de fase.

Solução Proposta: Calibração Ótima de Coeficientes (OCC)

Para abordar a transição de fase, os autores propõem a Calibração Ótima de Coeficientes (OCC), um esquema adaptativo que rastreia o coeficiente ótimo de MTP ($\lambda$) online.

1. Ótimo em Forma Fechada: A análise produz uma expressão em forma fechada para o $\lambda^*$ ótimo em cada etapa, que é proporcional à razão entre a correlação de gradientes ($c$) e a norma quadrada do gradiente ($v^2$).
2. Proxy de Log-Probabilidade: Calcular gradientes de modelo completo para $c$ e $v^2$ é proibitivamente caro em sistemas distribuídos em grande escala. A OCC utiliza um proxy de log-probabilidade derivado da mudança na log-probabilidade sob atualizações de pequeno passo ($\delta \approx \log \pi_\theta - \log \pi_{old}$).
   • Proxy de Alinhamento ($\hat{c}_t$): Estima o produto interno.
   • Proxy de Variância ($\hat{v}^2_t$): Estima a norma quadrada.
3. Atualização Adaptativa: O coeficiente é atualizado como $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$. Isso permite que a OCC ajuste dinamicamente o peso de MTP, mantendo ganhos positivos durante todo o treinamento sem calcular gradientes completos.

Contribuições Principais

1. Unificação Teórica: O artigo fornece uma análise teórica que unifica os regimes de Desacoplamento, Entropia Cruzada e Perda de Política, explicando seu sucesso ou fracasso com base na interação entre a correlação de primeira ordem e a perturbação de segunda ordem.
2. Algoritmo OCC: Propõe a Calibração Ótima de Coeficientes, um procedimento adaptativo que calibra o coeficiente de MTP online usando um proxy de log-probabilidade. Este método requer sobrecarga computacional negligenciável e evita o cálculo de gradientes de modelo completo.
3. Validação Empírica: Experimentos extensivos em múltiplos modelos (MiMo-7B, GLM-4.5-Air) e algoritmos (DAPO, GSPO) em benchmarks de raciocínio matemático de nível de competição demonstram que a OCC consistentemente iguala ou supera a linha de base de Desacoplamento.

Resultados Experimentais

Os autores avaliaram a OCC em seis benchmarks: AIME24, AIME25, AMC, MATH, Minerva e OlympiadBench.

• Desempenho: Em todos os benchmarks e escalas de modelo (7B denso e 106B MoE), a OCC superou consistentemente a linha de base de Desacoplamento e a Perda de Política de coeficiente fixo.
  • No MiMo-7B com DAPO, a OCC alcançou uma precisão média de 61,7%, superando o Desacoplamento (58,9%) e a Perda de Política (57,4%).
  • No GLM-4.5-Air, a OCC alcançou 67,6% de precisão média versus 65,9% para o Desacoplamento.
• Dinâmica de Treinamento: Enquanto a Perda de Política mostrou uma tendência de subida-então-descida (degradando nas etapas finais), a OCC manteve melhorias estáveis durante todo o treinamento. A Perda CE consistentemente teve desempenho inferior ao Desacoplamento.
• Eficiência: A OCC introduziu sobrecarga negligenciável. O tempo de treinamento por etapa (wall-clock) para a OCC (8,07s) foi comparável ao do Desacoplamento (8,11s), enquanto a retropropagação de modelo completo exigiu 45,06s (um desaceleração de 5,6×).
• Ablação: A análise de sensibilidade mostrou que o fator de escala pré-definido $\lambda_+$ é robusto em uma faixa de valores, com $\lambda_+ = 1,0$ produzindo o melhor desempenho médio.

Significado e Alegações

O artigo alega que a falha do treinamento conjunto MTP-RL não é inerente, mas decorre de uma incompatibilidade entre coeficientes estáticos e a deriva dinâmica do alinhamento de gradientes durante o treinamento. Ao fornecer uma decomposição teórica da paisagem de otimização e propor a OCC, os autores demonstram que o MTP pode ser reintroduzido com segurança no pós-treinamento de RL.

O significado reside em permitir que os modelos adquiram simultaneamente capacidades de raciocínio via RL e aproveitem representações de múltiplos tokens sem sacrificar o desempenho ou incorrer em custos computacionais proibitivos. O trabalho sugere que a calibração adaptativa é essencial para o treinamento conjunto em paisagens de aprendizado profundo não convexas, oferecendo uma alternativa principiante ao padrão industrial atual de desacoplamento de gradientes.

Limitações

Os autores observam que sua análise teórica depende de suposições de $L$-suavidade que podem não valer exatamente em paisagens não convexas, fornecendo orientação qualitativa em vez de quantitativa. O proxy de log-probabilidade é uma aproximação de primeira ordem que pode perder fidelidade sob taxas de aprendizado grandes. Além disso, os experimentos focam em raciocínio matemático com recompensas verificáveis; a generalização para configurações como RLHF com modelos de recompensa aprendidos ou geração aberta permanece a ser validada.