Resumo Técnico: Estendendo o Contexto de LLMs via Memória Recorrente Associativa

Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) são cada vez mais exigidos para processar entradas que abrangem centenas de milhares ou milhões de tokens para tarefas como análise de relatórios técnicos, desenvolvimento de software e raciocínio multi-documento. No entanto, as arquiteturas Transformer padrão enfrentam um gargalo fundamental: os custos computacionais e de memória da autoatenção escalam quadraticamente com o comprimento da sequência ($O(N^2)$). Além disso, o desempenho frequentemente degrada conforme o comprimento do contexto aumenta, mesmo dentro da janela nominal do modelo. Embora as arquiteturas recorrentes (ex: Mamba, RWKV) ofereçam escalonamento linear, elas geralmente exigem treinamento do zero, impedindo o aproveitamento de LLMs pré-treinados existentes, e muitas vezes têm dificuldade com tarefas algorítmicas complexas e seguimento de instruções em comparação aos Transformers.

Metodologia

Os autores propõem o Transformer de Memória Recorrente Associativa (ARMT) como uma solução prática para estender o comprimento do contexto mantendo a eficiência e o escalonamento de memória constante. O ARMT atua como um invólucro (wrapper) em torno de um LLM base pré-treinado, permitindo o processamento segmentado da entrada.

Arquitetura Central

O ARMT divide as entradas de contexto longo em segmentos não sobrepostos de comprimento fixo. Dentro de cada segmento, o modelo utiliza autoatenção total (memória de curto prazo/trabalho). Crucialmente, ele introduz um módulo de memória associativa por camada que propaga informações entre segmentos (memória de longo prazo). O mecanismo opera em três estágios:

1. Extração de Memória: Cada camada do transformer comprime um segmento de entrada em embeddings de memória.
2. Consolidação de Memória: Esses embeddings são consolidados em uma matriz associativa por camada como pares chave-valor.
3. Associação: Embeddings em segmentos subsequentes são transformados em vetores de consulta (query vectors) e multiplicados pela matriz associativa para recuperar informações relevantes de segmentos anteriores.

Receita de Treinamento

O artigo descreve uma estratégia de treinamento abrangente para adaptar LLMs pré-treinados ao ARMT:

• Pré-treinamento Continuado: Parâmetros de memória associativa não inicializados são inicializados via modelagem de linguagem não supervisionada em contextos longos (ex: 19B de tokens do FineWeb-Edu) para aprender a propagação de memória eficaz antes do ajuste fino (fine-tuning) específico da tarefa.
• Aprendizado por Currículo: Para abordar a dificuldade de aprender dependências de longo alcance do zero, o modelo é ajustado finamente aumentando progressivamente o número de segmentos (ex: de 2 para 4 para 8) e realizando o recozimento (annealing) da taxa de aprendizado.
• Geração de Dados Sintéticos: Para superar a escassez de dados em cenários de contexto longo, os autores geram instâncias de treinamento sintéticas concatenando passagens curtas de documentos longos e gerando pares de QA (pergunta e resposta) para cada uma, criando grandes grupos de comprimentos de contexto específicos.
• Poda e Seleção de Camadas: O estudo investiga se a memória associativa é necessária em todas as camadas. Ele propõe uma estratégia para manter a memória associativa apenas em um subconjunto de camadas (ex: camadas intermediárias e finais específicas), reduzindo os parâmetros treináveis e o custo computacional sem perda significativa de desempenho.

Contribuições Principais

1. Datasets Específicos de Domínio: A construção de dois novos conjuntos de dados, ManyTypes-long (MT) para previsão de tipos variáveis em código e GovReport-long (GR) para perguntas e respostas em documentos longos, projetados para avaliar cargas de trabalho realistas e de domínio restrito.
2. Receita de Treinamento: Um framework inovador para estender o contexto de LLM usando ARMT, combinando pré-treinamento continuado, geração de dados sintéticos, aprendizado por currículo e integração seletiva de camadas.
3. Validação Empírica: Um estudo experimental extensivo demonstrando que modelos aumentados com ARMT:
   • Processam entradas muito além dos limites originais de contexto (até 64k tokens) sem degradação de desempenho em relação às linhas de base dentro do limite.
   • Exibem generalização superior para comprimentos de contexto fora da distribuição (OOD) em comparação aos modelos base.
   • Requerem aproximadamente 30% menos FLOPs enquanto preservam o desempenho dentro da janela de contexto original.

Resultados Experimentais

Os autores avaliaram o ARMT utilizando os backbones Gemma-3-1B-IT e SmolLM-2-360M-IT nos datasets MT e GR.

• Desempenho: Os modelos ARMT mantiveram o desempenho estável em comprimentos de contexto de até 65k tokens. Em contraste, os modelos base (mesmo quando ajustados finamente) mostraram quedas acentuadas de desempenho além de suas janelas de contexto nativas (ex: 8k ou 32k). O ARMT superou significativamente os modelos base em regimes Long-OOD (32k–65k).
• Eficiência: O ARMT demonstrou uso constante de memória GPU independentemente do comprimento do contexto, enquanto o uso de memória dos modelos base cresceu linearmente. Para uma sequência de 32k tokens, o ARMT permitiu um aumento de 4x no tamanho do lote (batch size) (32 vs 8) sob o mesmo orçamento de memória.
• Redução de FLOPs: A análise teórica e as medições empíricas de tempo de inferência confirmaram uma redução nos FLFLOPs de atenção global por um fator de $T/S$ (comprimento da sequência dividido pelo tamanho do segmento), levando a uma redução total de aproximadamente 30% nos FLOPs totais em comparação com modelos de atenção total.
• Estudos de Ablação:
  • Poda de Camadas: Modelos com memória associativa em apenas ~20% das camadas (especificamente camadas intermediárias e finais pré-selecionadas) alcançaram desempenho comparável ou superior aos modelos ARMT completos.
  • Pré-treinamento: O pré-treinamento continuado foi considerado essencial para inicializar a memória associativa, melhorando significativamente o desempenho tanto em domínio quanto OOD.
  • Linhas de Base: O ARMT superou outras linhas de base de contexto longo, incluindo Mamba-2, DeltaNet e xLSTM, particularmente na generalização de contexto longo, enquanto exigiu menos pré-treinamento do que modelos treinados do zero.

Significância e Alegações

O artigo posiciona o ARMT como uma abordagem prática e eficiente de computação para permitir o processamento de contexto longo em LLMs de pequeno a médio porte (até 1 bilhão de parâmetros). Os autores alegam que esta abordagem preenche a lacuna entre o forte desempenho de contexto curto dos Transformers e o escalonamento linear dos modelos recorrentes.

As principais alegações de significância incluem:

• Privacidade e Implantação Local: Ao permitir o processamento eficiente de contexto longo em modelos menores, o ARMT facilita aplicações de preservação de privacidade que não dependem de APIs de LLM remotas e de grande escala.
• Escalabilidade: O método permite o processamento de contextos arbitrariamente longos com memória constante, um requisito crítico para análise de documentos do mundo real.
• Eficiência: A combinação de escalonamento de memória constante e redução de FLOPs torna o ARMT uma alternativa viável para ambientes com recursos limitados.
• Generalização: A receita de treinamento proposta aborda efetivamente o fenômeno "perdido no meio" (lost in the middle) e a degradação de desempenho observada em Transformers padrão ao lidar com sequências longas.

Os autores reconhecem limitações, observando que os experimentos foram restritos a modelos de até 1 bilhão de parâmetros e um conjunto específico de tarefas (código e QA de documentos), e que os mecanismos subjacentes da interação da memória associativa com as representações do transformer permanecem parcialmente compreendidos. No entanto, eles argumentam que a capacidade demonstrada de estender o contexto em modelos pequenos é um passo significativo em direção ao IA de contexto longo prática e local.