Maglev: Sliding Recurrent Memory

Resumo Técnico: Maglev – Memória Recorrente Deslizante

Problema
As arquiteturas Transformer atuais enfrentam um compromisso fundamental entre capacidade de memória e eficiência computacional. A autoatenção causal completa permite recuperação não linear poderosa e computação em contexto, mas incorre em custos que crescem linearmente com o comprimento do contexto. Por outro lado, a atenção de janela deslizante limita a atenção e os custos de cache KV, mas descarta informações distantes, falhando em fornecer uma memória persistente de longo prazo. Modelos recorrentes clássicos (ex: LSTMs) comprimem o histórico, mas propagam estados pequenos sequencialmente, limitando a capacidade e o paralelismo de sequência. Embora os Memory Transformers e os modelos de atenção linear tentem preencher essa lacuna, eles frequentemente impõem fronteiras de segmento artificiais, dependem de atualizações lineares estruturadas ou exigem o desenrolamento (unrolling) sequencial durante o treinamento. O desafio central é desenvolver um mecanismo prático de memória recorrente por token para Transformers que permita o pré-treinamento paralelo enquanto mantém custos de inferência limitados.

Metodologia: Arquitetura Maglev
O Maglev aborda isso desacoplando a construção da memória durante o treinamento da propagação da memória durante a inferência. A arquitetura consiste em dois modelos acoplados:

1. Prefiller ($Q$): Um Transformer causal que processa a sequência observada em paralelo usando um campo receptivo maior (ex: atenção intercalada total e de janela deslizante). Ele gera uma trajetória de vetores de memória alvo, $m'_t$, para cada posição de token.
2. Decoder ($P$): Um Transformer de janela deslizante que prevê o próximo token e gera seus próprios vetores de memória, $m_t$. Crucialmente, $P$ não tem acesso ao histórico completo; ele vê apenas uma janela limitada de tokens recentes e uma janela deslocada de vetores de memória.

Mecanismo de Treinamento: Treinamento Paralelo Elevado (Lifted Parallel Training)
Em vez de desenrolar o estado recorrente sequencialmente (o que impede o paralelismo), o Maglev emprega uma estratégia de treinamento de duas passagens:

• Passagem 1: O prefiller $Q$ gera a trajetória completa de memória alvo $m'_{1:T}$ em paralelo.
• Passagem 2: O decoder $P$ consome a trajetória alvo deslocada $m'_{0:T-1}$ juntamente com os tokens de entrada para prever o próximo token e produzir sua própria memória $m_{1:T}$.
• Objetivo: O modelo é otimizado usando uma função de perda composta:
  • Perda de Predição do Próximo Token ($L_{CE}$): Perda de entropia cruzada padrão para modelagem de linguagem.
  • Perda de Consistência de Memória ($L_{cons}$): Um termo que alinha a memória de saída do decoder $m_t$ com o alvo do prefiller $m'_t$. Isso ensina o decoder a reproduzir o estado de memória necessário para o próximo passo sem desenrolamento sequencial.

Mecanismo de Inferência
Na inferência, o prefiller $Q$ é descartado. O decoder $P$ executa de forma recorrente, alimentando suas próprias memórias geradas anteriormente $m_{t-1}$ de volta ao sistema.

• Injeção Recorrente de K/V: Em vez de usar tokens de memória dedicados, o Maglev injeta o estado de memória deslocado nas características de Chave/Valor (K/V) do mecanismo de atenção através de projeções aprendidas e mecanismos de gating.
• Complexidade Limitada: O decoder mantém uma janela deslizante de tamanho fixo de tokens e memórias. O estado recorrente é carregado adiante através do cache K/V, garantindo que os custos de inferência (computação e tamanho de cache) permaneçam limitados e independentes do comprimento total da sequência, correspondendo ao perfil de uma atenção de janela deslizante padrão.

Principais Contribuições

1. Recorrência Não Linear por Token: O Maglev introduz um mecanismo onde o estado do decoder de cada token serve como memória, atravessando toda a profundidade não linear do Transformer sem exigir tokens de memória extras ou fronteiras de segmento artificiais.
2. Treinamento Paralelo para Modelos Recorrentes: Ao usar um prefiller para gerar alvos em paralelo, o Maglev permite o pré-treinamento paralelo para arquiteturas recorrentes não lineares, evitando os gargalos sequenciais do treinamento tradicional de RNNs.
3. Compartilhamento de Parâmetros: Os autores demonstram que o prefiller $Q$ e o decoder $P$ podem compartilhar a maioria de seus parâmetros (especificamente os blocos Transformer) sem degradar significativamente o desempenho, reduzindo o overhead de memória de parâmetros do setup de treinamento.
4. Inferência Limitada: O modelo implantado retém a eficiência da atenção de janela deslizante enquanto ganha capacidades de memória persistente através da injeção recorrente de K/V.

Resultos Experimentais
Avaliado na arquitetura nanochat d20 (435M parâmetros, treinado em 43.52B tokens), o Maglev foi comparado contra Transformers de janela deslizante intercalada (total/deslizante), modelos puramente de janela deslizante e baselines de Latent Recurrent Transformer (LRT).

• Desempenho: A melhor variante Maglev (parâmetros separados, $\lambda=1$) melhorou os bits por byte (BPB) de validação do FineWeb-Edu de 0.7413 (baseline de janela deslizante) para 0.7251.
• Tarefas Downstream: Alcançou uma acurácia média em tarefas downstream de 56.4%, superando tanto o Transformer de janela deslizante correspondente (54.1%) quanto o baseline LRT correspondente.
• Peso de Consistência: O estudo observou que o peso de consistência ideal ($\lambda$) depende da estratégia de compartilhamento de parâmetros; um $\lambda$ mais alto beneficiou o modelo de parâmetros separados, mas pode restringir o modelo de parâmetros compartilhados.

Significância e Alegações
O artigo afirma que o Maglev fornece um caminho prático para equipar Transformers não lineares com memória persistente por token que pode ser reescrita após cada token. Ele consegue unir com sucesso o poder expressivo da atenção total e a eficiência dos modelos de contexto limitado. Ao separar a geração de memória do uso da memória durante o treinamento, o Maglev alcança o pré-treinamento paralelo para arquiteturas recorrentes, um feito anteriormente difícil para a recorrência não linear. Ao separar a geração de memória do uso da memória durante o treinamento, o Maglev alcança o pré-treinamento paralelo para arquiteturas recorrentes, um feito anteriormente difícil para a recorrência não linear. Os autores posicionam o Maglev como uma investigação preliminar neste espaço de design, observando que, embora melhore em relação aos baselines de janela deslizante e de recorrência latente testados, mais trabalho é necessário para escalar a abordagem, otimizar a eficiência de kernel para implantação e explorar mecanismos alternativos de injeção de memória (ex: injeção no fluxo residual ou atenção cruzada).