Resumo Técnico: Atenção de Janela Curta Habilita Memorização de Longo Prazo

Declaração do Problema

Modelos de Linguagem Grandes (LLMs) modernos baseados em atenção softmax dependem de um cache Chave-Valor (KV) que cresce linearmente com o comprimento da sequência, incorrendo em custos computacionais e de memória proibitivos para tarefas de contexto longo. Por outro lado, Redes Neurais Recorrentes (RNNs) lineares e Modelos de Espaço de Estado (SSMs) mantêm custos de memória e computação constantes, mas frequentemente sofrem com capacidades de recordação inferiores em comparação aos Transformers.

Arquiteturas híbridas recentes tentam combinar atenção de janela deslizante local (SWA) com atenção global ou camadas de RNN linear para equilibrar eficiência e desempenho. No entanto, o impacto do comprimento da janela deslizante na interação entre memória local e de longo prazo permanece pouco estudado. Especificamente, não está claro como o tamanho da janela influencia a capacidade do modelo de aprender dependências de longo prazo quando RNNs lineares estão presentes, e se janelas maiores melhoram inerentemente o desempenho ou prejudicam a capacidade do modelo de extrapolar para contextos mais longos.

Metodologia

Os autores investigam arquiteturas híbridas que intercalam camadas de Atenção de Janela Deslizante (SWA) com camadas de RNN linear, utilizando especificamente xLSTM (uma variante de mLSTM com memória matricial) como componente de memória de longo prazo. A arquitetura proposta, denominada SWAX, alterna entre camadas SWA e camadas xLSTM em uma proporção de 1:1.

O estudo foca em duas contribuições metodológicas principais:

1. Análise do Tamanho da Janela: Os autores treinam modelos SWAX com tamanhos de janela deslizante fixos variando de 128 a 2048 tokens. Eles avaliam esses modelos em benchmarks de raciocínio de contexto curto (por exemplo, ARC, Hellaswag, MBPP) e tarefas de recuperação de contexto longo (por exemplo, RULER Needle-In-A-Haystack).
2. Treinamento de Janela Estocástica: Para abordar o compromisso entre precisão de contexto curto (favorecida por janelas maiores) e extrapolação de contexto longo (favorecida por janelas menores), os autores propõem uma estratégia de treinamento onde o tamanho da janela deslizante é amostrado estocasticamente para cada lote. Em seus experimentos, o tamanho da janela alterna entre um tamanho curto (128) e um tamanho longo (2048) com uma probabilidade específica (0,5 para modelos de 1,4B, 0,75 para modelos de 7B). Além disso, eles empregam um cronograma de annealing onde a amostragem estocástica é removida nos últimos 10% do treinamento para estabilizar o desempenho de contexto curto.

Os experimentos são conduzidos em modelos com 1,4B e 7B parâmetros, treinados em 150 bilhões de tokens com um comprimento de sequência fixo de 16k durante o pré-treinamento.

Principais Contribuições

1. Descoberta Contra-Intuitiva sobre Tamanho da Janela: O artigo demonstra que, para arquiteturas híbridas que combinam SWA e RNNs lineares, janelas deslizantes maiores podem degradar o desempenho de contexto longo. Contrariamente à crença de que janelas maiores fornecem mais contexto, os autores descobrem que janelas grandes (por exemplo, 2048) permitem que o modelo dependa excessivamente da atenção softmax local precisa, fazendo com que as camadas de RNN linear subtreinem em dependências de longo prazo.
2. Mecanismo do Benefício da Janela Curta: Janelas mais curtas (por exemplo, 128) forçam o modelo a depender das camadas de RNN linear para recuperação de informações que caem fora da janela local. Esse "sinal de supervisão" obriga as camadas lineares a se especializarem na modelagem de dependências de longo prazo, melhorando assim a capacidade do modelo de extrapolar para comprimentos de sequência muito superiores à janela de treinamento.
3. Estratégia de Treinamento Estocástica: Os autores introduzem um procedimento de treinamento que varia estocasticamente o tamanho da janela. Essa abordagem força o modelo a aproveitar tanto a precisão de curto prazo da janela quanto a memória de longo prazo da RNN, atuando efetivamente como uma forma de dropout no mecanismo de atenção.

Resultados

• Desempenho de Contexto Longo: Modelos treinados com uma janela grande fixa (2048) exibem uma queda catastrófica no desempenho em tarefas de recuperação de contexto longo (por exemplo, RULER NIAH) quando testados em sequências mais longas que a janela. Em contraste, modelos treinados com janelas mais curtas (128) mantêm alta precisão mesmo em comprimentos de sequência de 131k tokens. O método de treinamento estocástico alcança desempenho comparável ou superior aos modelos apenas de janela curta nessas tarefas de contexto longo.
• Desempenho de Contexto Curto: Embora janelas curtas (128) sozinhas às vezes apresentem desempenho inferior em benchmarks de raciocínio de contexto curto em comparação com janelas maiores, o método de treinamento estocástico recupera esse desempenho. Tanto nas escalas de 1,4B quanto de 7B, o treinamento estocástico produz desempenho de contexto curto comparável, ou em alguns casos superior, ao de modelos treinados com uma janela fixa de 2048.
• Generalização: As descobertas mantêm-se em diferentes backbones de atenção linear (xLSTM e Gated DeltaNet) e diferentes configurações híbridas (alternando SWA com camadas de atenção completa). A estratégia de treinamento estocástico supera consistentemente as bases de janelas fixas em benchmarks de contexto curto e longo.

Significado e Alegações

O artigo afirma que a escolha do tamanho da janela deslizante em arquiteturas híbridas não é meramente um compromisso entre custo computacional e desempenho, mas um fator crítico em como o modelo aprende a distribuir responsabilidades de memória.

Os autores afirmam que:

• Janelas curtas são essenciais para treinar memória de longo prazo: Ao restringir o campo de atenção local, o modelo é forçado a utilizar suas camadas de RNN linear para dependências de longo alcance, levando a uma melhor extrapolação de comprimento.
• Treinamento estocástico resolve o compromisso: Ao alterar dinamicamente o tamanho da janela durante o treinamento, o modelo aprende a utilizar tanto a atenção local de alta fidelidade para contextos curtos quanto a memória linear global para contextos longos.
• Excesso de dependência na atenção local é prejudicial: Se um modelo é treinado com uma janela grande, ele falha em aprender os mecanismos necessários para recuperar informações além dessa janela, levando a uma generalização pobre em sequências mais longas.

O trabalho sugere que, para que arquiteturas híbridas alcancem desempenho de estado da arte em contexto longo sem custos computacionais ilimitados, o procedimento de treinamento deve ativamente impedir que o modelo dependa excessivamente de mecanismos de atenção local, mesmo que esses mecanismos sejam mais precisos para tarefas de curto alcance.