DWA-KD: Dual-Space Weighting and Time-Warped Alignment for Cross-Tokenizer Knowledge Distillation

Resumo Técnico: DWA-KD

1. O Problema

A Distilação de Conhecimento (KD) é uma técnica fundamental para comprimir Grandes Modelos de Linguagem (LLMs), transferindo conhecimento de um modelo "professor" grande para um modelo "aluno" menor. No entanto, a maioria dos métodos tradicionais de KD assume que professor e aluno compartilham o mesmo tokenizador (vocabulário e esquema de tokenização).

Quando os modelos utilizam tokenizadores diferentes (Cross-Tokenizers KD - CTKD), surgem dois desafios principais:

1. Incompatibilidade de Tokenização: As sequências de tokens geradas por professor e aluno têm comprimentos e alinhamentos diferentes, dificultando a comparação direta.
2. Desalinhamento de Sequência: Métodos existentes muitas vezes tratam todos os tokens como igualmente importantes ou falham em capturar o contexto semântico e temporal entre as sequências desalinhadas.
   • Métodos baseados em Optimal Transport (OT) não preservam necessariamente a ordem temporal.
   • Métodos baseados em Dynamic Time Warping (DTW) anteriores (como CDM) são computacionalmente caros e utilizam loops em CPU, subutilizando GPUs.
   • Abordagens atuais frequentemente negligenciam a importância variável de tokens individuais (alguns são mais informativos que outros).

2. Metodologia Proposta: DWA-KD

Os autores propõem o DWA-KD (Dual-Space Weighting and Time-Warped Alignment), um framework que alinha professor e aluno em dois níveis: token e sequência.

A. Nível de Token: Ponderação Dual-Espaço (Dual-Space Weighting)
O método introduz um esquema de ponderação assimétrico para focar o aprendizado nos tokens mais informativos, evitando atualizações ruidosas:

• Espaço do Aluno: O peso de cada token é calculado pelo produto da entropia do aluno (incerteza) e da confiança projetada do professor (probabilidade máxima após mapeamento para o vocabulário do aluno).
  • Lógica: Se o aluno está incerto e o professor é confiante, o token recebe alto peso. Se o aluno já sabe (baixa entropia) ou o professor é incerto, o peso é reduzido.
• Espaço do Professor: O peso é baseado apenas na confiança do professor (baixa entropia). Tokens onde o professor é incerto recebem pesos menores para evitar ruído.
• Objetivo: Concentrar a capacidade de aprendizado nos tokens que realmente importam, normalizando os pesos para a duração da sequência para garantir estabilidade.

B. Nível de Sequência: Alinhamento por Warping Temporal (Time-Warped Alignment)
Para resolver o desalinhamento de sequências de comprimentos diferentes, o DWA-KD utiliza Soft-DTW (Soft Dynamic Time Warping):

• Aplicação Dupla: O Soft-DTW é aplicado tanto nas representações de embeddings (nível lexical) quanto nos estados ocultos finais (nível semântico/contextual).
• Banda Adaptativa (Attention-Informed Soft Band): Para evitar a complexidade quadrática excessiva e manter o alinhamento plausível, o método utiliza uma "banda" suave ao redor da diagonal.
  • O centro e a largura dessa banda são adaptados dinamicamente com base na entropia da atenção cruzada (Cross-Model Attention) entre professor e aluno.
  • Alta incerteza na atenção resulta em uma banda mais larga; baixa incerteza resulta em uma banda mais estreita.
• Vantagem: Isso permite um alinhamento diferenciável e robusto que captura tanto dependências lexicais quanto semânticas, preservando a ordem temporal.

C. Função de Perda Global
A perda total combina:

1. Perda de Entropia Cruzada (Supervisão padrão).
2. Perda de Distilação Ponderada Dual-Espaço (Token-level).
3. Perda de Alinhamento Soft-DTW (Sequence-level).

3. Principais Contribuições

1. Framework DWA-KD: Uma nova arquitetura de distilação cross-tokenizer que integra ponderação seletiva de tokens com alinhamento de sequência via Soft-DTW em múltiplas camadas (embedding e hidden state).
2. Esquema de Ponderação Dual-Espaço: Um mecanismo inovador que pondera tokens com base na incerteza do aluno e na confiança do professor, focando o treinamento em exemplos difíceis e confiáveis.
3. Alinhamento Semântico e Lexical: Uso de Soft-DTW com bandas adaptativas baseadas em atenção para alinhar trajetórias semânticas e estruturais entre modelos com tokenizadores distintos, superando limitações de métodos anteriores que ignoravam o contexto profundo.

4. Resultados Experimentais

Os autores avaliaram o DWA-KD em diversos benchmarks de seguimento de instruções (Dolly, Self-Instruct, Vicuna, Super-Natural-Instructions) com pares professor-aluno de diferentes escalas (ex: Mistral-7B → TinyLlama-1.1B, Qwen → GPT-2).

• Desempenho Superior: O DWA-KD superou consistentemente os baselines de última geração (incluindo ULD, DSKD, MinED e MultiLevelOT) em métricas ROUGE-L e em avaliações de qualidade de resposta via GPT-4.
• Alinhamento de Representação: Estudos de similaridade mostraram que o DWA-KD preserva a estrutura de representação "token-a-token" do professor com maior fidelidade do que outros métodos (menor distância de similaridade).
• Ablação: Estudos demonstraram que tanto a ponderação por entropia quanto o alinhamento Soft-DTW (especialmente com a banda adaptativa) contribuem positivamente e de forma complementar para o desempenho final.

5. Significado e Conclusão

O DWA-KD representa um avanço significativo na compressão de LLMs heterogêneos. Ao resolver o problema de desalinhamento de tokenização não apenas mapeando espaços, mas também ponderando a importância dos tokens e alinhando trajetórias temporais semânticas, o método permite que modelos menores aprendam de forma mais eficiente de modelos maiores, mesmo sem compartilhar vocabulário.

Limitações Notadas:

• A complexidade quadrática do Soft-DTW em relação ao comprimento da sequência pode aumentar o custo computacional e de memória.
• Dependência de uma projeção estável entre professor e aluno; uma inicialização ruim pode enfraquecer a supervisão inicial.

Em suma, o trabalho oferece uma solução robusta para a distilação de conhecimento em cenários do mundo real onde modelos com diferentes arquiteturas e tokenizadores precisam ser integrados.