Resumo Técnico: Compressão de Contexto Ponta a Ponta em Escala

Declaração do Problema

A inferência de contexto longo em Grandes Modelos de Linguagem (LLMs) é cada vez mais limitada pela memória e latência devido ao crescimento linear do cache Key-Value (KV) com o comprimento do contexto. Embora as técnicas recentes de compressão de cache KV tentem mitigar isso desalojando entradas do cache ou treinando caches compactos, elas enfrentam limitações significativas:

1. Compromisso entre Qualidade e Eficiência: Muitos métodos degradam substancialmente o desempenho do modelo ou exigem treinamento específico para o domínio que não se generaliza.
2. Compatibilidade com Motores de Inferência: Técnicas que desalojam entradas de forma não uniforme entre cabeças de atenção ou camadas são difíceis de integrar com motores de inferência de produção modernos (ex: vLLM, SGLang), que assumem comprimentos de sequência compartilhados.
3. Sobrecarga de Prefill: Muitas abordagens exigem que o contexto completo seja preenchido antes da compressão, o que anula os ganhos potenciais de latência, ou exigem tempo de computação adicional substancial.
4. Limitações de Soft-Tokens: Os métodos existentes de "soft-tokens" encoder-decoder frequentemente dependem de pré-processamento offline, falham em preservar o comportamento amplo de in-context do modelo base, ou são avaliados apenas em tarefas de domínio específico estreito, em vez de benchmarks de contexto longo densos em informação.

A questão central abordada é se um compressor geral e agnóstico à tarefa pode ser treinado para preservar as capacidades originais de um modelo enquanto reduz significamente o uso de memória e o tempo de inferência.

Metodologia

Os autores introduzem os Latent Context Language Models (LCLMs), uma família de compressores encoder-decoder treinados de ponta a ponta. A arquitetura central mapeia uma sequência de tokens de entrada $x_{1:T}$ em uma sequência mais curta de embeddings latentes $z_{1:M}$ através de um encoder, que são então projetados por um adaptador e consumidos por um decoder no lugar do contexto original.

Design da Arquitetura

Através de uma busca de arquitetura em larga escala do zero, os autores identificaram escolhas de design ideais:

• Pooling: O mean pooling (agrupamento pela média) mostrou-se superior ao pooling baseado em tokens (EOS/CLS) e competitivo com a concatenação, particularmente em altas taxas de compressão (16×).
• Granularidade de Codificação: Um tamanho de janela do encoder ($W$) de 1024 tokens foi adotado, permitindo que o encoder atenda a um contexto local maior do que o bloco de compressão ($N$) sem sobrecarga proibitiva de memória.
• Mascaramento de Atenção: O mascaramento causal dentro do encoder superou consistentemente o mascaramento bidirecional.
• Adaptador: Um adaptador MLP leve (sem autoatenção) provou ser mais eficaz e eficiente do que adaptadores baseados em atenção para projetar os latentes do encoder para a dimensão do decoder.
• Inicialização: Inicializar o encoder a partir de um modelo de embedding pré-treinado (Qwen3-Embedding-0.6B) em vez de um LLM padrão forneceu melhores representações iniciais, embora a lacuna tenha diminuído com treinamento extensivo.

Receita de Treinamento

Os autores desenvolveram uma receita de treinamento em múltiplos estágios para evitar o esquecimento catastrófico e garantir estabilidade:

1. Estágio 0 (Warmup do Adaptador): Congelar encoder e decoder; treinar apenas o adaptador.
2. Estágio 1 (Treinamento do Encoder): Descongelar o encoder; manter o decoder congelado.
3. Estágio 2 (Pré-treinamento Contínuo): Descongelar o decoder com uma taxa de aprendizado pequena; treinar de ponta a ponta em uma mistura de predição de próximo token e dados de reconstrução.
4. Estágio 3 (Ajuste Fino Supervisionado - SFT): Treinar em dados de SFT para aumentar o seguimento de instruções e o raciocínio com entradas comprimidas.

Os dados de treinamento incluem uma mistura cuidadosamente selecionada de código, texto, raciocínio e documentos de contexto longo. Crucialmente, os dados de pré-treinamento usam um formato intercalado onde segmentos comprimidos (envoltos em tokens especiais) alternam com segmentos não comprimidos, forçando o modelo a condicionar-se em contexto latente em múltiplas posições, em vez de apenas no início.

Integração de Agente

Os autores propõem um sistema de agente onde o agente LCLM recebe todo o contexto comprimido. O agente pode invocar seletivamente uma ferramenta EXPAND para recuperar o texto original não comprimido para segmentos específicos, permitindo uma estratégia de "escanear e dar zoom" (skim-and-zoom) para tarefas de longo horizonte.

Resultados Principais

A família LCLM (especificamente um encoder de 0.6B e um decoder de 4B) foi avaliada em taxas de compressão de 4×, 8× e 16× em benchmarks incluindo RULER, LongBench, LongHealth e GSM8K.

• Fronteira de Pareto: Os LCLMs estabelecem uma nova fronteira de Pareto na relação entre precisão e eficiência. Eles alcançam um Time-to-First-Token (TTFT) significativamente mais rápido e menor pico de memória GPU em comparação com baselines de compressão de cache KV (ex: SnapKV, KVzip, Attention Matching), especialmente em altas taxas de compressão.
• Preservação de Precisão: Em compressão de 16×, os LCLMs mantêm alta precisão em benchmarks de contexto longo (ex: ~65% no RULER 16k vs. ~93% para KV total, superando significativamente os baselines que frequentemente caem abaixo de 20-40%).
• Eficiência de Memória: Ao contrário dos métodos de cache KV onde a memória escala com o comprimento total do contexto, o uso de memória do LCLM estabiliza em contextos longos porque o encoder processa janelas de tamanho fixo. Isso permite que os LCLMs lidem com contextos de até 1M de tokens em uma única GPU H200, enquanto os baselines esgotam a memória em 512K ou 1M de tokens.
• Desempenho de Grão Fino: No GSM8K (contextos curtos e densos), os LCLMs alcançam a maior precisão em todas as taxas de compressão, demonstrando que não estão limitados apenas a QA de contexto longo, mas podem comprimir entradas densas de informação efetivamente.
• Desempenho do Agente: O sistema de agente com expansão adaptativa melhora significativamente a precisão de correspondência exata de strings em tarefas de "agulha no palheiro" (needle-in-a-haystack), muitas vezes igualando o desempenho do modelo não comprimido enquanto retém a eficiência do contexto comprimido.

Significância e Alegações

O artigo afirma que os LCLMs oferecem uma solução prática e escalável para a inferência de contexto longo que é compatível com motores de inferência de produção padrão. Ao deslocar o fardo da compressão para um encoder menor e paralelizável, os LCLMs desacoplam os custos de memória e latência do comprimento do contexto do processo de geração do decoder.

Os autores enfatizam que sua abordagem:

1. Preserva Capacidades Gerais: Diferente dos métodos anteriores de soft-tokens que degradam o desempenho ou exigem ajuste específico para tarefas, os LCLMs mantem um forte desempenho em tarefas gerais, raciocínio e seguimento de instruções.
2. Permite Recuperação Adaptativa: A integração com frameworks de agentes permite a expansão dinâmica do contexto relevante, preenchendo a lacuna entre a visibilidade global do corpus e a recuperação de informação de grão fino.
3. Escalabilidade: O método escala efetivamente para contextos de um milhão de tokens em hardware comum (GPUs H200), superando os gargalos de memória que atualmente limitam a implantação de sistemas de LLM de longo horizonte.

O trabalho sugere que a compressão aprendida de ponta a ponta é uma alternativa viável ao desalojamento heurístico de cache KV, fornecendo um substrato robusto para futuros agentes de longo horizonte com memórias de trabalho persistentes.