Resumo Técnico: Compactação Rápida de KV via Correspondência de Atenção

1. Declaração do Problema

A escalabilidade de modelos de linguagem grandes (LLMs) para contextos longos está atualmente limitada pelas exigências de memória do cache de Chave-Valor (KV). Em arquiteturas Transformer autoregressivas, o cache KV deve reter chaves e valores para todos os tokens anteriores, frequentemente atingindo gigabytes por solicitação em cenários de longo horizonte (por exemplo, diálogos multi-sessão, codificação com agentes).

As abordagens existentes para reduzir o tamanho do cache KV geralmente se enquadram em duas categorias, ambas com limitações significativas:

• Métodos no espaço de tokens: Técnicas como evicção de tokens (H2O, SnapKV), fusão de tokens (KVMerger) e sumarização degradam-se rapidamente em altas taxas de compressão. A sumarização, em particular, é altamente perdedora e prejudica o desempenho downstream em tarefas que exigem extração precisa de informações.
• Métodos no espaço latente: Trabalhos recentes como Cartridges (Eyuboglu et al., 2025) demonstraram que caches KV compactos podem ser treinados no espaço latente para corresponder ao desempenho de contexto completo. No entanto, essa abordagem depende de otimização baseada em gradiente de ponta a ponta (prefix-tuning) para cada contexto, o que é proibitivamente lento e caro (requerendo várias horas de GPU por contexto).

O problema em aberto é alcançar compactação de contexto rápida e de alta qualidade no espaço latente que preserve o comportamento do modelo downstream sem o custo computacional da otimização de gradiente por contexto.

2. Metodologia: Correspondência de Atenção

Os autores propõem uma família de métodos baseada em Correspondência de Atenção, que constrói um cache KV compacto $(C_k, C_v)$ para reproduzir as saídas de atenção e a massa de atenção do cache original $(K, V)$ para um conjunto de consultas de referência.

Objetivo Central

O objetivo é encontrar um cache compacto $(C_k, C_v)$ de tamanho $t \ll T$ tal que, para qualquer consulta $q$ e quaisquer tokens fixos/futuros anexados $(K_{fixed}, V_{fixed})$, a saída de atenção permaneça aproximadamente preservada. Os autores aproveitam a decomposição da atenção sobre blocos concatenados, mostrando que basta corresponder duas quantidades sobre um conjunto de consultas de referência $Q_{ref}$:

1. Saída de Atenção Local: A soma ponderada dos valores.
2. Massa de Atenção: A soma dos logits de atenção não normalizados (o denominador no softmax).

Para lidar com a discrepância onde o cache compactado possui menos tokens ($t < T$) e, portanto, não pode corresponder naturalmente à massa total de atenção (por exemplo, para $q=0$, a massa é $T$ versus $t$), o método introduz um viés escalar por token $\beta$. Os alvos de otimização são:
$$\frac{\exp(qK^\top)V}{\sum \exp(qK^\top)} \approx \frac{\exp(qC_k^\top + \beta)C_v}{\sum \exp(qC_k^\top + \beta)}$$
$$\sum \exp(qK^\top) \approx \sum \exp(qC_k^\top + \beta)$$

Pipeline Algorítmico

O framework decompõe o problema em subproblemas eficientes e de forma fechada, evitando a descida de gradiente no momento da compactação:

1. Geração de Consultas de Referência ($Q_{ref}$):

   • Repeat-prefill: Executar uma passagem de prefill no próprio contexto (ou em uma versão repetida) para extrair ativações de consulta.
   • Self-study: Gerar interações sintéticas condicionadas ao contexto para ampliar a distribuição de consultas.
   • On-policy: Compactar camadas sequencialmente e extrair consultas do modelo parcialmente compactado para reduzir o deslocamento de distribuição.

2. Seleção de Chaves ($C_k$):

   • Chaves de Maior Atenção: Selecionar chaves originais que recebem as maiores pontuações de atenção entre as consultas de referência (agregadas via RMS).
   • Pursuit de Correspondência Ortogonal (OMP): Selecionar chaves de forma gananciosa para melhor satisfazer o objetivo de correspondência de massa de atenção, seguido por Mínimos Quadrados Não Negativos (NNLS) para ajustar os pesos.

3. Ajuste de Viés e Valores ($\beta, C_v$):

   • Ajuste de $\beta$: Resolvido via Mínimos Quadrados Não Negativos (NNLS) para corresponder à massa de atenção das chaves originais.
   • Ajuste de $C_v$: Resolvido via Mínimos Quadrados Ordinários (OLS) para corresponder às saídas de atenção dadas as chaves selecionadas e os vieses ajustados.

Compactação Não Uniforme

Os autores observam que diferentes cabeças de atenção exibem sensibilidades variadas à compactação. Eles propõem um cronograma de compactação não uniforme onde diferentes cabeças e camadas recebem diferentes taxas de retenção. Esse cronograma é pré-calculado uma vez por modelo usando um algoritmo de troca ganancioso baseado em curvas de sensibilidade por cabeça, permitindo aplicação robusta em diferentes contextos.

Compactação em Blocos (Chunked)

Para suportar contextos muito longos (por exemplo, 60k+ tokens), o método aplica compactação independentemente a blocos contíguos da entrada. A abordagem de fragmentação baseada em KV (preenchimento do contexto completo, fatiamento dos estados KV, compactação e mesclagem) é usada por padrão para preservar semânticas posicionais globais e interações entre blocos.

3. Contribuições Principais

• Formulação de Correspondência de Atenção: Um objetivo inovador para compactação no espaço latente que otimiza diretamente a preservação do comportamento de atenção em vez de verossimilhanças de saída, decompondo-se naturalmente em problemas de álgebra linear eficientes.
• Soluções de Forma Fechada: O método elimina a necessidade de otimização baseada em gradiente de ponta a ponta durante a compactação, substituindo-a por solucionadores de mínimos quadrados e NNLS.
• Fronteira de Pareto Velocidade-Qualidade: A abordagem alcança taxas de compactação de até 50× em segundos (em uma única GPU), ordens de magnitude mais rápido que Cartridges, mantendo precisão downstream comparável ou superior.
• Agendamento Não Uniforme: Uma estratégia de alocação de orçamento pré-calculada e específica do modelo que prioriza capacidade para cabeças de atenção mais sensíveis à compactação.
• Mecanismo de Viés: A introdução de vieses escalares por token para corrigir discrepâncias de massa de atenção, permitindo correspondência exata de massa com um conjunto reduzido de chaves.

4. Resultados

Os métodos foram avaliados em QuALITY (QA de documentos longos), LongHealth (registros clínicos densos), QASPER (QA generativa) e LongBench v2.

• Desempenho vs. Velocidade: Em QuALITY com uma taxa de compactação de 50×, os métodos de Correspondência de Atenção traçam a fronteira de Pareto, superando as linhas de base de evicção de tokens (H2O+, KVzip) e sumarização. Eles igualam ou excedem a precisão de Cartridges enquanto são duas ordens de magnitude mais rápidos (minutos versus horas).
• Alta Compressão: Em taxas extremas (até 100×), os métodos retêm precisão downstream significativa (por exemplo, ~60% em LongHealth versus ~30% para linhas de base de sumarização).
• Generalização: A abordagem generaliza para tarefas generativas (QASPER) e benchmarks intensivos em recuperação (RULER), onde a variante mais rápida (AM-HighestAttnKeys-fast) superou o estado da arte KVzip no RULER.
• Eficiência: A compactação de um contexto de 60k tokens em Gemma-3-12B leva aproximadamente 100–600 segundos, dependendo da estratégia de seleção de chaves (OMP versus Maior Atenção), dominada pelas etapas de geração de consultas e seleção de chaves.
• Compactação Online: Resultados preliminares em AIME mostram que a compactação repetida no meio da trajetória preserva o desempenho de raciocínio, desacoplando a profundidade de raciocínio dos limites físicos de memória.

5. Significado e Alegações

O artigo afirma que a Correspondência de Atenção fornece um primitivo prático e eficiente para compactação de cache KV que preenche a lacuna entre métodos rápidos e perdedores no espaço de tokens e métodos lentos e de alta qualidade no espaço latente.

• Praticidade: Ao evitar a descida de gradiente por contexto, o método torna a compactação de alta fidelidade viável para implantação no mundo real, onde latência e custo computacional são restrições.
• Integração de Sistema: A abordagem é compatível com motores de inferência existentes (por exemplo, FlashAttention, empacotamento de KV de comprimento variável) e pode ser realizada fora do caminho crítico (compactação desagregada).
• Limitações: Os autores reconhecem que, embora seu método seja significativamente mais rápido que Cartridges, ainda é mais lento que heurísticas simples. Além disso, em taxas de compressão extremas (por exemplo, 100× em LongHealth), métodos baseados em gradiente como Cartridges podem ainda superar devido à sua capacidade de buscar um espaço de representação mais amplo não restrito às chaves originais.
• Direções Futuras: O artigo sugere trabalhos futuros na otimização da geração de consultas, indo além da seleção de subconjuntos para otimização direta de chaves compactas, e integrando a compactação em motores de inferência para configurações de agentes online e de longo horizonte.

Em resumo, este trabalho estabelece que a compactação de contexto rápida e de alta qualidade é alcançável ao corresponder diretamente o comportamento de atenção, oferecendo uma solução escalável para o gargalo de memória em LLMs de contexto longo.