Resumo Técnico: ARKD – Destilação de Divergência KL Bidirecional Guiada por Aprendizado por Reforço Adaptativo

1. Definição do Problema

A Destilação de Conhecimento (KD) é uma técnica primária para comprimir Grandes Modelos de Linguagem (LLMs) para reduzir os custos de inferência enquanto mantém o desempenho. No entanto, os métodos tradicionais de KD que dependem de um único objetivo de divergência de Kullback-Leibler (KL) enfrentam um compromisso fundamental na geração de texto aberta:

• KL Direta (FKL): Incentiva o comportamento de "cobertura de modo" (mode-covering), onde o estudante tenta cobrir todos os modos da distribuição do professor. Isso frequentemente leva à superestimação de probabilidades em regiões de baixa confiança e cauda longa, degradando a qualidade da geração.
• KL Reversa (RKL): Incentiva o comportamento de "busca de modo" (mode-seeking), focando nos modos dominantes e ignorando regiões de baixa probabilidade. Embora isso evite a superestimação da cauda longa, compromete a diversidade e a cobertura, podendo levar ao colapso de modo (mode collapse).

Os métodos de interpolação estática existentes (ex: ponderação fixa de 0.5:0.5) falham em capturar a natureza dinâmica do hiato entre estudante e professor conforme o treinamento evolui. Um coeficiente estático não consegue se adaptar às exigências variáveis de alinhamento de distribuição, que necessitam de diferentes equilíbrios entre FKL e RKL em diferentes estágios do treinamento.

2. Metodologia: ARKD

Os autores propõem o ARKD (Destilação de Divergência KL Bidirecional Guiada por Aprendizado por Reforço Adaptativo), um framework que atribui pesos dinamicamente ao FKL e ao RKL usando um agente de aprendizado por reforço (RL).

Fundamentação Teórica

O artigo estabelece primeiro a consistência teórica do FKL e do RKL. A análise mostra que ambos os objetivos convergem para a mesma solução ($q_\theta = p$) quando os gradientes são zero. Isso confirma que ambas as divergências são válidas para o pareamento de distribuições, mas exibem comportamentos complementares durante a trajetória de otimização, justificando seu uso conjunto.

Framework Baseado em RL

O ARKD reformula o problema de atribuição de pesos como um processo de decisão sequencial (Processo de Decisão de Markov):

• Estado ($s$): Um vetor de 6 dimensões que captura a dinâmica atual do treinamento, incluindo entropias do professor/estudante ($H_T, H_S$), variâncias ($\sigma^2_T, \sigma^2_S$) e valores atuais de perda FKL/RKL ($L_{FKL}, L_{RKL}$).
• Ação ($\alpha$): Uma rede de política ($\pi_\phi$) gera um peso escalar $\alpha \in (0, 1)$, onde $\alpha$ pondera o termo FKL e $(1-\alpha)$ pondera o termo RKL.
• Recompensa ($r$): Definida como a perda de destilação negativa ($r = -L_{distil}$) para o lote (batch) atual. O objetivo é maximizar a recompensa cumulativa (minimizar a perda total) ao longo da trajetória de treinamento.
• Otimização:
  • Atualização do Estudante: Minimiza a perda de destilação ponderada mais uma perda de modelagem de linguagem.
  • Atualização da Política: Utiliza o algoritmo REINFORCE com uma base de Média Móvel Exponencial (EMA) e regularização de entropia para estabilizar o treinamento e evitar a saturação de borda.

Dinâmica de Treinamento

A rede de política aprende a navegar adaptativamente pelo compromisso:

1. Fase de Exploração: No início do treinamento, a política explora uma gama mais ampla de $\alpha$ (ex: 0.3–0.6), enfatizando o FKL para garantir ampla cobertura e evitar o colapso de modo.
2. Fase de Convergência: À medida que a distribuição do estudante se aproxima da do professor, a política muda para valores de $\alpha$ mais baixos (dominância de RKL) para focar no alinhamento preciso e evitar a superestimação da cauda longa.
3. Fase de Estabilidade: A política estabiliza em um $\alpha$ específico (ex: $\approx 0.18$) para o refinamento final.

3. Principais Contribuições

1. Análise Teórica: O artigo fornece uma análise teórica estabelecendo a complementaridade e a consistência de convergência do FKL e do RKL, formando a base para a destilação adaptativa.
2. Framework ARKD: Introduz um novo framework de otimização conjunta baseado em RL, onde uma rede de política pondera dinamicamente o FKL e o RKL com base em características dependentes do estado, indo além de heurísticas estáticas.
3. Validação Empírica: Experimentos extensos demonstram que o ARKD melhora consistentemente a qualidade da geração, o ajuste de distribuição e a generalização em múltiplos benchmarks e escalas de modelos.

4. Resultados Experimentais

Os autores avaliaram o ARKD em tarefas de ajuste de instrução (instruction tuning) utilizando configurações GPT-2 (1.5B $\to$ 120M/340M) e LLaMA (13B $\to$ 7B).

• Desempenho In-Distribution (ID): No benchmark DollyEval, o ARKD alcançou o maior BertScore e pontuações de Rouge-L competitivas ou superiores em comparação com as linhas de base (SFT, SeqKD, FKL, RKL, FKL+RKL e MiniLLM). Por exemplo, no GPT-2 1.5B $\to$ 120M, o ARKD alcançou um Rouge-L de 24.5 e BertScore de 0.815, superando a forte linha de base MiniLLM.
• Generalização Out-of-Distribution (OOD): O ARKD demonstrou generalização superior em conjuntos de dados não vistos (SelfInst, VicunaEval, S-NI, UnNI), superando consistentemente todas as linhas de base em ambas as famílias de modelos.
• Estudos de Ablação:
  • vs. Escalonamento Estático/Linear: O ARKD superou significativamente a ponderação estática ( $\alpha$ fixo) e as estratégias de escalonamento linear (decaimento/aumento baseado no tempo), que falharam em capturar dinâmicas de alinhamento complexas.
  • vs. Heurísticas Gananciosas (Greedy): O ARKD superou as heurísticas gananciosas (selecionar a divergência de menor perda em cada etapa) em 0.4–0.6 pontos. A abordagem gananciosa foi considerada míope, muitas vezes degenerando para o RKL puro, enquanto o ARKD manteve um componente benéfico de pequeno FKL através da mistura contínua.
  • Análise de Trajetória: A trajetória de $\alpha$ aprendida confirmou uma estratégia inteligente: ênfase inicial no FKL para cobertura, transicionando para o RKL para precisão, um padrão não alcançável por design manual.

5. Significância e Alegações

O artigo afirma que o ARKD avança o estado da arte da destilação de conhecimento adaptativa ao abordar as limitações dos métodos baseados em KL estáticos. Ao utilizar aprendizado por reforço para equilibrar dinamicamente os comportamentos de cobertura de modo e busca de modo, o ARKD alcança um melhor alinhamento de distribuição que melhora tanto a qualidade da geração quanto a generalização.

Os autores posicionam o ARKD como uma abordagem principiológica que valida a necessidade de uma ponderação dinâmica e consciente do estado sobre heurísticas fixas. Eles observam que, embora a política atual dependa de um vetor de estado construído manualmente e recompensas de lote imediatas, o framework é compatível com famílias mais amplas de $f$-divergência e paradigmas on-policy. O trabalho fornece uma base sólida para a compressão eficiente de LLMs ao demonstrar que a otimização adaptativa guiada por RL pode consistentemente superar as linhas de base estáticas e gananciosas em diversas tarefas e escalas de modelos.