Resumo Técnico: X-Token: Destilação de Conhecimento Cruzada entre Tokenizadores Guiada por Projeção

Declaração do Problema

A Destilação de Conhecimento (KD) tipicamente exige um tokenizador compartilhado entre os modelos professor e aluno, efetivamente prendendo os praticantes a famílias específicas de modelos. Essa restrição impede o aproveitamento de professores mais fortes ou especializados com vocabulários incompatíveis (por exemplo, destilar de um modelo Qwen ou Phi para um aluno Llama) e dificulta a destilação multi-professor entre diversas famílias de tokenizadores.

Os métodos existentes de KD cruzada entre tokenizadores dividem-se em duas categorias: baseados em representação (alinhamento de estados ocultos) e baseados em distribuição de logits (operando sobre probabilidades de saída). O último é preferido para pré-treinamento contínuo, pois não requer alterações arquiteturais. No entanto, o método baseado em logits mais avançado atual, GOLD, sofre de duas limitações estruturais críticas ao lidar com tokenizadores incompatíveis:

1. Falha em Tokens Incomuns: Quando os tokenizadores fragmentam o texto de forma diferente (por exemplo, um professor dividindo numerais de múltiplos dígitos enquanto o aluno os agrupa), tokens críticos caem em um subconjunto "não correspondente". Esses tokens são degradados por ruído agnóstico à identidade proveniente do emparelhamento baseado em rank e suprimidos por gradientes do termo KL comum, levando a quedas severas de desempenho (por exemplo, a precisão no GSM8k caindo de 12,89 para 2,56).
2. Emparelhamento Excessivamente Conservador: A igualdade estrita de strings 1-a-1 exclui alinhamentos semânticos válidos onde as formas superficiais diferem (por exemplo, um token de aluno "Centenas" correspondendo a tokens de professor "Hund" + "reds"), deixando sinais de supervisão úteis não explorados.

Metodologia: X-Token

O X-Token aborda essas questões através de um framework de três componentes: Alinhamento de Span, uma Matriz de Projeção ($W$) e Formulações de Perda Complementares.

1. Alinhamento de Sequência

Para lidar com comprimentos de sequência diferentes e desalinhamento posicional, o X-Token emprega alinhamento de span. Tokens são agrupados em blocos que decodificam para o mesmo texto subjacente. Uma fusão pela regra da cadeia é aplicada dentro de cada bloco para combinar probabilidades por token em uma única distribuição de nível de bloco, criando unidades alinhadas para destilação.

2. Matriz de Projeção ($W$)

Uma matriz de projeção esparsa $W \in \mathbb{R}^{|\mathcal{V}_S| \times |\mathcal{V}_T|}$ mapeia as probabilidades dos tokens do aluno para o espaço de vocabulário do professor. Ela é construída deterministicamente em duas passagens:

• Passagem de Correspondência Exata: Define $W[s, t] = 1$ para tokens com correspondências de strings canonizadas.
• Passagem de Regra de Múltiplos Tokens: Para tokens de aluno não correspondidos, o texto decodificado é re-tokenizado pelo professor. O token do aluno é mapeado para a sequência resultante de tokens do professor com pesos decrescentes (por exemplo, $\beta=0,9, \gamma=0,1$).
  A matriz é inicializada a partir de regras de tokenizador e pode ser opcionalmente ajustada finamente durante o treinamento.

3. Formulações de Perda Complementares

O X-Token introduz dois modos de perda, selecionados com base em uma auditoria de cobertura de tokens críticos (por exemplo, numerais) no vocabulário do aluno:

• P-KL (KL de Projeção): Projetado para cenários onde tokens críticos caem no conjunto "incomum" (por exemplo, professor Qwen). Ele remove a partição inteiramente. A distribuição do aluno é projetada via $W$ para o espaço de vocabulário do professor, e uma divergência KL direta é computada entre a distribuição do professor e a distribuição do aluno projetada. Isso elimina gradientes baseados em rank errôneos e efeitos supressivos.
• H-KL (Híbrido-KL): Projetado para cenários onde a partição é estruturalmente sólida (por exemplo, professor Phi-4-mini). Ele mantém a estrutura híbrida, mas relaxa a definição do conjunto "comum". Em vez de igualdade estrita de strings, inclui pares $(s, t^*)$ onde $t^*$ é o token do professor de maior rank para o token do aluno $s$ sob $W$. Isso captura alinhamentos quase equivalentes enquanto preserva a supervisão mais nítida do KL direto em pares alinhados por identidade.

4. Destilação Multi-Professor

O X-Token estende-se naturalmente a múltiplos professores agregando perdas por professor. Ele emprega pesagem estática (que superou esquemas adaptativos em experimentos) para combinar sinais de professores heterogêneos, aproveitando suas forças complementares (por exemplo, raciocínio vs. senso comum).

5. Escalonamento Dinâmico

Para estabilizar o treinamento, a perda de destilação é reescalada dinamicamente a cada passo para corresponder à magnitude da perda de entropia cruzada, prevenindo instabilidade de otimização causada por escalas de perda variáveis.

Contribuições Principais

1. Alinhamento Cruzado entre Tokenizadores Determinístico: Introdução de uma matriz de projeção esparsa $W$ construída via regras de strings e decodificação de múltiplos tokens, permitindo alinhamento direto de vocabulário sem componentes auxiliares treináveis.
2. Formulações de Perda Duplas: A proposta de P-KL e H-KL, que se adaptam a diferentes regimes de incompatibilidade de tokenizadores. O P-KL recupera o sinal para tokens não correspondidos críticos, enquanto o H-KL refina o emparelhamento para tokens bem alinhados.
3. Capacidade Multi-Professor: Demonstração de que o X-Token permite destilação eficaz a partir de professores heterogêneos, com resultados mostrando que combinações complementares de professores produzem desempenho superior ao de emparelhamentos de professor único ou homogêneos.
4. Alinhamento de Sequência Robusto: Um método escalável de alinhamento de span baseado em programação dinâmica que garante unidades consistentes com o texto para destilação.

Resultados Experimentais

Os experimentos foram conduzidos usando Llama-3.2-1B como aluno e Qwen3-4B, Phi-4-mini-Instruct e Llama-3.2-3B como professores.

• Desempenho vs. SOTA: O X-Token supera o estado da arte atual (GOLD) em +3,82 pontos médios ao usar um professor Qwen3-4B (com uma melhoria de 6x no GSM8k, de 2,56 para 15,54) e em +0,5 pontos com um professor Phi-4-mini.
• Superioridade Cruzada entre Famílias: No regime Qwen, o X-Token (P-KL) superou a KD do mesmo tokenizador a partir de um professor mais fraco Llama-3.2-3B, demonstrando que a destilação cruzada entre tokenizadores pode exceder os limites da mesma família quando a formulação de perda está correta.
• Ganhos Multi-Professor: Uma configuração de dois professores (Phi-4-mini + Llama-3.2-3B) melhorou a destilação de professor único em +1,3 pontos, destacando o valor de combinar forças diversas de professores.
• Descobertas de Ablação:
  • A seleção de perda é crítica: o P-KL falha no Phi-4-mini, e o H-KL falha no Qwen3-4B.
  • Aprender conjuntamente a matriz de projeção $W$ produz ganhos modestos, mas consistentes, sobre uma inicialização baseada em regras congelada.
  • A pesagem estática de professores superou esquemas adaptativos (baseados em entropia ou confiança).

Significado

O artigo afirma que o X-Token demonstra que um alinhamento cuidadoso tanto no nível da sequência quanto no nível do vocabulário, combinado com um design de perda adaptativo, é fundamental para desbloquear todo o potencial da destilação de conhecimento cruzada entre tokenizadores. Ao abordar os modos de falha específicos de métodos baseados em partição, o X-Token permite que os praticantes aproveitem professores mais fortes, especializados ou diversos sem serem constrangidos pela compatibilidade de tokenizadores, quebrando efetivamente o "aprisionamento por tokenizador" que atualmente limita a escalabilidade e especialização de modelos.