Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation

Título: Token-Selective Dual Knowledge Distillation (TSD-KD)

Autores: Minsang Kim e Seung Jun Baek (Korea University e SK Telecom)

1. O Problema

A Distilação de Conhecimento (Knowledge Distillation - KD) é uma técnica fundamental para transferir a capacidade de raciocínio de modelos grandes ("professores") para modelos menores ("estudantes"), reduzindo custos de inferência, especialmente em tarefas que exigem Cadeias de Pensamento (Chain-of-Thought - CoT).

No entanto, métodos tradicionais de KD enfrentam desafios críticos em tarefas complexas de raciocínio:

• Mismatch de Distribuição: Métodos off-policy (treinamento em dados estáticos gerados pelo professor) sofrem com a discrepância entre a distribuição de dados do professor e a distribuição real gerada pelo estudante durante a inferência.
• Supervisão Excessiva e Forçada: Métodos on-policy recentes tentam mitigar o shift de distribuição, mas frequentemente forçam o estudante a imitar a distribuição completa do professor em todos os tokens gerados. Isso pode sobrecarregar modelos com capacidade limitada, causando um "mismatch" onde o estudante tenta seguir um caminho de raciocínio que não se alinha com seu nível de compreensão (ex: um estudante tentando imitar um professor em passos que ele ainda não domina).
• Falta de Autonomia: Forçar a imitação em cada token impede que o estudante desenvolva seu próprio processo de raciocínio, limitando a generalização.

2. Metodologia Proposta: TSD-KD

Os autores propõem o TSD-KD, um framework de distilação centrado no estudante que combina distilação indireta e distilação direta, ambas aplicadas de forma seletiva por token. O objetivo é fornecer feedback direcionado apenas nos pontos críticos, permitindo que o estudante raciocine com suas próprias palavras no restante do processo.

O método baseia-se em três pilares principais:

A. Distilação Indireta Guiada por Preferência (Teacher-Guided Preference Ranking)

• Conceito: Em vez de o professor ditar a distribuição de probabilidade de todos os tokens, o estudante gera seus próprios candidatos de tokens (top-k) em cada passo. O professor atua como um "oráculo de classificação", reordenando esses candidatos com base em qual caminho de raciocínio é preferível.
• Seletividade (Openers): A distilação indireta é aplicada apenas no início da resposta, definido como o "Opener". O Opener é a sequência inicial de tokens onde a entropia cumulativa atinge um certo percentil (ex: 10%). A hipótese é que os pontos de decisão mais críticos (onde o raciocínio pode desviar) ocorrem no início.
• Mecanismo: Utiliza o modelo Plackett-Luce para maximizar a probabilidade de o estudante alinhar sua preferência de candidatos com a do professor. Isso é uma forma de feedback "fraco" ou indireto, que preserva a autonomia do estudante.

B. Distilação Direta Baseada na Lacuna de Incerteza (Uncertainty Gap)

• Conceito: Aplica-se uma distilação direta (baseada em divergência) apenas em tokens onde há uma grande discrepância de confiança entre professor e estudante.
• Critério de Seleção: O foco são tokens onde o estudante está altamente incerto (alta entropia) mas o professor está altamente confiante (baixa entropia).
• Mecanismo: Utiliza uma função de gating (porta) sigmoide baseada na diferença de entropia ($H_{student} - H_{teacher}$) para modular a força da perda de distilação (usando Divergência Jensen-Shannon Generalizada). Isso evita forçar o estudante em tokens onde ele já é competente, reduzindo o shift de distribuição.

C. Regularização por Entropia Seletiva

• Conceito: Para aumentar a confiança do estudante em tokens críticos, aplica-se uma minimização de entropia seletiva.
• Mecanismo: Minimiza a entropia apenas dos tokens mais incertos (top-10% de entropia), tornando o estudante mais determinístico nesses pontos-chave sem restringir a exploração em outras partes da sequência.

Função de Perda Final:
A perda total combina os três componentes:
$$\mathcal{L} = \alpha \mathcal{L}_{Indirect} + \mathcal{L}_{Direct} + \mathcal{L}_{EM}$$
Onde $\alpha$ pondera a distilação indireta.

3. Contribuições Principais

1. Abordagem Centrada no Estudante: O framework muda o paradigma de "imitação cega" para "aprendizado reflexivo", onde o estudante propõe soluções e recebe feedback seletivo.
2. Seleção de Tokens por Entropia: Identifica e foca nos tokens de alta importância (início da cadeia e pontos de alta incerteza), ignorando a maioria dos tokens onde a imitação é desnecessária ou prejudicial.
3. Dualidade Indireta/Direta: Combina a flexibilidade da preferência (indireta) com a precisão da correspondência de distribuição (direta), criando um efeito sinérgico.
4. Superação do Professor: Demonstração empírica de que um modelo pequeno, treinado com TSD-KD, pode superar o próprio modelo professor em tarefas complexas de raciocínio.

4. Resultados Experimentais

Os experimentos foram conduzidos em 10 benchmarks desafiadores de raciocínio (Matemática, STEM, Síntese de Código, Raciocínio Lógico).

• Modelos Testados:
  • Qwen2.5: Professor (14B) -> Estudante (1.5B).
  • Gemma2: Professor (9B) -> Estudante (2B).
  • Qwen3: Professor (8B) -> Estudante (1.7B).
• Desempenho Geral: O TSD-KD alcançou o estado da arte (SOTA) em todos os benchmarks testados.
  • Superou a linha de base (estudante sem distilação) em até 54,4% de precisão.
  • Superou o segundo melhor método (GKD, MiniLLM, etc.) em até 40,3%.
• Superação do Professor (Student > Teacher):
  • Em 4 casos de benchmarks, o estudante TSD-KD superou seu próprio professor.
  • Exemplo notável: No benchmark MATH (Qwen2.5), o estudante (1.5B) atingiu 26.1%, superando o professor (14B) que teve 21.7% (melhoria de 20.3% sobre o professor).
  • No benchmark SciQ (Gemma2), o estudante superou o professor em 7.6%.
• Ablação: A remoção de qualquer componente (Distilação Indireta, Seletividade Direta ou Regularização de Entropia) resultou em queda significativa de desempenho, confirmando a necessidade da abordagem dual e seletiva.

5. Significado e Impacto

O trabalho TSD-KD representa um avanço significativo na compressão de modelos de linguagem para tarefas de raciocínio complexo.

• Eficiência Computacional: Permite que modelos pequenos (ex: 1.5B - 2B parâmetros) realizem raciocínios complexos com desempenho superior a modelos muito maiores, reduzindo drasticamente o custo de inferência.
• Mudança de Paradigma: Sugere que a distilação de conhecimento não deve ser sobre forçar a distribuição de probabilidade do professor, mas sim sobre guiar o processo de raciocínio do estudante nos momentos críticos, permitindo que ele internalize a lógica em vez de apenas copiar os tokens.
• Aplicabilidade: O método é robusto em diferentes arquiteturas (Qwen, Gemma) e escalas, sendo eficaz tanto para treinamento completo quanto para Parameter-Efficient Fine-Tuning (PEFT/LoRA).

Em resumo, o TSD-KD demonstra que a inteligência em modelos pequenos pode ser amplificada não pela imitação total, mas por uma supervisão inteligente, seletiva e que respeita a capacidade de auto-aperfeiçoamento do estudante.