Post-Training is About States, Not Tokens: A State Distribution View of SFT, RL, and On-Policy Distillation

Resumo Técnico: Pós-treinamento como Modelagem da Distribuição de Estados

Declaração do Problema

As análises atuais dos métodos de pós-treinamento de Modelos de Linguagem de Grande Escala (LLMs) — especificamente Ajuste Fino Supervisionado (SFT), Aprendizado por Reforço (RL) e Distilação — focam principalmente nos objetivos de otimização (por exemplo, máxima verossimilhança, gradientes de política, divergência KL). No entanto, essa visão centrada no objetivo tem dificuldade em explicar vários fenômenos empíricos:

• Por que um SFT leve pode melhorar o desempenho com esquecimento mínimo, enquanto um SFT agressivo ("estresse") causa esquecimento catastrófico e degradação do desempenho no alvo.
• Por que o RL, apesar de usar recompensas esparsas ou ruidosas, frequentemente produz melhorias estáveis.
• Por que um modelo estudante na distilação pode, às vezes, superar um professor degradado, mesmo quando o professor é a única fonte de supervisão.

O artigo argumenta que essas discrepâncias surgem porque as análises padrão ignoram um eixo crítico de variação: a distribuição de estados na qual a supervisão é aplicada. Em modelos autoregressivos, um "estado" é definido como o prompt mais o prefixo gerado ($s_t = (x, y_{<t})$). Diferentes métodos de pós-treinamento aplicam seus sinais a distribuições de estados fundamentalmente diferentes (prefixos de conjuntos de dados fixos versus trajetórias induzidas pelo aprendiz), o que impacta significativamente o comportamento do modelo, a localidade e a retenção.

Metodologia

Os autores propõem um quadro unificado que vê o pós-treinamento como modelagem da distribuição de estados. Eles decompõem qualquer método de pós-treinamento em dois eixos distintos:

1. Fonte do Estado ($q_k$): De onde originam os contextos de treinamento (estados).
2. Fonte do Sinal ($S$): Qual supervisão (tokens, recompensas, logits) é fornecida nesses estados.

Formulação Teórica:

• SFT: Ajuste de estado fora de política. O modelo minimiza a perda de tokens em estados fixos do conjunto de dados ($d_{data}$). Isso cria uma incompatibilidade se as trajetórias do conjunto de dados divergirem das próprias trajetórias do modelo, levando a viés de exposição e instabilidade potencial sob estresse.
• RL: Melhoria local em política. O modelo amostra trajetórias de sua política atual ($d_{\pi_\theta}$) e atualiza com base em recompensas. As atualizações são naturalmente localizadas para estados que o modelo realmente visita.
• Distilação em Política (OPD): Um híbrido onde o estudante amostra estados ($d_{\pi_S}$), mas o professor fornece o sinal (logits ou continuação). Isso desacopla a fonte do estado da fonte do sinal, permitindo que o estudante aprenda reparos locais sem herdar a distribuição global de trajetórias do professor.

Configuração Experimental:

• Modelo: Qwen3-0.6B-Base com adaptadores LoRA.
• Hardware: Uma única GPU RTX 3090.
• Tarefas:
  • Alvo: GSM8K (raciocínio matemático).
  • Retenção: TruthfulQA e MMLU (para medir esquecimento catastrófico).
• Métricas: Precisão no alvo, taxa de retenção, magnitude do esquecimento e deriva de estado (medida via Discrepância Média Máxima, MMD, em características léxicas dos estados de trajetória).

Condições Experimentais:

1. SFT Leve: Treinamento padrão no GSM8K.
2. SFT Estressado: Treinamento agressivo (5 épocas, alta taxa de aprendizado, alto rank LoRA) para induzir degradação.
3. OPD: Treinamento de um estudante usando um professor "SFT Estressado" degradado como fonte de sinal, mas amostrando estados da própria política do estudante.
4. RL em Política: RL leve usando recompensas de resposta exata relativas a grupos no GSM8K.

Resultados Principais

1. O SFT é Dependente do Contexto: O SFT leve melhorou o GSM8K (0,448 $\to$ 0,512) com perda de retenção negligenciável. Por outro lado, o SFT Estressado degradou tanto a tarefa alvo (0,420) quanto as tarefas de retenção (TruthfulQA/MMLU caíram significativamente), demonstrando que o SFT não é inerentemente destrutivo, mas torna-se sob pressão fora de política que não se alinha com a distribuição natural de estados do modelo.
2. OPD Supera Professores Degradados: Um estudante treinado via OPD usando o modelo "SFT Estressado" como professor superou o professor em todas as métricas (GSM8K: 0,466 vs. 0,420; MMLU: 0,430 vs. 0,364). Isso confirma que os estudantes não herdam as falhas globais de trajetória do professor se forem treinados em seus próprios estados, recebendo apenas orientação condicional local.
3. RL Preserva Localidade: O RL leve em política melhorou o GSM8K (0,472) enquanto mantinha níveis de retenção próximos à base. Isso apoia a visão de que o RL atua como um procedimento de melhoria local fundamentado nos estados alcançáveis pelo aprendiz.
4. Métricas de Deriva são Insuficientes: Métodos com deriva escalar de estado quase idêntica (MMD) exibiram resultados de retenção drasticamente diferentes. Por exemplo, SFT Estressado e OPD a partir do professor Estressado tiveram valores de MMD similares (~0,0109), mas taxas de retenção muito diferentes (0,82 vs. 0,95). Isso indica que a fonte e a localidade dos estados de treinamento importam mais do que a magnitude da mudança distribucional isoladamente.

Contribuições

1. Quadro Unificado: Formulou SFT, RL e OPD sob uma visão comum de "distribuição de estados", separando explicitamente a fonte do estado da fonte do sinal.
2. Pipeline Empírico Controlado: Implementou uma configuração experimental de GPU única medindo precisão no alvo, retenção, esquecimento e deriva de estado de trajetória para isolar os efeitos da distribuição de estados.
3. Evidência para Afirmações Centradas no Estado: Forneceu evidências empíricas de que:
   • O OPD pode recuperar o desempenho e superar um professor degradado ao desacoplar a amostragem de estados da geração de sinais.
   • O RL em política melhora o desempenho no alvo com perda mínima de retenção devido à sua localidade.
   • Métricas de deriva escalar (como MMD) não podem caracterizar totalmente a dinâmica de pós-treinamento sem considerar a fonte do estado.

Significado e Afirmações

O artigo afirma que o comportamento de pós-treinamento é governado pela interação entre o sinal de supervisão e a distribuição de estados de treinamento. Os autores argumentam que a "fonte, localidade e dependência do aprendiz dos estados de treinamento" são tão críticas quanto a forma do próprio sinal de supervisão.

• Reenquadramento do SFT: O esquecimento catastrófico não é apenas uma função da máxima verossimilhança, mas surge quando atualizações densas fora de política empurram o modelo para regiões incompatíveis com sua própria distribuição futura de estados.
• Reenquadramento da Distilação: Um estudante pode superar um professor porque os erros do professor estão frequentemente acoplados à própria distribuição de trajetória do professor. Ao consultar o professor nos estados do estudante, o estudante acessa reparos locais úteis sem herdar falhas globais de trajetória.
• Reenquadramento do RL: A estabilidade do RL é parcialmente atribuída à sua natureza em política, garantindo que as atualizações sejam aplicadas apenas a estados que o modelo pode realmente alcançar e dos quais pode se recuperar.

O artigo conclui que uma visão centrada no estado é necessária para explicar a estabilidade e reversões de desempenho que a análise apenas no nível do objetivo não consegue justificar. Os autores enfatizam que essas descobertas baseiam-se em um estudo de pequena escala e devem ser vistas como evidência para uma hipótese mecânica, e não como uma afirmação definitiva de benchmark.