Towards On-Policy SFT: Distribution Discriminant Theory and its Applications in LLM Training

Resumo Técnico: Rumo ao SFT On-Policy: Teoria Discriminante de Distribuição e suas Aplicações no Treinamento de LLMs

1. O Problema

O artigo aborda a lacuna de desempenho entre o Ajuste Fino Supervisionado (SFT) e o Aprendizado por Reforço (RL) na pós-treinamento de Grandes Modelos de Linguagem (LLMs).

• SFT: É computacionalmente eficiente, mas frequentemente resulta em generalização inferior e esquecimento catastrófico (catastrophic forgetting) em comparação ao RL. Isso ocorre porque o SFT padrão força o modelo a ajustar-se igualmente a todos os dados externos, ignorando a distribuição nativa do modelo pré-treinado.
• RL: Oferece melhor generalização ao utilizar dados on-policy (gerados pelo próprio modelo), preservando sua distribuição nativa. No entanto, o RL sofre de ineficiência de dados (sinais de recompensa esparsos), alta complexidade computacional e dificuldade de aplicação em cenários onde a verificação de feedback é difícil ou enviesada.
• Questão Central: É possível fechar a lacuna entre SFT e RL, alinhando o processo de treinamento com a distribuição interna do modelo (tanto no nível de dados quanto no nível de objetivo), mantendo a eficiência do SFT e preservando o conhecimento pré-treinado?

2. Metodologia e Fundamentos Teóricos

Os autores propõem um framework baseado na Teoria Discriminante de Distribuição (DDT - Distribution Discriminant Theory) para quantificar o alinhamento entre os dados e a distribuição induzida pelo modelo.

A. Teoria Discriminante de Distribuição (DDT)

• Definição do Problema: A detecção de alinhamento de distribuição é tratada como um problema de teste de hipóteses estatísticas.
  • $H_0$ (In-distribution): O token é amostrado da distribuição do próprio modelo ($p_t$).
  • $H_1$ (Out-of-distribution): O token vem de um mecanismo externo desconhecido ($q_t$).
• Critério Ótimo (CLL): Utilizando a Teoria de Detecção de Sinal (Signal Detection Theory), os autores provam que o Log-Likelihood Centralizado (CLL) é o critério estatístico ótimo para discriminar entre $H_0$ e $H_1$.
  • Fórmula: $\phi_t = \log p_t(x_t) + H[p_t]$, onde $H[p_t]$ é a entropia da distribuição condicional.
  • Justificativa: O CLL maximiza a Razão Sinal-Ruído (SNR) ao subtrair a variabilidade da dificuldade do contexto (entropia), focando apenas na discrepância intrínseca do token. Isso supera o uso simples de log-likelihood (LL), que é sensível à complexidade do contexto.
• Propriedades Sequenciais: Sob $H_0$, a soma cumulativa de $\phi_t$ comporta-se como uma martingale de média zero. Sob $H_1$, exibe um desvio linear negativo proporcional à Divergência KL.

B. Aplicações Práticas
Com base na DDT, são introduzidas duas técnicas complementares:

1. In-Distribution Finetuning (IDFT) - Nível de Perda:

   • Objetivo: Mitigar o esquecimento catastrófico e a superajuste (overfitting) em dados ruidosos ou fora de distribuição (OOD).
   • Mecanismo: Introduz um coeficiente de modulação adaptativa $\gamma_t = \exp(-\phi_t)$ na função de perda do SFT.
     • Se $\phi_t \ll 0$ (token OOD/Ruído): O gradiente é suprimido ($\gamma_t > 1$), protegendo a estrutura pré-treinada.
     • Se $\phi_t \approx 0$ (In-distribution): O aprendizado padrão é mantido.
     • Se $\phi_t > 0$ (Fortemente in-distribution): O aprendizado é reforçado para consolidar conhecimento.
   • Vantagem: Diferente de métodos anteriores que usam apenas probabilidade bruta ou entropia absoluta, o IDFT usa o estatístico ótimo $\phi_t$ para distinguir entre dificuldade intrínseca e desalinhamento de distribuição.

2. Hinted Decoding (Decodificação Sugerida) - Nível de Dados:

   • Objetivo: Re-alinhar um corpus de treinamento fixo para a distribuição do modelo, gerando respostas que são corretas e esteticamente consistentes com o modelo.
   • Mecanismo: Um algoritmo de decodificação em tempo de inferência que combina duas distribuições:
     • $p_I$: Distribuição de um "imitador" (baseado em self-distillation com a resposta correta visível).
     • $p_m$: Distribuição original do modelo (sem ver a resposta).
   • Fusão Adaptativa: A distribuição final é uma mistura geométrica ponderada pela entropia do imitador ($H_I$).
     • Em tokens de alta incerteza (alta entropia), o modelo original domina para preservar o estilo.
     • Em tokens críticos (baixa entropia), o imitador domina para garantir a correção lógica.
   • Detecção de Falsos Positivos: Um mecanismo monitora delimitadores (ex: < /think >) para garantir que a cadeia de raciocínio (CoT) e a resposta final sejam consistentes, descartando amostras onde a lógica falha.

3. Resultados Experimentais

Os autores realizaram extensos experimentos em modelos de ponta (Qwen2.5, DeepSeek-R1, Mistral, LLaMA) em tarefas de raciocínio matemático e código.

• Desempenho vs. SFT Aprimorado: O IDFT superou consistentemente métodos existentes como DFT (Dynamic Fine-Tuning) e EAFT (Entropy-Adaptive Fine-Tuning) em benchmarks matemáticos (AIME, AMC, Olympiad) e de raciocínio geral (MMLU, ARC).
• Desempenho vs. RL Offline: A combinação de Hinted Decoding + IDFT superou algoritmos de RL offline populares (DPO, SimPO, Rejection Sampling) em todos os benchmarks testados.
  • Eficiência: A abordagem proposta atingiu desempenho superior com um custo computacional (GPU hours) significativamente menor que o RL.
  • Robustez: O método demonstrou ser eficaz mesmo em modelos instruídos ("instruct models") e modelos de raciocínio ("thinking models"), onde o gap de distribuição é maior e o SFT padrão falha.
• Mitigação de Esquecimento: O framework manteve ou melhorou o desempenho em benchmarks gerais, indicando uma mitigação eficaz do esquecimento catastrófico, algo que o SFT padrão e o RL muitas vezes têm dificuldade em equilibrar simultaneamente.

4. Contribuições Chave

1. Teoria (DDT): Estabelecimento rigoroso da Teoria Discriminante de Distribuição, provando matematicamente que o Log-Likelihood Centralizado (CLL) é o critério ótimo para discriminar dados on-policy de dados off-policy em LLMs.
2. Algoritmo (IDFT): Proposta de uma nova função de perda que modula dinamicamente os gradientes com base na distribuição, protegendo o conhecimento pré-treinado sem sacrificar a capacidade de aprendizado.
3. Técnica de Dados (Hinted Decoding): Desenvolvimento de um método de decodificação que transforma dados estáticos em amostras on-policy corretas, permitindo o treinamento eficiente sem a necessidade de RL complexo.
4. Validação Empírica: Demonstração de que é possível superar o RL offline em tarefas de raciocínio complexo usando apenas SFT aprimorado, oferecendo uma alternativa prática para domínios onde o RL é inviável.

5. Significado e Impacto

Este trabalho oferece uma alternativa prática e eficiente para o pós-treinamento de LLMs em cenários onde o RL é difícil de aplicar (devido à falta de feedback verificável ou custos computacionais). Ao formalizar teoricamente o conceito de "dados in-distribution" e criar ferramentas para alinhar dados e objetivos a essa distribuição, os autores mostram que é possível obter a generalização do RL mantendo a eficiência do SFT. Isso tem implicações profundas para o treinamento de modelos em áreas como matemática, código e agentes autônomos, onde a precisão e a estabilidade são críticas.

O código e os dados foram disponibilizados publicamente, incentivando a reprodutibilidade e pesquisas futuras na interseção entre SFT e RL.