Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch

1. Problema e Contexto

Os Grandes Modelos de Linguagem (LLMs) alcançam desempenho de última geração (SOTA), mas seu custo computacional e de recursos dificulta a implantação. A Distilação de Conhecimento (KD) é uma técnica comum para treinar modelos menores (Alunos) que imitam modelos maiores (Professores), melhorando a eficiência.

No entanto, um desafio crítico surge quando o Modelo Professor e o Modelo Aluno utilizam tokenizadores e vocabulários distintos. Isso resulta em:

• Sequências de tokens de comprimentos diferentes.
• Espaços preditivos incompatíveis.
• Dificuldade em alinhar e comparar estados internos ou saídas.

O método atual SOTA para lidar com isso é o DSKD-CMA (Dual-Space Knowledge Distillation with Cross-Model Attention). Embora eficaz, seu funcionamento interno é uma "caixa preta", dificultando a otimização e a compreensão de como o alinhamento ocorre entre os dois espaços de vocabulário.

2. Metodologia Proposta

Os autores propõem uma análise sistemática do mecanismo de atenção do DSKD-CMA e introduzem uma nova abordagem para melhorar o alinhamento.

A. Análise do DSKD-CMA Original

O DSKK-CMA original usa um mecanismo de Atenção Cruzada entre Modelos (CMA):

• As embeddings do Aluno são projetadas no espaço do Professor para formar as Queries ($Q$).
• As embeddings do Professor servem como Keys ($K$).
• Matrizes de atenção são calculadas para projetar os estados ocultos de um modelo para o outro, permitindo o cálculo de perdas de distilação em espaços compartilhados.

Para entender o que o CMA aprende, os autores desenvolveram:

1. Probes de Alinhamento de Blocos (Chunks): Um processo manual de alinhamento baseado em deslocamento de caracteres para identificar quais tokens do Professor e do Aluno cobrem o mesmo trecho de texto.
2. Visualização de Heatmaps: Comparação entre as matrizes de atenção aprendidas pelo CMA e as matrizes de alinhamento de blocos explícitas.

B. Novas Abordagens: DSKD-CMA-GA e DSKD-CMA-CT

O problema central identificado é que as Keys (do Professor) e as Queries (do Aluno) vêm de distribuições diferentes, o que pode enfraquecer a eficácia da atenção. Para resolver isso, os autores propõem integrar o casamento de Key-Query usando aprendizado adversarial:

1. DSKD-CMA-GA (Generative Adversarial):

   • Utiliza um discriminador para distinguir entre Keys do Professor e Queries do Aluno.
   • O projetor de Queries do Aluno atua como um gerador, tentando "enganar" o discriminador para que as distribuições se tornem indistinguíveis.
   • Objetivo: Alinhar as distribuições de $p_K$ e $p_Q$ para melhorar a atenção.

2. DSKD-CMA-CT (Conditional Transport):

   • Utiliza uma versão restrita do Transporte Ótimo com um crítico treinável para encorajar o casamento bidirecional entre as distribuições de Key e Query, minimizando o custo de transporte.

A função de perda total combina a perda de distilação original com a perda de alinhamento de Key-Query ($L_{train} = L_{DSKD-CMA} + L_{KQ}$).

3. Principais Contribuições

1. Framework de Análise para DSKD-CMA:

   • Demonstraram que o CMA aprende implicitamente estruturas de "blocos" (chunks) de tokens, mas sofre com ruído em tokens repetidos e falta de localização precisa.
   • Mostraram que substituir o CMA por projeção de nível de bloco (CLP) mantém o desempenho, mas é computacionalmente mais caro, enquanto a atenção restrita a blocos (CLA) degrada o desempenho, sugerindo que o CMA beneficia-se de alinhamentos mais "grossos" e estáveis.

2. Novos Objetivos de Distilação (DSKD-CMA-GA e CT):

   • Introduziram métodos baseados em aprendizado adversarial para alinhar as distribuições de Keys e Queries de modelos com vocabulários diferentes.
   • Demonstraram que o alinhamento adversarial é superior ao transporte condicional neste contexto específico.

4. Resultados Experimentais

Os experimentos foram realizados utilizando:

• Aluno: GPT-2 (124M parâmetros).
• Professores: GPT-2 XL (mesmo tokenizador) e Qwen1.5 (tokenizador diferente, vocabulário muito maior).
• Dados: Dolly 15K (treino) e vários conjuntos de teste in-distribution (ID) e out-of-distribution (OOD).

Principais descobertas:

• Desempenho Geral: O método proposto DSKD-CMA-GA obteve ganhos consistentes e modestos na qualidade de geração de texto (medido por ROUGE-L) em comparação com o DSKD-CMA original.
• Dados OOD: O ganho foi particularmente notável em dados fora da distribuição, com uma melhoria média de +0,37 pontos no ROUGE-L.
• Fechamento da Lacuna: O DSKD-CMA-GA reduziu significativamente a lacuna de desempenho entre a KD com tokenizadores diferentes e a KD com o mesmo tokenizador. Em alguns casos, superou até mesmo a distilação com mesmo tokenizador.
• Interação com Divergência: A escolha da função de divergência (ex: KL, SKL, SRKL) é crucial e deve ser co-projetada com o método de alinhamento. O GA funcionou melhor com divergência KL e SKL.

5. Significado e Conclusão

Este trabalho é significativo porque:

• Transparência: Torna o mecanismo de atenção do DSKD-CMA interpretável, revelando suas forças (alinhamento de blocos) e fraquezas (localização).
• Eficiência em Vocabulários Diferentes: Oferece uma solução prática e eficaz para o problema de "mismatch" de vocabulário, que é comum ao integrar modelos de diferentes famílias (ex: GPT, Llama, Qwen).
• Avanço no Estado da Arte: A abordagem baseada em aprendizado adversarial (GA) prova ser uma ferramenta poderosa para alinhar representações latentes de modelos heterogêneos, permitindo que modelos menores aprendam de professores maiores mesmo sem compartilhar o mesmo espaço de vocabulário, aproximando-se ou superando o desempenho de cenários com vocabulário compartilhado.

Em suma, o artigo valida que o alinhamento explícito das distribuições de Keys e Queries via aprendizado adversarial é um passo fundamental para tornar a distilação de conhecimento entre LLMs heterogêneos mais robusta e eficiente.