Escaping the KL Agreement Trap in On-Policy Distillation

Resumo Técnico: Escapando da Armadilha de Acordo de KL na Destilação On-Policy

1. Definição do Problema

A destilação on-policy (OPD) surgiu como um método poderoso para transferir capacidades de grandes modelos professores para modelos estudantes menores, fornecendo supervisão densa ao nível de tokens em rollouts gerados pelo estudante. No entanto, este artigo identifica um modo de falha crítico denominado armadilha de baixo acordo de KL (low-KL agreement trap).

Na OPD padrão, uma baixa divergência de Kullback-Leibler (KL) reversa entre o estudante e o professor é tipicamente interpretada como um sinal de progresso ou alinhamento. Os autores observam que, conforme um rollout do estudante progride, ele pode derivar para um "prefixo corrompido" (por exemplo, um erro de raciocínio precoce ou um padrão de degeneração repetitiva). Condicionado a esse estado degradado, o modelo professor pode produzir continuações que são localmente plausíveis, mas não informativas, resultando em uma baixa divergência KL. Isso cria um acordo degenerado: o professor e o estudante concordam, mas o sinal é fraco porque o professor está apenas seguindo o erro do estudante em vez de fornecer orientação corretiva.

O artigo argumenta que os métodos de OPD existentes tratam todo o rollout como supervisão útil. Consequentemente, eles desperdiçam computação em tokens que seguem a entrada na armadilha, onde o sinal de supervisão é diluído ou prejudicial, potencialmente desviando o processo de otimização.

2. Metodologia: KAT (Terminação de Armadilha de Acordo de KL)

Para enfrentar isso, os autores propõem o KAT, uma regra de terminação de rollout online projetada para detectar e filtrar regiões de acordo degenerado sem alterar o objetivo de OPD subjacente ou exigir modelos auxiliares.

Mecanismo Central:
O KAT reutiliza a divergência KL reversa ($d_t$) já computada durante o processo de OPD. Em vez de aplicar um limiar a um único token de KL, o KAT monitora uma estatística de janela deslizante ($z_t$) da média de KL sobre uma janela de $W$ tokens.

Limiar Adaptativo:
Reconhecendo que a escala da divergência KL muda ao longo do treinamento, o KAT emprega um limiar dinâmico e baseado em dados ($s_k$):

• Fase de Warmup: Durante os $K$ passos iniciais do otimizador, o KAT coleta os valores mínimos de KL da janela deslizante ($m(\hat{y})$) de rollouts concluídos em um buffer FIFO.
• Calibração: Após o warmup, o limiar de parada é definido como o quantil $(1-\eta)$ dos mínimos recentes do buffer. Isso garante que a regra se adapte à política evolutiva do estudante e à escala de KL.

Lógica de Terminação:
Um rollout é terminado online quando a estatística de KL da janela deslizante permanece abaixo do limiar adaptativo por $T$ janelas consecutivas. Crucialmente, o KAT inclui um comprimento de isenção ($L_0$) para evitar a terminação prematura durante a fase inicial de um rollout, onde o KL naturalmente cai. Uma vez confirmado o gatilho, o rollout é truncado no início da primeira janela de baixo KL, e nenhum token adicional é gerado ou usado para supervisão.

3. Principais Contribuições

O artigo apresenta três contribuições primárias:

1. Identificação de Armadilhas de Baixo Acordo de KL: Os autores caracterizam um modo de falha específico em OPD onde o baixo KL persistente indica um "acordo degenerado" sobre prefixos corrompidos, em vez de um alinhamento benéfico.
2. Análise da Qualidade da Supervisão: Através de análise de geometria de gradiente (projetando gradientes em subespaços principais de atualização) e visualização de vocabulário, o artigo demonstra que:
   • Os gradientes nas fases de "acordo" e "pós-acordo" alinham-se mal com as direções principais de atualização de parâmetros.
   • O vocabulário do professor muda de tokens relacionados ao raciocínio para fragmentos superficiais locais ou continuações genéricas nessas regiões.
   • A qualidade da supervisão não se recupera mesmo se o KL aumentar ligeiramente após a armadilha.
3. Proposta do KAT: Uma regra de terminação online e adaptativa que filtra sufixos não informativos. Ela melhora a eficiência e a precisão do treinamento ao focar as atualizações em prefixos informativos pré-armadilha, exigindo apenas $O(1)$ de contabilidade por token.

4. Resultados Experimentais

Os autores avaliaram o KAT em quatro benchmarks de raciocínio matemático (AMC, MATH500, MinervaMath, AIME24) usando duas escalas de estudante (Qwen3-0.7B/1.6B e Qwen3-4B) com um professor Qwen3-8B.

Ganhos de Desempenho:

• Acurácia: O KAT melhorou a acurácia média avg@k em 2,66% e pass@k em 3,43% em comparação com a OPD padrão.
• Eficiência: O comprimento médio do rollout foi reduzido em 59,73% (de ~1480 para ~596 tokens).
• Custo de Treinamento: No estudante 4B, o KAT reduziu o tempo de treinamento em 2,4x (26,28h para 11,00h) e o uso de computação em 59,4% em comparação com a OPD padrão.

Análise Comparativa:

• vs. Prefixo Fixo: O KAT superou a truncagem de comprimento fixo (ex: 1024 ou 2048 tokens), demonstrando que a qualidade do ponto de truncagem importa mais do que o comprimento em si. Prefixos fixos frequentemente cortam raciocínios úteis ou retêm sufixos inúteis.
• vs. Terminação Aleatória: O KAT superou significativamente a truncagem aleatória, confirmando que o sinal de baixo KL é um indicador específico de supervisão fraca, e não apenas uma métrica de comprimento.
• Ablação: Remover as etapas de "acordo" e "pós-acordo" do treinamento (como faz o KAT) resultou em maior acurácia do que incluí-las, validando que essas regiões fornecem sinais fracos ou prejudiciais.

5. Significância e Alegações

O artigo afirma que o KAT oferece uma abordagem mais eficaz e eficiente para a destilação on-policy ao repensar fundamentalmente como os sinais de supervisão são utilizados. Ao identificar que baixo KL nem sempre é benigno, o método evita a diluição dos sinais de aprendizado com tokens não informativos.

Os autores enfatizam que o KAT alcança esses ganhos sem alterar o objetivo de OPD ou introduzir novas funções de perda ou modelos auxiliares. Ele aproveita a computação de KL reversa já existente para tomar decisões online e adaptativas. O artigo conclui que o KAT torna a OPD tanto mais eficaz (maior acurácia) quanto mais eficiente (menor computação), particularmente em tarefas de raciocínio matemático, onde caminhos de raciocínio podem facilmente degradar-se em estados irrecuperáveis.

Limitações Notadas:
Os autores observam modestamente que o trabalho foca em raciocínio matemático. O comportamento da regra de terminação em domínios mais amplos, como geração de código, seguimento de instruções de texto aberto ou interações de múltiplos turnos, ainda precisa ser estudado. Além disso, embora os hiperparâmetros sejam intuitivos, uma orientação mais sistemática para o ajuste em diversos cenários poderia melhorar a usabilidade.