Resumo Técnico: Alinhamento de Trajetória Multi-Granular (MTA) para Destilação de LLMs

Declaração do Problema

Os métodos existentes de Destilação de Conhecimento (KD) para Grandes Modelos de Linguagem (LLMs) focam principalmente no alinhamento de distribuições de probabilidade de saída (logits) ou estados ocultos intermediários em camadas fixas e estáticas. Embora eficazes para transferir sinais locais, essas abordagens frequentemente empregam uma estratégia de alinhamento uniforme, ao nível de token, em todas as profundidades da rede.

O artigo argumenta que esse tratamento uniforme está desalinhado com a natureza hierárquica da linguagem humana e a organização funcional dos modelos Transformer. Pesquisas linguísticas e estudos de interpretabilidade de Transformer indicam que:

1. Camadas inferiores enfatizam características de superfície, lexicais e de recuperação factual.
2. Camadas superiores codificam progressivamente informações semânticas abstratas e raciocínio composicional.

Ao tratar todos os tokens de forma isolada e aplicar o mesmo objetivo de alinhamento de maneira uniforme, os métodos atuais falham em capturar a trajetória evolutiva progressiva das representações à medida que se transformam de unidades lexicais em estruturas semânticas abstratas. Essa limitação restringe a capacidade do modelo estudante de generalizar para tarefas downstream.

Metodologia: Alinhamento de Trajetória Multi-Granular (MTA)

Para abordar essas limitações, os autores propõem o MTA, um framework modular que alinha a trajetória representacional dos modelos professor e estudante utilizando uma estratégia adaptativa às camadas. O MTA opera sob a hipótese de que a destilação deve respeitar os papéis semânticos dependentes da profundidade das camadas do Transformer.

Componentes Principais

1. Intervalos Multi-Granulares Adaptativos às Camadas:

   • Camadas Inferiores: Os alinhamentos são realizados ao nível de palavra (Intervalos de Palavras) para preservar a ancoragem lexical e dependências de alta granularidade.
   • Camadas Superiores: Os alinhamentos deslocam-se para o nível de frase (Frases Nominais e Verbais) para capturar semântica composicional e relações estruturais de ordem superior.
   • Construção de Intervalos: Unidades semânticas são extraídas usando um analisador sintático (por exemplo, spaCy). Pesos de importância de token são computados via um mecanismo de auto-atenção bidirecional (excluindo laços próprios) para agregar representações de token em representações de intervalo.

2. Função de Perda de Alinhamento Estrutural Dinâmico ($L_{DSA}$):

   • Este objetivo impõe consistência geométrica entre o professor e o estudante.
   • Em vez de corresponder vetores de token individuais, $L_{DSA}$ minimiza a discrepância nas distâncias relacionais pareadas entre intervalos semânticos (palavras ou frases) dentro de uma camada específica.
   • Prioriza intervalos semanticamente salientes ponderando o alinhamento com base na importância agregada dos tokens constituintes.
   • Matematicamente, alinha a estrutura topológica do espaço de representação:
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     onde $d(\cdot)$ é a distância cosseno, e $U$ representa a representação de intervalo ponderada.

3. Função de Perda de Alinhamento de Representação Oculta ($L_{Hid}$):

   • Para complementar as restrições estruturais, uma perda ao nível de recurso alinha estados ocultos específicos.
   • Como as dimensões do estudante são tipicamente menores que as do professor, um projetor linear aprendível mapeia os estados do estudante para o espaço do professor.
   • A perda minimiza a distância cosseno ponderada entre os estados projetados do estudante e os estados do professor, focando nos tokens mais informativos.

4. Otimização:

   • O MTA é projetado como um módulo plug-in. A perda total combina a perda de destilação base ($L_{Base}$) de frameworks existentes (por exemplo, DistiLLM, FDD) com as perdas estruturais e ocultas propostas:
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

Contribuições Principais

1. Identificação de Limitações: Os autores identificam que o alinhamento intermediário uniforme ignora a evolução hierárquica das representações, levando a uma transferência de conhecimento subótima.
2. Framework Adaptativo às Camadas: Eles introduzem o MTA, que transita dinamicamente do alinhamento ao nível de palavra em camadas inferiores para alinhamento estrutural ao nível de frase em camadas superiores, espelhando a natureza composicional de baixo para cima da linguagem.
3. Módulo Generalizado: Demonstra-se que o MTA é uma camada de otimização generalizada compatível com métodos de destilação de última geração (DistiLLM, DistiLLM-2, FDD) que compartilham tokenização.
4. Validação Empírica: Experimentos extensivos mostram que o MTA melhora consistentemente o desempenho em diversas arquiteturas de modelos (GPT-2, Qwen1.5, OPT) e benchmarks.

Resultados Experimentais

Os autores avaliaram o MTA em tarefas de seguimento de instruções usando três pares estudante-professor:

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

Métricas de Desempenho:

• ROUGE-L: O MTA superou consistentemente as linhas de base (FDD, DistiLLM, DistiLLM-2) em todos os conjuntos de dados (Dolly, SelfInst, Vicuna, S-NI). Por exemplo, na configuração GPT-2 1.5B $\to$ 120M, adicionar o MTA ao DistiLLM melhorou o ROUGE-L médio de 20,21 para 21,45.
• LLM como Juiz: Usando o GPT-4o-mini para avaliar precisão factual, coerência e utilidade, os modelos aprimorados com MTA mostraram melhorias consistentes nas pontuações em relação às linhas de base.

Estudos de Ablação:

• Componentes de Perda: Tanto $L_{DSA}$ (estrutural) quanto $L_{Hid}$ (recurso) contribuíram positivamente, com sua combinação produzindo os melhores resultados.
• Granularidade: A estratégia "Nível Completo" (adaptativa às camadas) superou estratégias estáticas (Apenas Palavra ou Apenas Frase), confirmando que alinhar a ancoragem lexical em camadas inferiores e semântica composicional em camadas superiores é crucial.
• Seleção de Camadas: O método é robusto a diferentes agendamentos de camadas, embora uma alocação híbrida (1 camada ao nível de palavra : 2 camadas ao nível de frase) tenha sido encontrada como ótima para as configurações testadas.

Eficiência:

• O MTA introduz sobrecarga de treinamento principalmente devido à extração de intervalos sintáticos (reduzida para ~0,27s/etapa via otimização) e correspondência geométrica.
• Os custos de inferência permanecem inalterados, pois a extração de intervalos não é utilizada no momento do teste.
• Comparações ajustadas por tempo confirmaram que os ganhos de desempenho devem-se ao objetivo de alinhamento, e não a um orçamento de treinamento aumentado.

Significado e Alegações

O artigo alega que o MTA oferece uma abordagem semanticamente fiel à destilação de conhecimento ao modelar explicitamente a trajetória representacional hierárquica dos LLMs. Ao ir além da correspondência estática ao nível de token para um alinhamento estrutural dinâmico e multi-granular, o MTA permite que os modelos estudantes repliquem melhor as dinâmicas internas de processamento de informações do professor.

Os autores posicionam o MTA não como um substituto para frameworks de destilação existentes, mas como um módulo complementar, plug-and-play que aprimora a generalização e a robustez de LLMs comprimidos. O trabalho sublinha a importância de respeitar a hierarquia linguística e a evolução de recursos dependente da profundidade no design de futuras técnicas de compressão de modelos.