Hybrid Policy Distillation for LLMs

1. O Problema

A compressão de Grandes Modelos de Linguagem (LLMs) através de Distilação de Conhecimento (KD) é essencial para a implantação prática, mas enfrenta desafios fundamentais relacionados a três eixos de design interconectados:

1. Direção da Divergência:
   • KL Direto (Forward KL - FKLD): Promove a cobertura de todos os modos da distribuição do professor, mas tende a gerar previsões excessivamente suaves (over-smoothing), perdendo a precisão.
   • KL Reverso (Reverse KL - RKLD): Foca nos modos de alta probabilidade do professor (mode-seeking), mas pode levar a instabilidade no treinamento quando há uma grande lacuna entre aluno e professor, ignorando saídas válidas de baixa probabilidade.
2. Estratégia de Otimização: O uso de funções de perda (loss) versus recompensas (reward) gera dinâmicas de otimização e vieses de gradiente diferentes.
3. Regime de Dados:
   • Off-policy: Usa dados estáticos (gerados pelo professor ou ground-truth). É eficiente, mas sofre de desalinhamento entre o comportamento de treino e inferência.
   • On-policy: Amostra dados da própria política do aluno. Oferece melhor alinhamento, mas é computacionalmente caro e sofre de train-inference mismatch se não for bem gerido.

A escolha isolada de um desses fatores é insuficiente. Métodos existentes frequentemente otimizam apenas uma direção ou regime, resultando em instabilidade, ineficiência computacional ou perda de qualidade na geração.

2. Metodologia: Distilação de Política Híbrida (HPD)

Os autores propõem a HPD, um método unificado que integra as vantagens complementares do KL Direto e Reverso, combinando dados off-policy com amostragem on-policy leve.

Visão Unificada

O trabalho reformula a KD como um objetivo de verossimilhança ponderada (reweighted log-likelihood) no nível do token. A função de perda geral é definida como:
$$L(\theta) = \min_{\theta} -\mathbb{E}_{(s_t, a_t) \sim D_\pi} [w(a_t | s_t) \log q_\theta(a_t | s_t)]$$
Onde $w(a_t | s_t)$ é um sinal de recompensa que captura a discrepância local entre o aluno e o professor.

Mecanismo Híbrido

A HPD opera em dois componentes principais para cada estado $s_t$:

1. Token de Especialista (Ground Truth/Professor):

   • Utiliza um estimador de KL Reverso negativo ($k_1$) para medir a lacuna entre o professor e o aluno no token correto.
   • Se o aluno subestima o token especialista ($k_1 > 0$), aplica-se um peso baseado no KL Direto (usando a probabilidade do professor) para forçar o aluno a cobrir esse modo.
   • Se o aluno superestima o token ($k_1 \leq 0$), o peso do KL Direto é mascarado para evitar gradientes conflitantes, mantendo apenas o sinal de correção do KL Reverso.

2. Token Amostrado pelo Aluno (Não-Especialista):

   • O aluno amostra um token diferente do especialista ($a_t \sim q_\theta$).
   • Calcula-se o estimador $k'_1$ para este token.
   • Se o aluno superestima um token não especialista ($k'_1 < 0$), aplica-se um peso negativo (KL Reverso) para suprimir essa probabilidade.
   • Reforço (Reinforce): Quando o aluno suprimir um token não especialista, o peso do token especialista é duplicado (se $k_1 > 0$). Isso redistribui a massa de probabilidade suprimida diretamente para a ação correta do especialista.

Vantagens Técnicas

• Eficiência: Mantém a eficiência da supervisão "one-hot" (como no SFT) ao evitar o cálculo exato da divergência sobre todo o vocabulário, usando estimadores Monte Carlo no nível do token.
• Estabilidade: O mecanismo de mascaramento evita gradientes conflitantes, enquanto a amostragem leve on-policy corrige o desalinhamento de distribuição sem o custo de rollouts completos de sequência.

3. Contribuições Principais

1. Visão Unificada: Apresenta uma perspectiva de verossimilhança ponderada que conecta SFT, FKLD e RKLD, permitindo a análise sistemática de métodos de distilação.
2. Algoritmo HPD: Propõe um método que hibridiza KL Direto e Reverso em regimes de dados off-policy e on-policy leve, preservando a eficiência computacional.
3. Inicialização Forte: Demonstra que a HPD serve como uma excelente inicialização para métodos on-policy puros (OPD), levando a ganhos adicionais de desempenho.

4. Resultados Experimentais

Os autores validaram a HPD em três domínios: Raciocínio Matemático (longa geração), Diálogo (personalização) e Código.

• Raciocínio Matemático (Off-policy):
  • A HPD superou consistentemente SFT, SeqKD, RKLD e JSD em modelos Qwen 2.5 e LLaMA 3.
  • Destaque: Modelos de 3B de parâmetros com HPD alcançaram desempenho comparável a modelos maiores, com melhorias de 41% (Qwen 3B) e 77,9% (LLaMA 3 3B) em métricas de raciocínio.
  • Mostrou maior estabilidade de entropia durante o treinamento, evitando o colapso de entropia comum no SFT.
• Diálogo e Código (Off-policy):
  • Em tarefas de diálogo (MT-Bench, AlpacaEval), a HPD manteve melhor coerência conversacional e capacidade de multi-turno.
  • Em código (HumanEval, MBPP), alcançou o melhor desempenho médio, demonstrando robustez e menor variância entre tarefas.
• Dados On-policy (Raciocínio):
  • Quando combinada com OPD, a HPD superou a linha de base de dois estágios (SFT + OPD), atingindo os melhores escores médios e melhorando a convergência.
  • Evitou o colapso excessivo de entropia no tempo de teste, alinhando melhor o comportamento de inferência do aluno ao do professor.
• Ablação: A remoção da amostragem do aluno ou da operação de reforço resultou em convergência prematura ou redução mais lenta da divergência KL, confirmando a necessidade de ambos os componentes.

5. Significado e Impacto

O trabalho da HPD é significativo porque:

• Resolve o Dilema de Trade-off: Equilibra a cobertura de modos (FKLD) e a busca por modos (RKLD) de forma dinâmica e adaptativa, em vez de usar coeficientes fixos.
• Eficiência Computacional: Oferece uma alternativa viável e mais barata aos métodos de distilação on-policy completos (como MiniLLM ou GKD), que exigem rollouts sequenciais caros.
• Versatilidade: Funciona bem em tarefas de geração curta (diálogo) e longa (raciocínio), sendo aplicável tanto para inicialização quanto para treinamento contínuo.
• Futuro: Abre caminho para aplicar distilação em estágios mais precoces do treinamento (mid-training ou pré-treinamento), potencialmente melhorando a eficiência geral do desenvolvimento de LLMs.

Em resumo, a HPD representa um avanço na teoria e prática de compressão de LLMs, oferecendo um método que é simultaneamente estável, eficiente e de alto desempenho, superando as limitações das abordagens de distilação tradicionais.