Resumo Técnico: Desacoplamento de KL e Trajetórias em Distilação de LLM

Declaração do Problema

A distilação de conhecimento tornou-se uma pedra angular do pós-treinamento de Modelos de Linguagem de Grande Escala (LLM), particularmente para tarefas de raciocínio, operando frequentemente em conjunto com Aprendizado por Reforço (RL). No entanto, o espaço de design da distilação permanece pouco compreendido. Os pipelines atuais tipicamente seguem um de dois paradigmas: distilação off-policy (treinamento em traços gerados pelo professor, frequentemente via Ajuste Fino Supervisionado) ou distilação on-policy (OPD) (treinamento em rollouts gerados pelo aluno com feedback do professor).

O artigo argumenta que esses paradigmas predominantes acoplam implicitamente duas escolhas de design ortogonais:

1. Fonte do Prefixo: Se a supervisão é aplicada a prefixos gerados pelo professor (off-policy) ou pelo aluno (on-policy).
2. Direção do KL ao Nível de Token: Se o objetivo minimiza o KL direto ($KL(p_{teacher} \parallel q_{student})$) ou o KL reverso ($KL(q_{student} \parallel p_{teacher})$).

A prática convencional acopla prefixos do professor com KL direto e prefixos do aluno com KL reverso. Os autores contendem que esse acoplamento não é intrínseco ao nível de token e que desacoplar esses eixos revela um espaço de design mais amplo com quatro objetivos válidos, cada um correspondendo a regimes de aprendizado distintos.

Metodologia

Estrutura Teórica: Desacoplamento dos Eixos

Os autores decompõem a divergência KL ao nível de sequência usando a regra da cadeia autoregressiva. Eles demonstram que:

• KL Direto ($KL(p_T \parallel q_S)$) acopla naturalmente com prefixos do professor (off-policy) porque a expectativa é tomada sobre a distribuição de trajetória do professor.
• KL Reverso ($KL(q_S \parallel p_T)$) acopla naturalmente com prefixos do aluno (on-policy) porque a expectativa é tomada sobre a distribuição de trajetória do aluno.

Ao tratar esses como eixos independentes, os autores definem quatro objetivos:

1. Prefixo do Professor + KL Direto: Corresponde ao SFT off-policy (distilação padrão).
2. Prefixo do Aluno + KL Direto: Corresponde ao SFT on-policy estilo DAgger.
3. Prefixo do Professor + KL Reverso: Corresponde a um objetivo de distilação estilo RL offline.
4. Prefixo do Aluno + KL Reverso: Corresponde à Distilação On-Policy (OPD), vista como RL on-policy com recompensa densa.

O artigo estabelece identidades ao nível de gradiente mostrando que:

• KL Direto produz um gradiente de entropia cruzada estilo SFT com alvos suaves do professor.
• KL Reverso produz um gradiente de política estilo RL com uma recompensa densa definida pela razão logarítmica professor-aluno ($r = \log p_T - \log q_S$).

Configuração Experimental

Os autores conduzem um estudo controlado sobre raciocínio matemático usando modelos Qwen3 (aluno de 0.6B, professores de 4B/8B). Eles avaliam os quatro objetivos em dois regimes:

1. Distilação Autônoma: Comparando o desempenho final de modelos treinados exclusivamente com cada objetivo.
2. Distilação-seguida-de-RL: Usando checkpoints distilados como inicializações para Otimização de Política Relativa em Grupo (GRPO).

Eles acompanham precisão, entropia preditiva e comprimento de resposta através de vários comprimentos de sequência (128 vs. 4096 tokens) e benchmarks (MATH500, GSM8K, AIME24, AMC23).

Principais Contribuições

1. Identificação de Três Compensações

A análise empírica revela três compensações fundamentais que governam o design da distilação:

• Compensação Precisão–Entropia (Direção do KL): O KL reverso consistentemente alcança maior precisão autônoma (Avg@k), mas leva a entropia preditiva ao colapso, reduzindo a diversidade (menor Pass@k) e tornando o modelo menos estável para RL subsequente. O KL direto preserva a entropia, suportando melhoria de RL mais estável, embora com menor precisão autônoma inicial.
• Compensação Qualidade–Cálculo (Fonte do Prefixo): Prefixos do aluno (on-policy) produzem maior precisão sob passos de treinamento correspondentes ao supervisionar estados que o aluno realmente visita. No entanto, prefixos do professor (off-policy) são mais eficientes em cálculo sob FLOPs correspondentes porque permitem reutilizar logits do professor em cache e trajetórias offline, evitando o custo da geração online do aluno.
• Compensação Precisão–Estabilidade (Duração do Treinamento): A distilação de sequência longa (4096 tokens) melhora a precisão de raciocínio, mas, particularmente sob KL reverso, causa colapso severo de entropia e "inflação de comprimento" (respostas atingindo o limite de geração). O treinamento de sequência curta é mais estável, mas menos preciso.

2. Métodos Propostos para Equilibrar Compensações

Motivados por essas descobertas, os autores propõem dois métodos práticos:

• Mistura de KL: Uma combinação ponderada de KL direto e reverso em cada prefixo. Os autores descobrem que, na distilação de sequência longa, uma mistura pesada em direção ao KL direto é crucial. Ela retém os benefícios de precisão do KL reverso enquanto previne o colapso de entropia e a inflação de comprimento, estabilizando a dinâmica de treinamento.
• Currículo de Comprimento Acionado por Entropia: Um cronograma de treinamento que começa com horizontes curtos e aumenta o comprimento da sequência apenas enquanto a entropia preditiva permanece acima de um limiar de estabilidade. Essa abordagem evita o colapso de entropia associado ao treinamento de horizonte longo fixo.

Resultados

Distilação Autônoma

• KL Reverso geralmente supera o KL Direto em precisão autônoma (Avg@k), com ganhos de +2,45 pontos em média.
• Prefixos do Aluno superam Prefixos do Professor em precisão, especialmente em configurações de sequência longa (+3,55 ganho em Avg@k).
• Colapso de Entropia: KL reverso puro em configurações de sequência longa leva a entropia próxima de zero e causa inflação dos comprimentos de resposta até o limite máximo (ex: 8192 tokens), degradando o Pass@k.

Distilação-seguida-de-RL (GRPO)

• Estabilidade de Inicialização: Modelos inicializados com KL Direto (mesmo que com menor precisão autônoma) mostram melhoria mais confiável durante o GRPO porque retêm maior entropia.
• Armadilha do KL Reverso: Modelos inicializados com KL Reverso frequentemente começam com maior precisão, mas degradam ou estagnam durante o RL devido à baixa entropia e capacidade reduzida de exploração.
• Sucesso do Currículo: O currículo de comprimento acionado por entropia melhora significativamente o desempenho. Comparado ao treinamento fixo de 4096 tokens, ele melhorou Avg@k em 3,6 pontos e Pass@k em até 5,8 pontos, enquanto reduziu o comprimento médio de resposta em aproximadamente 3× (ex: de ~8185 para ~2688 tokens).

Resultados da Mistura de KL

Em configurações de sequência longa, uma mistura com alto peso de KL direto (ex: 0,8 Direto / 0,2 Reverso) preservou a maior parte da precisão do KL reverso enquanto estabilizava a entropia e prevenia a inflação de comprimento.

Significado e Alegações

O artigo alega fornecer uma estrutura teórica unificada que conecta receitas de treinamento diversas (SFT, DAgger, RL Offline, OPD) como instâncias específicas de duas escolhas de design ortogonais. Ao desacoplar a fonte do prefixo e a direção do KL, os autores argumentam que os praticantes podem ir além de paradigmas rígidos para projetar objetivos de distilação de raciocínio que equilibrem explicitamente:

• Precisão (via KL reverso e sequências longas),
• Diversidade/Exploração (via KL direto e preservação de entropia),
• Eficiência Computacional (via prefixos do professor/cache), e
• Treinabilidade de RL a jusante (via entropia estável).

Os métodos propostos (mistura de KL e currículo acionado por entropia) são apresentados como ferramentas práticas para navegar nas compensações identificadas, permitindo o treinamento de modelos de raciocínio menores que são tanto precisos quanto estáveis para otimização de política adicional. Os autores enfatizam que o "melhor" objetivo autônomo não é necessariamente a melhor inicialização para um pipeline de RL, destacando a necessidade de uma visão holística do processo de pós-treinamento.