WDL-OPD: Weak-Driven On-Policy Distillation via Mixture-Constrained Co-Training

Resumo Técnico: WDL-OPD

Definição do Problema
A destilação on-policy (OPD) aborda o descompasso entre treino e teste inerente à destilação offline, treinando um modelo estudante em trajetórias que ele mesmo gera, utilizando um professor congelado para supervisão ao nível de token. No entanto, a OPD introduz uma instabilidade dinâmica: cada atualização de política altera o comportamento do estudante, o que, por sua vez, desloca a distribuição de estados visitados em subsequentes rollouts. Esse ciclo de feedback pode levar a um treinamento instável, crescimento de entropia ou degradação de trajetórias, particularmente quando os padrões de raciocínio do professor e do estudante são incompatíveis ou quando o conjunto compartilhado de tokens de alta probabilidade é pequeno. Métodos existentes como OPD2 e W2S-OPD tentam estabilizar o aprendizado usando sinais "delta" estáticos (diferenças entre um professor e seu modelo base) ou proxies contrastivos congelados, mas dependem de uma única política treinável.

Metodologia: WDL-OPD
Os autores propõem o WDL-OPD (Weak-Driven On-Policy Distillation), um framework de co-treinamento com restrição de mistura que utiliza duas políticas treináveis: uma política âncora ($\pi_A$) e uma política auxiliar ($\pi_U$).

1. Rollout e Construção de Suporte: Apenas a política âncora gera trajetórias ($y \sim \pi_A$). Para cada estado visitado (prefixo) $s_t$, um conjunto de suporte compartilhado $K_t$ é construído, consistindo nos top-$k$ tokens com os maiores logits da âncora. Tanto a política âncora quanto a auxiliar são avaliadas neste mesmo suporte para garantir consistência distributiva.
2. Mistura Geométrica: O método define um log-score misturado $m_t$ como uma média geométrica ponderada das log-probabilidades da âncora e da auxiliar no suporte compartilhado:
   $$m_t(v) = \lambda \ell_A(v) + (1-\lambda) \ell_U(v)$$
   Esta mistura é renormalizada para formar uma distribuição $q_M$.
3. Função de Perda: O objetivo de treinamento minimiza a divergência Kullback-Leibler reversa (KL) entre esta distribuição de mistura $q_M$ e a distribuição do professor congelado $\pi_T$:
   $$L = D_{KL}(q_M \parallel \pi_T)$$
   Crucialmente, os gradientes desta perda atualizam tanto a política âncora quanto a política auxiliar simultaneamente.
4. Limite Teórico: O artigo demonstra que, se a política auxiliar for congelada, o método reduz-se algebricamente a um alvo de proxy "âncora-mais-contraste", intimamente relacionado ao OPD2 e W2S-OPD. No entanto, ao manter a auxiliar treinável, o método introduz "graus de liberdade de nível de ramo", permitindo que as duas políticas dividam o sinal de correção do professor de acordo com seus distintos Jacobianos e estados do otimizador.

Principais Contribuições

• Formulação: Os autores formulam um algoritmo de OPD de duas políticas com mistura restrita, onde a distribuição auxiliar é parte do estado de otimização, em vez de um alvo fixo.
• Caracterização Teórica: Eles caracterizam a relação entre este método e abordagens delta estáticas, mostrando que o treinamento conjunto cria um espaço nulo no nível do ramo que um delta estático não consegue expressar.
• Consolidação Empírica: O artigo fornece um registro experimental consolidado em duas escalas de modelo (1.7B e 4B) e dois domínios (raciocínio matemático e geração de código), incluindo execuções negativas e reavaliações independentes para destacar modos de falha.

Resultos Experimentais
Os experimentos foram conduzidos utilizando modelos Qwen3 (1.7B e 4B) em raciocínio matemático (MATH500, GSM8K, etc.) e geração de código (HumanEval+, MBPP+, etc.).

• Desempenho: Em todas as quatro configurações de escala-domínio, o WDL-OPD produziu o checkpoint de estudante mais forte em comparação com as linhas de base de single-policy OPD registradas.
  • Matemática (4B): A precisão no MATH500 melhorou de 0.630 (melhor single-policy) para 0.685 (auxiliar do WDL-OPD).
  • Matemática (1.7B): A precisão no MATH500 melhorou de 0.521 para 0.585.
  • Código (4B): Configurações de single-policy OPD exibiram instabilidade (crescimento de entropia ou repetição), falhando em produzir um candidato final saudável. O WDL-OPD alcançou um score Dev500 reavaliado de 0.637.
  • Código (1.7B): O desempenho da single-policy degradou (0.106 vs base 0.184), enquanto o WDL-OPD atingiu 0.375.
• Dinâmica de Ramos: Os dois ramos não se comportaram como cópias redundantes. Na configuração de matemática 4B, a auxiliar (inicialmente mais fraca) ultrapassou a âncora para se tornar o melhor artefato final. Na configuração de matemática 1.7B, a âncora permaneceu superior. Isso sugere que a identidade e a especialização dos ramos são empiricamente significativas.
• Estabilidade: Na geração de código, onde a OPD de single-policy frequentemente sofreu com crescimento de entropia e degradação de trajetória, a abordagem de co-treinamento evitou essas assinaturas específicas de falha tempo suficiente para gerar checkpoints úteis, embora não tenha eliminado todas as patologias (ex: a repetição eventualmente apareceu na execução 4B).

Significância e Alegações
O artigo posiciona o WDL-OPD não como uma solução causal universal para todas as instabilidades de OPD, mas como um método que suporta uma hipótese de estabilização. Os autores afirmam explicitamente que seus resultados são descritivos em vez de uma alegação causal controlada, observando confusões como currículos distintos, estados de inicialização e orçamentos de computação entre seu método e as linhas de base de single-policy.

A principal significância reside na hipótese do buffer de otimização: o ramo auxiliar treinável pode absorver parte de um sinal de correção agudo, evitando o deslocamento imediato e potencialmente desestabilizador da distribuição do estado de rollout que ocorre quando uma única política deve lidar tanto com a correção quanto com a geração do estado. O artigo conclui que o próximo passo científico não é fazer alegações mais amplas, mas sim uma matriz de comparação controlada que isole o acoplamento de mistura de fatores como inicialização, capacidade e design de alvo.