Resumo Técnico: SRA (Alinhamento de Representação de Span para Destilação de Modelos de Linguagem Grandes)

Declaração do Problema

O artigo aborda o desafio da Destilação de Conhecimento entre Tokenizadores Diferentes (CTKD), onde um grande modelo de linguagem "professor" e um menor modelo "aluno" utilizam tokenizadores e vocabulários distintos. Os métodos existentes de CTKD dependem principalmente de estratégias de alinhamento ao nível de token ou de logits (por exemplo, usando distância de edição, transporte ótimo ou camadas de projeção). Os autores argumentam que essas abordagens são frequentemente frágeis e sensíveis a discrepâncias entre tokenizadores, pois tentam alinhar unidades discretas e atômicas que podem não corresponder semanticamente entre diferentes sistemas de vocabulário. Além disso, os métodos atuais frequentemente falham em capturar a dinâmica subjacente da evolução de características dentro da arquitetura Transformer quando os tokenizadores diferem.

Metodologia: Framework SRA

O framework SRA proposto reformula a CTKD ao deslocar a unidade fundamental de alinhamento de tokens individuais para spans robustos e agnósticos ao tokenizador. Esta abordagem fundamenta-se na analogia física de Sistemas Dinâmicos de Múltiplas Partículas (MPDS), onde o Transformer é visto como um sistema de partículas interagentes (tokens) evoluindo através das camadas.

1. Construção e Mapeamento de Span

• Definição de Span: Em vez de alinhar tokens diretamente, o SRA agrupa tokens em "spans" que correspondem às mesmas unidades textuais na entrada bruta.
• Alinhamento via LCS: Para corresponder spans entre os modelos professor e aluno, os autores computam a Longest Common Subsequence (LCS) das deslocamentos de caracteres produzidos por ambos os tokenizadores. Isso identifica segmentos correspondentes do texto de entrada, definindo índices de início e fim para spans alinhados em ambos os modelos sem exigir um mapeamento direto token-a-token.

2. Representação de Span via Centro de Massa (CoM)

Baseando-se na perspectiva MPDS, cada span é tratado como um aglomerado de partículas. O estado desse aglomerado é representado pelo seu Centro de Massa (CoM):

• Pooling Ponderado: A representação de um span é computada como uma média ponderada por atenção dos estados ocultos dos tokens constituintes.
• Pesos de Atenção: O peso $w_i$ para um token é derivado dos escores de atenção recebidos pelo último token da sequência através de todos os cabeçalhos de atenção. Isso mimetiza o conceito físico onde partículas "mais pesadas" (mais semanticamente salientes) dominam o centro do aglomerado.
• Fórmula: $C_i = \sum w_t H_t$, onde $H_t$ é o estado oculto e $w_t$ é o peso derivado da atenção.

3. Objetivos de Destilação

O SRA emprega uma função de perda combinada para transferir conhecimento ao nível de span:

• Transferência de Estado Oculto ao Nível de Span ($L_{HS}^{Span}$):

  • Alinha as representações CoM dos spans do aluno e do professor usando uma perda de similaridade de cosseno.
  • Regularizador Geométrico ($L_{Geo}$): Para preservar a integridade estrutural do espaço de representação, um regularizador é adicionado para manter as distâncias relativas (distância de cosseno) entre pares de spans no modelo do aluno, correspondendo às do modelo professor.
  • Perda Ponderada: A perda é ponderada pela saliência semântica dos spans (derivada da soma dos pesos dos tokens dentro do span), garantindo que o aluno se concentre nas regiões mais informativas.

• Destilação de Logits de Span Alinhados ($L_{KD}^{Span}$):

  • Aborda a incompatibilidade de vocabulário ($V_{tea} \neq V_{stu}$) restringindo a transferência de conhecimento ao subespaço compartilhado $\tilde{V} = V_{tea} \cap V_{stu}$.
  • Logits ao nível de span são agregados a partir de logits de tokens dentro do vocabulário compartilhado, e uma divergência de Kullback-Leibler (KL) é computada entre as distribuições do professor e do aluno sobre este espaço compartilhado.

• Objetivo Total: A perda final combina a perda de entropia cruzada padrão sobre rótulos verdadeiros com as perdas de destilação ao nível de span (estados ocultos e logits), equilibradas por um hiperparâmetro $\alpha$.

Contribuições Principais

1. Novo Framework Conceitual: O artigo introduz um framework de CTKD fundamentado na visão de Sistema Dinâmico de Múltiplas Partículas dos Transformers, modelando spans como aglomerados de partículas e utilizando seu Centro de Massa para representações estáveis e aditivas.
2. Destilação ao Nível de Span: Propõe um método para destilar conhecimento ao nível de span, abrangendo tanto representações ocultas quanto logits. Isso inclui um esquema de ponderação derivado de atenção para priorizar spans salientes e um regularizador geométrico para preservar a geometria relativa do espaço de representação.
3. Validação Empírica: Os autores demonstram que o SRA supera significativamente as bases de referência (baselines) de CTKD mais avançadas (incluindo ULD, DSKD, MinED e MultilevelOT) em múltiplos pares professor-aluno e benchmarks.

Resultados Experimentais

Os autores avaliaram o SRA em diversos pares de arquiteturas e tokenizadores diferentes (por exemplo, Qwen1.5/2.5, Mistral-7B, GPT-2, OPT, TinyLLaMA) usando conjuntos de dados de seguimento de instruções (Dolly, Vicuna, SelfInst, S-NI, DialogSum).

• Desempenho: O SRA alcançou consistentemente as pontuações ROUGE-L mais altas e o melhor desempenho médio em todas as configurações em comparação com bases de referência robustas.
• Qualidade Semântica: Avaliações semânticas usando GPT-4o-mini como juiz confirmaram que o SRA produz respostas mais fiéis semanticamente e alinhadas contextualmente.
• Estudos de Ablação:
  • A remoção do regularizador geométrico ou do mecanismo de pooling ponderado resultou em degradação de desempenho, confirmando a necessidade de preservar relações estruturais e focar em spans salientes.
  • O método demonstrou robustez mesmo em cenários de mesmo tokenizador, superando métodos baseados em distância e outras abordagens de CTKD.
• Eficiência: Embora o SRA exija ligeiramente mais memória de GPU do que algumas bases de referência devido ao emparelhamento e agregação de spans, alcançou o menor tempo médio de treinamento por passo entre os métodos comparados.

Significado e Alegações

O artigo alega que o SRA fornece uma abordagem robusta e fundamentada fisicamente para destilação de conhecimento que supera a fragilidade do alinhamento ao nível de token. Ao deslocar a unidade de alinhamento para spans e aproveitar o conceito de Centro de Massa, o SRA suporta naturalmente a transferência da evolução dinâmica de características, camada por camada. Os autores afirmam que seu método valida o princípio de que a destilação deve focar na dinâmica da evolução geométrica do sistema e não apenas nas saídas finais, oferecendo uma solução escalável para implantar modelos eficientes em ambientes heterogêneos de tokenizadores.

O trabalho é apresentado como um passo em direção a uma transferência de conhecimento mais estável e abrangente, embora os autores reconheçam limitações, como a natureza estática do mapeamento atual de logits e o custo computacional da inferência do professor sob demanda necessária para o alinhamento de spans.