Trust-Region Behavior Blending for On-Policy Distillation

Resumo Técnico: Mistura de Comportamento de Região de Confiança para Destilação On-Policy

Declaração do Problema

A destilação on-policy (OPD) aborda o descompasso de prefixo inerente à destilação de conhecimento offline, treinando um modelo estudante em prefixos amostrados de sua própria política atual enquanto busca corresponder a um professor mais forte. No entanto, essa abordagem introduz uma fragilidade crítica nos estágios iniciais do treinamento: modelos estudantes fracos costam gerar prefixos de baixa qualidade. Quando a distribuição de trajetórias do estudante é ruim, a supervisão do professor é aplicada a prefixos de baixa qualidade ou "inutilizáveis", limitando a eficácia do sinal de aprendizado.

Tentativas anteriores de estabilizar esse processo frequentemente envolvem mover a coleta para off-policy (por exemplo, injetando tokens do professor ou alterando distribuições alvo). Por outro lado, rollouts puramente de estudantes preservam a distribuição, mas sofrem com a baixa qualidade inicial. O desafio central é melhorar a qualidade dos rollouts iniciais sem deslocar permanentemente o processo de coleta para off-policy ou alterar o objetivo fundamental de destilação por prefixo.

Metodologia: Mistura de Comportamento de Região de Confiança (TRB)

Os autores propõem a Mistura de Comportamento de Região de Confiança (TRB), um método de warmup projetado para controlar a política de comportamento durante a fase inicial de coleta de rollouts. A TRB opera no princípio de encontrar uma política de comportamento que seja o mais próxima possível do professor, mas estritamente contida dentro de uma "região de confiança" centrada na política atual do estudante.

1. Otimização por Prefixo
Em cada prefixo de geração $h$, a TRB define uma política de comportamento $\mu^*(\cdot | h)$ resolvendo um problema de otimização restrita. O objetivo é minimizar a divergência de Kullback-Leibler (KL) em relação à política do professor $\pi_T$, sujeita a uma restrição de divergência KL em relação à política do estudante $\pi_S$:

$$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{sujeito a} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$$

Aqui, $\epsilon$ representa o desvio local permitido em relação ao estudante. Esta formulação garante que a distribuição de amostragem seja puxada em direção aos tokens suportados pelo professor, mas permaneça limitada dentro de uma região de confiança centrada no estudante.

2. Solução de Forma Fechada
A otimização resulta em uma solução de forma fechada onde a política de comportamento é uma interpolação geométrica entre as políticas do estudante e do professor:
$$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$$
O coeficiente $\beta^*(h)$ é o maior valor viável em $[0, 1]$ tal que a restrição de KL seja satisfeita. Se o professor já estiver dentro da região de confiança, $\mu^*$ torna-se o professor; se $\epsilon=0$, retorna ao estudante.

3. Warmup com Anelamento (Annealed Warmup)
A TRB é aplicada apenas durante uma fase inicial de warmup. O orçamento de KL ($\epsilon$) é anulado linearmente para zero ao longo de um horizonte $K$:
$$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$$
Isso permite uma forte orientação do professor no início do treinamento para melhorar a qualidade do prefixo, enquanto retorna gradualmente a coleta de rollouts para a amostragem pura do estudante conforme o estudante melhora. Crucialmente, a perda de destilação de reverse-KL por prefixo permanece inalterada; apenas a política de coleta de dados é modificada.

Principais Contribuições

• Uma Nova Estratégia de Warmup: A TRB introduz um método para controlar o comportamento inicial de rollout ao otimizar uma política guiada pelo professor sob uma restrição explícita de KL centrada no estudante, em vez de simplesmente injetar tokens ou alterar a função de perda.
• Preservação de Propriedades On-Policy: Ao contrário de métodos que misturam persistentemente as políticas do professor e do estudante ou alteram a distribuição alvo, a TRB garante que o treinamento eventualmente retorne aos rollouts puramente on-policy, preservando a distribuição alvo para a fase principal de treinamento.
• Formulação Teórica: O artigo fornece uma solução de forma fechada para a política de comportamento que equilibra a orientação do professor com a estabilidade do estudante, justificada pela monotonicidade da divergência KL em relação ao coeficiente de interpolação.

Resultados Experimentais

Os autores avaliaram a TRB em dois cenários de destilação de raciocínio matemático:

1. Qwen3-1.7B-Base destilado de Qwen3-8B.
2. Qwen3-0.6B-Base destilado de Qwen3-4B.

Desempenho:

• A TRB alcançou o melhor desempenho médio entre todos os métodos comparados em ambos os cenários, superando a vanilla OPD, a mistura de $\epsilon$-fixo, Veto, SKD (injeção intercalada de professor) e warmup de temperatura e warmup de SFT.
• No cenário Qwen3-1.7B, a TRB alcançou um pass@1 médio de 33,2%, comparado a 32,3% da vanilla OPD.
• No cenário Qwen3-0.6B, a TRB alcançou 44,4%, comparado a 44,0% da vanilla OPD.

Análise:

• Dinâmica de Treinamento Inicial: A análise de trajetória (Figura 2) mostra que, embora alguns métodos (como o SKD) subam mais rápido inicialmente, eles não necessariamente sustentam o desempenho final mais alto. A TRB proporciona um "início rápido" que leva a resultados finais superiores.
• Entropia e Qualidade do Prefixo: Durante a fase de warmup, a TRB reduz a entropia média dos tokens do professor nos prefixos visitados, indicando uma orientação de maior qualidade. Os testes de Step-0 (Figura 4) demonstram que os prefixos gerados sob a TRB apresentam taxas de sucesso mais altas quando continuados tanto pelo professor quanto pelo estudante em comparação com os prefixos da vanilla OPD, sugerindo que a TRB desloca com sucesso a distribuição de estados iniciais para trajetórias mais aprendíveis.

Significância e Alegações

O artigo afirma que a TRB aborda efetivamente a fragilidade da destilação on-policy precoce ao melhorar a qualidade da distribuição de dados inicial sem se desviar permanentemente do objetivo on-policy.

• Escopo Modesto: Os autores declaram explicitamente que seu estudo é limitado a dois cenários de raciocínio matemático com pares Qwen3-Base e um protocolo de avaliação baseado em correção. Eles não alegam que os cronogramas de warmup específicos se transfiram sem alterações para outros domínios ou lacunas professor-estudante.
• Compromisso de Custo: O artigo reconhece que a TRB aumenta os custos de tempo de treinamento durante a fase de warmup devido à necessidade de decodificação online do professor e coexistência estudante-professor. No entanto, eles observam que esses custos são temporários, pois o treinamento retorna ao perfil de execução padrão e mais eficiente da OPD após o horizonte de warmup.
• Insight Central: Os resultados sugerem que o comportamento off-policy guiado pelo professor é mais benéfico quando os prefixos visitados pelo estudante estão significativamente desalinhados com o professor. Uma vez resolvido esse descompasso, a orientação off-policy contínua pode oferecer retornos decrescentes, apoiando a estratégia de anelar a política de comportamento de volta ao estudante.