Resumo Técnico: PReM (Preserve and Refresh Memory)

1. Declaração do Problema

A inferência eficiente de contexto longo em Grandes Modelos de Linguagem (LLMs) enfrenta uma tensão fundamental entre a redução dos custos computacionais e a manutenção do acesso às evidências contextuais necessárias. À medida que o comprimento do contexto aumenta, o cache Key-Value (KV) cresce, aumentando significativamente a latência de decodificação e o uso de memória.

Abordagens existentes para mitigar isso, como compressão de cache KV e compressão de contexto, sofrem de duas limitações primárias:

1. Retenção de Evidência Estática: A maioria dos métodos toma uma decisão única e precoce sobre qual contexto reter. Essa compressão estática falha em se adaptar às necessidades evolutivas do raciocínio de múltiplos passos, onde diferentes passos de raciocínio podem exigir diferentes partes da evidência.
2. Seleção de Evidência Desacoplada: Métodos que dependem de compressores externos separam o processo de compressão da geração. Isso introduz computação adicional e impede que o modelo personalize a seleção de evidências dinamicamente para passos de raciocínio específicos.

O artigo argumenta que a compressão eficaz deve ir além da simples redução de tamanho; ela deve preservar dinamicamente a evidência útil e atualizá-la (refresh) quando passos de raciocínio posteriores exigirem informações diferentes.

2. Metodologia: PReM

PReM (Preserve and Refresh Memory) é um framework de compressão de contexto que mantém contextos longos como memória KV interna, camada por camada, dentro do LLM. Em vez de depender de compressores externos, o PReM aprende a decidir o que preservar e quando atualizar com base no estado interno do modelo.

Mecanismos Principais

• Representação de Memória KV Camada por Camada: O contexto longo é codificado como estados KV por camada ($M_\ell$) armazenados dentro do modelo. Esses estados são particionados em blocos contíguos.
• Camada de Memória Dedicada: Uma camada específica ($\ell_m$) é designada para tomar decisões de seleção de memória. Esta camada utiliza uma matriz de projeção de consulta de seleção dedicada ($W^c_\ell$) para pontuar blocos de contexto.
  • Seleção de Blocos: O modelo pontua os blocos com base na consulta atual (da pergunta ou de um token de atualização).
  • Preservar-e-Agrupar (Preserve-and-Pool): Os blocos selecionados são preservados com granularidade ao nível de token. Os blocos não selecionados são comprimidos em um único representante KV médio.
• O Token de Memória $\langle m \rangle$: Um token especial serve como um sinal interno explícito para atualizar a memória. Quando o modelo gera $\langle m \rangle$, a camada de memória reavalia os blocos de contexto, seleciona novas evidências relevantes e atualiza a memória KV comprimida para as camadas subsequentes. Isso permite que o modelo controle quando atualizar com base no passo de raciocínio atual.

Treinamento: Treinamento de Atualização com Separação de Fase

Para treinar o modelo para alinhar a seleção de memória com a geração, o PReM introduz uma estratégia de Treinamento de Atualização com Separação de Fase (Phase-Separated Refresh Training). Cada passo de raciocínio é dividido em duas passagens progressivas distintas:

1. Fase de Seleção de Memória: O modelo processa a entrada (pergunta ou $\langle m \rangle$) para calcular as pontuações dos blocos e selecionar os $k$ blocos de evidência superiores. Esta fase é supervisionada por uma Perda de Seleção de Memória ($L_{ms}$) que incentiva o modelo a classificar os blocos de evidência acima dos blocos de não-evidência.
2. Fase de Geração Condicionada à Memória: O modelo gera o segmento da resposta ($y_s$) condicionado ao estado de memória recentemente atualizado resultante da fase de seleção.

Para garantir a consistência entre treinamento e inferência, o objetivo de treinamento inclui termos de entropia cruzada de fronteira ($L_{bd}$). Esses termos conectam a fase de seleção ao primeiro token da resposta e o fim de um segmento de resposta ao próximo token $\langle m \rangle$, garantindo que o modelo aprenda os elos autorregressivos entre as atualizações de memória e a geração.

3. Principais Contribuições

1. Formulação: O artigo reformula a compressão de contexto como um problema aprendido de preservação e atualização de memória, eliminando a necessidade de compressores externos.
2. Arquitetura: O PReM introduz uma camada de memória dedicada e um token especial $\langle m \rangle$ para permitir atualizações de memória dinâmicas e específicas por passo dentro dos estados KV internos do LLM.
3. Estratégia de Treinamento: A proposta de Treinamento de Atualização com Separação de Fase, que desacopla a seleção de memória da geração durante o treinamento para garantir que o modelo aprenda a condicionar a geração ao estado de memória específico necessário para aquele passo.
4. Validação Empírica: Experimentos abrangentes demonstrando que atualizações de memória aprendidas e dinâmicas superam a compressão estática e o prompting de contexto total.

4. Resultados Experimentais

Experimentos foram conduzidos em variantes de 3B e 7B do Qwen2.5-Instruct usando contextos de 32K tokens em seis benchmarks (TriviaQA, SQuAD, NaturalQuestions, 2WikiMQA, HotpotQA, MuSiQue).

• Desempenho vs. Baselines: Sob razões de compressão de 16× e 32×, o PReM superou oito baselines fortes, incluindo métodos de cache KV (SnapKV, StreamingLLM, CAKE) e métodos de compressão de contexto (LongLLMLingua, EXIT, ICAE).
  • No backbone de 3B com compressão de 16×, o PReM melhorou a média de Exact Match (EM) e F1 em 5.30 e 5.15 pontos, respectivamente, sobre o baseline mais forte.
  • No backbone de 7B com compressão de 32×, as melhorias atingiram 8.20 (EM) e 9.89 (F1) pontos.
• Comparação ao Contexto Total: O PReM superou o prompting direto de 32K tokens (sem compressão) tanto nos modelos de 3B quanto nos de 7B, sugerindo que a seleção de memória aprendida é mais eficaz do que expor o modelo ao contexto completo e não filtrado.
• Superioridade de Modelos Pequenos: O modelo PReM de 3B superou baselines de compressão de contexto suave de 7B lançados (ex: Activation Beacon), indicando que a preservação e atualização aprendidas podem compensar escalas menores de backbone.
• Eficiência: O PReM reduziu significativamente o pico de memória da GPU (ex: ~19.57GB vs. 30.0GB para memória total) enquanto manteve uma latência de inferência competitiva. O uso de tokens $\langle m \rangle$ permite um trade-off controlável entre computação de decodificação (frequência de atualização) e economia de memória.

5. Significância e Alegações

O artigo afirma que o PReM demonstra que a inferência de contexto longo eficaz requer atualizações de memória acopladas à geração, em vez de apenas o encurtamento estático do contexto. Ao manter o contexto como memória KV interna e aprender a atualizá-la dinamicamente, o modelo pode acessar adaptativamente a evidência específica necessária para cada passo de raciocínio.

Os autores enfatizam que sua abordagem aborda as limitações da retenção estática e da seleção desacoplada, mostrando que um modelo pode aprender a "saber o que preservar e quando atualizar", levando a um desempenho superior em tarefas de raciocínio de múltiplos saltos (multi-hop) e a melhores trocas entre eficiência e qualidade em comparação com as técnicas de compressão existentes. O trabalho sugere que o gerenciamento de memória interna, treinado conjuntamente com a geração, é um caminho viável para escalar LLMs para contextos mais longos sem custos computacionais proibitivos.