Source-linked AI summary
Maglev: Sliding Recurrent Memory
Bo Liu, Qiang Liu
TL;DR
Training nonlinear recurrent Transformers in parallel while retaining bounded recurrent memory at inference is challenging. Maglev uses a causal prefiller to create parallel memory targets for a sliding-window decoder, improving validation BPB and downstream accuracy over matched baselines.
Problem
The central question is how to preserve token-wise recurrent memory at deployment without sequentially unrolling it during pretraining.
Method
Maglev trains a sliding-window decoder to reproduce causal prefiller-generated memory targets through a consistency objective, enabling bounded-memory recurrent inference.
Results
0.7251 FineWeb-Edu validation BPB from 0.7413 and 56.4 average downstream accuracy from 54.1 were achieved over a matched sliding-window Transformer.
Takeaways & Limitations
Maglev provides a way to train nonlinear recurrent models with fixed inference memory while preserving parallel training.
Takeaways & Limitations
The compute-constrained experiments remain a preliminary rather than definitive study of this design space.
Abstract
from arXiv · showhide
We introduce \ours{}, a recurrent Transformer architecture with fixed-size memory that generalizes sliding-window attention while remaining parallelizable during training. \ours{} consists of two coupled models: a prefiller $Q$, which leverages full attention\footnote{In practice, we use interleaved full and sliding-window attention for $Q$, as this yields stronger performance. The essential requirement is that $Q$ be more expressive than $P$, with access to the full history.} to produce memory targets $m'_t$, and a decoder $P$, which uses only sliding-window attention and recurrent K/V injection to produce decoder memories $m_t$ for next-token prediction. We train \ours{} with a memory consistency loss that aligns $m_t$ with $m'_t$, allowing inference to use $P$ alone. Empirically, \ours{} improves validation loss and downstream pretraining benchmarks over sliding-window and latent recurrent transformer baselines. Moreover, sharing parameters between $P$ and $Q$ reduces parameter memory while preserving most of the gains.
1 Introduction
Maglev enables a nonlinear Transformer to maintain token-wise recurrent memory with bounded sliding-window inference while preserving parallel pretraining. It separates parallel memory construction by a stronger prefiller from recurrent memory propagation by a decoder, yielding improvements over matched sliding-window and latent recurrent Transformer baselines.
- Motivation: Maglev addresses the lack of persistent long-term memory in sliding-window attention by giving a nonlinear Transformer token-wise latent memory that can be rewritten after every token.Each update traverses the full nonlinear Transformer depth rather than using only structured linear or affine updates.
- Architecture: The architecture keeps token and memory windows fixed, so recurrent memories carry information beyond the local KV cache without increasing inference size.The deployed decoder augments sliding-window attention with a gated K/V pathway for shifted recurrent memory.
- Training: Maglev restores parallel pretraining by using a stronger causal prefiller Q to generate target memories in parallel, while decoder P predicts tokens and learns to reproduce them through consistency loss.Training uses two sequence-parallel passes; at inference, Q is removed and P closes the loop with its own memories.
- Results: 0.7251 FineWeb-Edu validation BPB, versus 0.7413 for a matched sliding-window Transformer, was achieved by the best 435M-scale deployed decoder trained on 43.52B tokens.The same model improved average downstream accuracy from 54.1 to 56.4 over the matched sliding-window Transformer.
- Results: 56.4 average downstream accuracy, versus 54.1 for the matched sliding-window Transformer, accompanied performance gains while also outperforming a matched latent recurrent Transformer.Sharing most parameters between Q and P preserved most of the gains while reducing parameter memory.
2 Method
Maglev trains a causal prefiller Q and deployable decoder P in two sequence-parallel passes, using prefiller memories as decoder targets. At inference, Q is discarded while P recurrently feeds back its own memories through bounded sliding-window attention.
- Two-pass training: Q constructs a complete causal auxiliary memory trajectory in parallel, while P learns to reproduce it one step at a time.Q may have a larger receptive field than P but remains causal, preventing future information from entering its memories.
- Two-pass training: Maglev replaces a length-T recurrent training unroll with two sequence-parallel passes while retaining recurrent inference.All shifted prefiller targets are available before the decoder pass, allowing P to evaluate every position simultaneously.
- Training objective: The decoder optimizes next-token prediction and consistency between decoder memory m_t and prefiller target m′_t.This teaches P to produce the memory needed for the next window, enabling decoder memories to replace prefiller memories during inference.
- Prefiller and decoder architecture: The prefiller uses SLSL with a 512-token sliding window and full attention over the 2048-token training context, whereas the decoder uses SSSS with W = 512.Thus Q has broader context for constructing targets, while P preserves the bounded receptive field needed for deployment.
- Recurrent K/V injection: Recurrent K/V injection mixes shifted memory features with local token features, exposing the memory window without adding sequence positions.During training the memory is m′_j−1; during inference it is m_j−1, and the resulting decoder state m_t serves as both recurrent memory and language-model input.
- Inference: At inference, P appends its memory to the recurrent window, matching ordinary sliding-window attention’s cache size and attention cost independent of sequence length.Q is discarded, and the W cached entries represent both token and memory histories.
3 Related Work
Prior work addresses long-context modeling through efficient attention, fixed-size recurrence, hybrid architectures, online memory, and latent computation. Maglev combines nonlinear recurrent memory with parallel prefiller training, distinguishing it from approaches that directly unroll or refine recurrence.
- Efficient and long-context attention: Long-context Transformers reduce full-attention cost through local attention, compressed memory, recurrence, retrieval, or hybrid memory mechanisms.Transformer-XL and Compressive Transformers reuse or compress segment activations, while Infini-attention and landmark methods provide compressed or random-access memory.
- Linear recurrent and convolutional sequence models: Linear attention, state-space, convolutional, and gated recurrent models replace full key-value caches with fixed-size states or efficient sequence operators.Examples include H3, S5, Hyena, Mamba, RetNet, RWKV, Griffin/Hawk, gated delta networks, and xLSTM.
- Hybrid attention and recurrence: Hybrid models combine local attention with recurrent or state-space layers, preserving access to recent tokens while compressing older context.Feedback Transformers, TransformerFAM, Recurrent Memory Transformers, and Block-Recurrent Transformers expose past representations, memory tokens, or block states to future computation.
- Online and test-time memory: Online-memory research and parallel recurrent computation motivate Maglev, whose nonlinear recurrent update avoids sequential training through a parallel prefiller.This differs from directly unrolling the decoder recurrence during training.
- Latent recurrence and extra computation: Latent-recurrence methods add computation by looping blocks, inserting auxiliary tokens, or refining previous-token memory states through parallel passes.LRT trains recurrence through parallel refinement passes rather than Maglev’s parallel prefiller formulation.
4 Experiments
Experiments evaluate Maglev in a controlled nanochat pretraining setup against full/sliding-window, sliding-window, and latent recurrent baselines. Maglev improves fixed-window performance, with separate prefiller capacity yielding the strongest reported results while consistency-weight effects depend on parameter sharing.
- Experimental setup: All models use the same tokenizer, data pipeline, optimizer family, sequence length, and evaluation scripts, isolating architecture-specific differences.The setup uses the nanochat d20 architecture with 20 layers, width 1280, maximum sequence length 2048, and short-window size W = 512.
- Models: The comparison includes SLSL full/sliding-window, SSSS sliding-window, LRT variants, and shared- or separate-parameter Maglev models.Maglev uses an SLSL prefiller Q and SSSS decoder P, with consistency weights λ ∈ {0.1, 1.0}.
- Evaluation: Evaluation reports FineWeb-Edu validation BPB, LAMBADA perplexity and accuracy, and downstream accuracy across PIQA, HellaSwag, WinoGrande, ARC, SocialIQA, and BoolQ.The average column is the mean of the reported downstream accuracies.
- Results: 0.7295 FW BPB and 56.2 average downstream accuracy are achieved by shared Maglev with λ = 0.1, improving over SSSS and the corresponding LRT baseline.Maglev improves the fixed-window decoder without relying on full attention at inference.
- Results: 0.7251 FW BPB and 56.4 average downstream score are achieved by the separate-parameter variant with λ = 1, the best reported results.The comparison suggests that additional prefiller capacity can improve the consistency target.
- Results: A larger λ improves the separate-parameter model but hurts the shared model on several downstream tasks, so consistency weight effects are not purely monotone.An overly strong consistency term can constrain shared decoder representations, whereas separate prefiller capacity can provide a more useful training signal.
5 Conclusion and Future Work
Maglev is a fixed-memory recurrent Transformer trained with prefiller–decoder consistency, enabling bounded-memory recurrent inference without the prefiller. In nanochat d20 experiments, it improves validation BPB and downstream pretraining benchmarks over sliding-window and LRT baselines, while scaling and deployment efficiency remain open questions.
- Conclusion: Maglev trains a decoder to produce recurrent memories matching parallel targets from a more expressive prefiller, then removes the prefiller at inference.The decoder uses sliding-window attention with recurrent K/V injection and runs as a bounded-memory recurrent model.
- Conclusion: Maglev improves validation BPB and downstream pretraining benchmarks over sliding-window and LRT baselines in the nanochat d20 setting.The experiments establish these gains in a preliminary study constrained by available compute.
- Future Work: Future work includes studying prefiller strength versus decoder capacity, efficient recurrent-injection kernels, and using pretrained or lightly fine-tuned models to provide memory targets.The paper frames Maglev as a preliminary rather than definitive investigation of the design space.