Compressing Sequences in the Latent Embedding Space: $K$-Token Merging for Large Language Models

Título: Compressão de Sequências no Espaço de Embedding Latente: K-Token Merging para Grandes Modelos de Linguagem

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) enfrentam custos computacionais e de memória proibitivos ao processar prompts longos. Isso ocorre porque o mecanismo de atenção total (self-attention) escala quadraticamente com o comprimento da entrada ($O(n^2)$).

Embora a compressão de tokens seja uma solução natural, as abordagens existentes apresentam limitações:

• Compressão Hard (Dura): Remove ou resume tokens. É eficaz para tarefas com informação esparsa, mas falha em tarefas densas (como revisão de documentos, raciocínio matemático ou código), onde cada token pode ser crítico.
• Compressão Soft (Suave): Aprende novos tokens ou adapta parâmetros para representações mais compactas, mas opera quase exclusivamente no espaço de tokens, ignorando a redundância inerente ao espaço de embedding latente.
• Ineficiência de Embedding: O artigo destaca uma enorme lacuna de eficiência: um token consome milhares de bits (ex: 28.672 bits no Qwen 2.5), enquanto a informação teórica necessária para identificá-lo é muito menor (~18 bits). O espaço de embedding atual é subutilizado e redundante.

2. Metodologia: K-Token Merging

O artigo propõe o K-Token Merging, um framework de compressão que atua diretamente no espaço de embedding latente, em vez de manipular a sequência de tokens diretamente.

Principais Componentes:

1. Codificador Leve (Encoder $f$):
   • Um pequeno MLP (Multi-Layer Perceptron) que recebe um bloco contíguo de $K$ embeddings de tokens e os funde em um único embedding comprimido.
   • Inicialização Inteligente: O codificador é inicializado para que sua saída seja a média (pooling) dos $K$ embeddings originais. O MLP é inicializado com pesos próximos de zero, garantindo que, no início do treinamento, a representação seja estável e semanticamente fundamentada, acelerando a convergência.
2. Adaptação via LoRA:
   • O LLM base é ajustado (fine-tuned) usando LoRA (Low-Rank Adaptation) para entender e processar esses embeddings comprimidos.
   • A função de perda é calculada apenas nas posições dos tokens originais (não comprimidos) gerados, garantindo que a qualidade de geração no vocabulário original seja mantida.
3. Fluxo de Operação:
   • Fase de Preenchimento (Prefill): O encoder comprime a entrada de $N$ tokens em $N/K$ embeddings. Isso reduz drasticamente o custo de atenção quadrática.
   • Fase de Geração: O modelo gera tokens originais (não comprimidos). Esses novos tokens são adicionados ao prefixo misto (comprimido + não comprimido) para a geração autoregressiva subsequente.

Diferenciação Crítica:
Diferente da expansão de vocabulário (que criaria tokens de "cauda" raros e inviáveis para grandes $K$), o K-Token Merging não expande o vocabulário. Ele generaliza naturalmente para combinações de tokens nunca vistas durante a inferência, pois o encoder aprende a mapear sequências, não a memorizar tokens fixos.

3. Contribuições Principais

• Novo Framework: Introdução do K-Token Merging, uma abordagem de compressão suave que visa redundâncias no espaço de embedding latente.
• Receita de Treinamento: Proposta de uma estratégia de fine-tuning com LoRA e inicialização específica do encoder para adaptação eficiente a entradas comprimidas.
• Desempenho Superior: Evidência experimental de que o método atinge a fronteira de Pareto entre compressão e desempenho, reduzindo o comprimento da entrada em até 75% com degradação mínima de desempenho.
• Redução de Custo Computacional: Em cenários onde o comprimento de saída é muito menor que a entrada, a redução de 75% no comprimento de entrada resulta em uma redução estimada de 94% nos FLOPs totais (devido à natureza quadrática da atenção no pré-preenchimento).

4. Resultados Experimentais

Os experimentos foram realizados em três tarefas distintas usando o modelo base Qwen-2.5 0.5B:

1. Raciocínio Estrutural (Textualized Tree):
   • Tarefa: Identificar relações pai-filho em árvores de texto.
   • Resultado: Com $K=4$ (75% de redução), o modelo atingiu 98,38% de precisão (vs. 99,97% do baseline sem compressão). O método superou todos os baselines (Hard e Soft) no score combinado P-L F1.
2. Classificação de Sentimento (Amazon Reviews):
   • Tarefa: Classificar reviews como positivos ou negativos.
   • Resultado: O modelo com $K=4$ atingiu 91,05% de precisão com 75% de redução. O método superou o baseline mais forte em 25,5% no score P-L F1.
3. Edição de Código (CommitPackFT):
   • Tarefa: Gerar código modificado baseado em instruções.
   • Resultado: O método alcançou a melhor relação entre perplexidade e redução de comprimento. Com 75% de redução, o modelo manteve a capacidade de seguir instruções complexas e reutilizar variáveis corretamente (validado em estudo de caso).

Métrica Chave (P-L F1):
Os autores introduziram o P-L F1 (F1 harmônico entre Desempenho e Redução de Comprimento) para quantificar o trade-off. O K-Token Merging ocupou consistentemente a fronteira de Pareto ótima em todos os conjuntos de dados, superando métodos de compressão dura (que perdem informação) e métodos de compressão suave existentes (que não comprimem o suficiente).

5. Significado e Conclusão

O trabalho demonstra que a ineficiência no espaço de embedding latente é uma fonte subexplorada de redundância. Ao comprimir sequências de tokens em embeddings únicos antes da entrada no LLM, é possível:

• Reduzir drasticamente a memória e o custo computacional de inferência para contextos longos.
• Manter a integridade da informação em tarefas densas, onde a exclusão de tokens (métodos duros) seria fatal.
• Oferecer uma solução escalável que não sofre com a explosão de vocabulário.

Limitações e Futuro:
O método atual não comprime tokens gerados (apenas a entrada), e utiliza uma taxa de compressão fixa ($K$). Trabalhos futuros devem explorar compressão recursiva, estratégias adaptativas (variar $K$ baseado na densidade de informação) e a aplicação em modelos de maior escala.

Em suma, o K-Token Merging representa um avanço significativo na eficiência de LLMs, provando que a compressão no espaço latente é uma via promissora para viabilizar o processamento de contextos extremamente longos.