Resumo Técnico: CoMem (Comprehension Memory)

Declaração do Problema

A modelagem de linguagem de contexto longo enfrenta dois gargalos primários: o custo quadrático de preenchimento (prefill) da autoatenção e o crescimento linear de memória do cache de Chave-Valor (KV). Os métodos eficientes existentes normalmente abordam isso através da expulsão ou compressão de estados KV ao longo do eixo de tokens (retendo um subconj conjunto de tokens) ou pela recuperação de unidades de texto externas. No entanto, essas abordagens frequentemente têm dificuldade em equilibrar a eficiência de memória com a fidelidade necessária para o raciocínio complexo sobre contextos longos. Além disso, muitos métodos dependem do treinamento de redes de memória especializadas ou da compressão de todo o contexto, o que pode introduzir um overhead significativo ou perda de fidelidade.

Os autores postulam que as camadas Transformer exibem uma divisão de trabalho: a informação semântica torna-se acessível em camadas inferiores e médias, enquanto as camadas superiores especializam cada vez mais as representações para a consulta (query) e predição atuais. Os métodos atuais não exploram plenamente essa estrutura baseada na profundidade para organizar a memória.

Metodologia: CoMem

O CoMem (Comprehension Memory) reorganiza a memória de contexto longo ao longo do eixo de profundidade em vez de apenas o eixo de tokens. Ele combina seleção de tokens limitada com reuso de computação particionada por profundidade.

Mecanismo Central

1. Fase de Escrita (Particionamento de Profundidade):

   • O contexto de entrada é transmitido em blocos disjuntos (ex: 512 tokens).
   • Cada bloco é processado apenas através de uma camada intermediária específica $j$ (a "profundidade de divisão" ou split depth).
   • Em vez de armazenar estados KV de profundidade total, o CoMem armazena o tensor residual $h_j$ (o estado oculto) na camada $j$ para cada token no bloco.
   • Esses tensores são armazenados em um índice externo indexado pelo texto do bloco.
   • Eficiência de Armazenamento: Para o Qwen3-8B, armazenar $h_j$ requer apenas 1/18 dos bytes de um par KV bf16 completo por token.

2. Fase de Leitura (Recuperação Limitada e Recomputação):

   • Ao receber uma consulta, um seletor (usando BM25) recupera um conjunto limitado de $k$ blocos relevantes de topo.
   • Os tensores residuais $h_j$ armazenados para esses blocos são empacotados junto com um token "sink" e a consulta atual.
   • O modelo recomputa apenas as camadas superiores $[j:L]$ sobre este pacote de tamanho fixo usando atenção cross-chunk completa.
   • Complexidade: Para um orçamento de recuperação ($k$) e tamanho de bloco fixos, o conjunto de trabalho de computação de leitura online e memória do lado do modelo são independentes do comprimento total do contexto armazenado.

Treinamento e Adaptação

• Backbone: O modelo base (Qwen3-8B) permanece congelado.
• Auto-destilação: Para mitigar a perda de fidelidade causada pelo salto de camadas inferiores durante a fase de leitura, um LoRA leve (Rank-32) é treinado via auto-destilação. O modelo estudante (na divisão $j=12$) aprende a corresponder aos logits de um modelo professor congelado (em $j=0$) no suporte de top-64 logits.
• Dados de Treinamento: O LoRA é treinado em texto PG-19 puro, sem supervisão de tarefas específicas ou dados de recuperação, permitindo que o modelo generalize a capacidade de "divisão de profundidade".

Principais Contribuições

1. Profundidade como Eixo de Reuso: O artigo estabelece que a acessibilidade semântica atinge o pico próximo à profundidade média, enquanto a "fronteira de leitura zero-shot" (onde estados armazenados podem ser usados diretamente sem recomputação) depende da escala do modelo. O CoMem trata essa divisão de profundidade ($j$) como um controle ajustável entre o trabalho pré-computado e o custo no momento da consulta.
2. Memória Particionada por Profundidade Limitada: O CoMem armazena um estado residual por token e retoma exatamente a parte superior do transformer. Ele combina a recuperação limitada com atenção cross-chunk completa, alcançando uma pegada de armazenamento de 1/18 do cache KV completo por token no Qwen3-8B.
3. Trade-off Controlado entre Qualidade e Eficiência: Os autores fornecem evidências de que o armazenamento em profundidade maior reduz a computação por consulta, mas acarreta uma perda de fidelidade, que a auto-destilação repara substancialmente. O método isola os benefícios da recuperação dos custos do reuso de profundidade.

Resultados Experimentais

Avaliado em um protocolo unificado, livre de chat-template, utilizando o Qwen3-8B:

• Desempenho em Benchmarks:
  • RULER: O CoMem atinge 97.05, superando a linha de base KV-Direct de contexto total (78.80).
  • LoCoMo (Memória de Diálogo): O CoMem pontua 38.27 contra 34.59 do KV-Direct. Essa vantagem persiste sob reamostragem de cluster de conversação e auditorias de juiz independentes.
  • LongEval & LongBench: O CoMem iguala ou supera as linhas de base (69.0 no LongEval; 12.15 macro F1 no LongBench).
• Eficiência (NVIDIA H20 com contexto de 128k):
  • Memória: O CoMem utiliza 18.26 GB de memória de pico comparado a 89.36 GB para o KV-Direct de contexto total.
  • Velocidade: O CoMem alcança um ganho de velocidade de 7.83× no tempo de prefill incluindo a escrita total.
  • Decode: O decode de banda retomada permanece dentro de 11% da inferência densa até 64k e torna-se 1.07× mais rápido em 128k.
• Ablações:
  • Recuperação vs. Profundidade: A seleção limitada sozinha (sem reuso de profundidade) melhora a memória de diálogo, mas o reuso de profundidade é crítico para a eficiência.
  • Destilação: A auto-destilação recupera ~13.75 pontos de LoCoMo perdidos pela divisão congelada em $j=12$.
  • Atenção Cross-Chunk: A atenção cross-chunk completa é essencial para a desambiguação de múltiplos fatos; o reuso de bloco-diagonal causa quedas significativas de desempenho.

Significância e Alegações

O artigo afirma que a memória de contexto longo pode ser efetivamente organizada ao longo do eixo de camada, não apenas o eixo de tokens. O CoMem demonstra que:

• A informação semântica é acessível em profundidades intermediárias, permitindo o armazenamento de estados independentes de query que são posteriormente especializados pela recomputação das camadas superiores.
• Uma estratégia de recuperação limitada combinada com computação particionada por profundidade pode desacoplar o custo de inferência online do comprimento do contexto armazenado.
• A auto-destilação leve pode reparar a perda de fidelidade inerente ao saltar as camadas inferiores, permitendo alta precisão sem o retreinamento do backbone.

Os autores enfatizam que o CoMem é um método de eficiência que oferece um trade-off qualidade-memória-computação, e não uma solução de precisão estado-da-arte para cada tarefa dentro da janela. Eles declaram explicitamente que seus resultados em contextos que excedem a janela nativa (40.960 tokens) são "referências de estresse não estendidas" e não reivindicam superioridade sobre modelos que utilizam técnicas de extensão de posição como o YaRN. O método é apresentado como uma abordagem complementar aos métodos existentes de compressão de eixo de tokens e recuperação.