🤖 machine learning

Maglev: Sliding Recurrent Memory

The paper introduces Maglev, a recurrent Transformer architecture that combines a full-attention prefiller with a sliding-window decoder trained via a memory consistency loss to achieve fixed-size memory, parallelizable training, and superior performance over existing sliding-window and latent recurrent baselines.

Original authors: Bo Liu, Qiang Liu

Published 2026-08-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bo Liu, Qiang Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Memory Problem: Why AI Needs a Better Notebook

Imagine you are trying to write a story, but every time you finish a sentence, you have to forget everything you wrote before it. That is the basic struggle of many modern artificial intelligence models. These models, known as Transformers, are incredibly good at understanding language, but they have a tricky memory strategy. The "perfect" way for them to remember is to keep every single word they've ever seen in front of them, like a giant, open notebook. This lets them connect ideas from the beginning of a book to the end, but it gets so heavy and slow that the computer runs out of energy and time if the story gets too long.

To fix this, engineers often use a "sliding window" approach. Imagine the AI only keeps the last few sentences in its notebook and throws away the rest. This makes the computer fast and efficient, but it means the AI loses the ability to remember important details from earlier in the story. On the other hand, older types of AI models tried to compress their whole history into a tiny, single summary note. This is fast, but the summary often gets messy and loses the rich details needed to understand complex plots. The big question in this corner of computer science is: Can we build an AI that has the speed of the sliding window but the deep, rich memory of the full notebook, without getting bogged down?

Maglev: The Train That Remembers Without the Baggage

This paper introduces a new architecture called Maglev (short for Magnetic Levitation), which acts like a high-speed train that glides over a track of memory without needing to carry a heavy baggage car. The researchers, Bo Liu and Qiang Liu from the University of Texas at Austin, propose a clever two-step training process that allows the AI to learn how to remember things perfectly, even though it only keeps a small, fixed amount of information during actual use.

Think of the training process as a rehearsal with a director and an actor.

The Director (Prefiller Q): First, a powerful "director" model reads the entire story from start to finish. Because it has access to the whole text, it can write down the perfect "memory notes" for every single sentence. It knows exactly what the actor needs to remember to make the next line make sense.
The Actor (Decoder P): Next, a "sliding window" actor tries to perform the story. This actor is only allowed to look at the last few sentences and the memory notes the director just handed them. The actor tries to predict the next word and, crucially, tries to write their own memory notes for the next step.

Here is the magic trick: The researchers teach the actor to match the notes they write with the perfect notes the director wrote. They use a "consistency loss," which is basically a grading system that says, "If your memory note doesn't look like the director's perfect note, you lose points." Over time, the actor learns to write notes so good that they don't actually need the director anymore.

What happens at the end?
Once the training is done, the director is fired. The actor is left alone on the stage. Now, the actor runs the show by reading the last few sentences and using the memory notes they wrote for the previous sentence to understand the current one. This creates a loop where the AI has a "sliding window" of immediate context but also carries a rich, compressed memory of everything that came before, all while keeping its memory size fixed and small.

The Results: Speed Meets Smarts

The authors tested this Maglev system on a massive dataset of 43.52 billion tokens (a huge amount of text). They compared their model against standard sliding-window models and other advanced memory models.

The results suggest that Maglev is a strong contender. In their experiments:

The Maglev model improved the FineWeb-Edu validation bits per byte (BPB) from 0.7413 (the standard sliding-window baseline) down to 0.7251. In the world of AI, a lower BPB score means the model is making fewer mistakes and understanding the text better.
It also boosted the average accuracy on various downstream tasks (like answering questions or finishing sentences) from 54.1% to 56.4%.

One of the most surprising findings is that the "Director" and the "Actor" can actually share most of their brain power (parameters). The researchers found that even when the two models shared the majority of their internal weights, the Maglev system still kept most of its performance gains. This means the system could be much smaller and cheaper to run than if the two models were completely separate.

Why This Matters

The paper suggests that Maglev offers a practical way to give AI models a "nonlinear" memory—one that can be rewritten and updated with every single word, just like a human thinking—without the massive computational cost of looking at the entire history every time. Unlike some other methods that force the AI to stop and think in big chunks or use rigid mathematical formulas to update memory, Maglev lets the model update its memory continuously and creatively.

While the authors note that this is a preliminary investigation and that scaling it up to even larger models will require more work, the core idea is solid: by using a parallel "teacher" to guide a "student" during training, we can teach an AI to be both fast and deeply remembering. It's a bit like teaching a student to take perfect notes in a classroom so that, when they are alone in the library later, they can recall the whole lecture just by looking at their own shorthand.

Technical Summary: Maglev – Sliding Recurrent Memory

Problem Statement
Current Transformer architectures face a fundamental trade-off between memory capacity and computational efficiency. Full causal self-attention enables powerful nonlinear retrieval and in-context computation but incurs costs that grow linearly with context length. Conversely, sliding-window attention bounds attention and KV-cache costs but discards distant information, failing to provide persistent long-term memory. Classical recurrent models (e.g., LSTMs) compress history but propagate small states sequentially, limiting capacity and sequence parallelism. While Memory Transformers and linear attention models attempt to bridge this gap, they often impose artificial segment boundaries, rely on structured linear updates, or require sequential unrolling during training. The core challenge is to develop a practical, token-wise nonlinear recurrent memory mechanism for Transformers that allows for parallel pretraining while maintaining bounded inference costs.

Methodology: Maglev Architecture
Maglev addresses this by decoupling memory construction during training from memory propagation during inference. The architecture consists of two coupled models:

Prefiller ( $Q$ ): A causal Transformer that processes the observed sequence in parallel using a larger receptive field (e.g., interleaved full and sliding-window attention). It generates a trajectory of target memory vectors, $m'_t$ , for every token position.
Decoder ( $P$ ): A sliding-window Transformer that predicts the next token and generates its own memory vectors, $m_t$ . Crucially, $P$ does not have access to the full history; it only sees a bounded window of recent tokens and a shifted window of memory vectors.

Training Mechanism: Lifted Parallel Training
Instead of unrolling the recurrent state sequentially (which prevents parallelization), Maglev employs a two-pass training strategy:

Pass 1: The prefiller $Q$ generates the complete target memory trajectory $m'_{1:T}$ in parallel.
Pass 2: The decoder $P$ consumes the shifted target trajectory $m'_{0:T-1}$ alongside the input tokens to predict the next token and produce its own memory $m_{1:T}$ .
Objective: The model is optimized using a composite loss function:
- Next-Token Prediction Loss ( $L_{CE}$ ): Standard cross-entropy loss for language modeling.
- Memory Consistency Loss ( $L_{cons}$ ): A term that aligns the decoder's output memory $m_t$ with the prefiller's target $m'_t$ . This teaches the decoder to reproduce the necessary memory state for the next step without sequential unrolling.

Inference Mechanism
At inference, the prefiller $Q$ is discarded. The decoder $P$ runs recurrently, feeding its own previously generated memories $m_{t-1}$ back into the system.

Recurrent K/V Injection: Rather than using dedicated memory tokens, Maglev injects the shifted memory state into the Key/Value (K/V) features of the attention mechanism via learned projections and gating mechanisms.
Bounded Complexity: The decoder maintains a fixed-size sliding window of tokens and memories. The recurrent state is carried forward through the K/V cache, ensuring that inference costs (compute and cache size) remain bounded and independent of the total sequence length, matching the profile of standard sliding-window attention.

Key Contributions

Token-Wise Nonlinear Recurrence: Maglev introduces a mechanism where every token's decoder state serves as memory, traversing the full nonlinear depth of the Transformer without requiring extra memory tokens or artificial segment boundaries.
Parallel Training for Recurrent Models: By using a prefiller to generate targets in parallel, Maglev enables the training of nonlinear recurrent models with standard sequence-parallel throughput, avoiding the sequential bottlenecks of traditional RNN training.
Parameter Sharing: The authors demonstrate that the prefiller $Q$ and decoder $P$ can share the majority of their parameters (specifically Transformer blocks) without significantly degrading performance, reducing the parameter memory overhead of the training setup.
Bounded Inference: The deployed model retains the efficiency of sliding-window attention while gaining persistent memory capabilities through the recurrent K/V injection.

Experimental Results
Evaluated on the nanochat d20 architecture (435M parameters, trained on 43.52B tokens), Maglev was compared against interleaved full/sliding-window Transformers, purely sliding-window models, and Latent Recurrent Transformer (LRT) baselines.

Performance: The best Maglev variant (separate parameters, $\lambda=1$ ) improved the FineWeb-Edu validation bits-per-byte (BPB) from 0.7413 (sliding-window baseline) to 0.7251.
Downstream Tasks: It achieved an average downstream accuracy of 56.4%, outperforming both the matched sliding-window Transformer (54.1%) and the matched LRT baseline.
Consistency Weight: The study noted that the optimal consistency weight ( $\lambda$ ) depends on the parameter sharing strategy; a higher $\lambda$ benefited the separate-parameter model but could constrain the shared-parameter model.

Significance and Claims
The paper claims that Maglev provides a practical pathway to equip nonlinear Transformers with persistent, token-wise memory that can be rewritten after every token. It successfully bridges the gap between the expressive power of full attention and the efficiency of bounded-context models. By separating memory generation from memory usage during training, Maglev achieves parallel pretraining for recurrent architectures, a feat previously difficult for nonlinear recurrence. The authors position Maglev as a preliminary investigation into this design space, noting that while it improves upon sliding-window and latent recurrent baselines in the tested setting, further work is needed to scale the approach, optimize kernel efficiency for deployment, and explore alternative memory injection mechanisms (e.g., residual-stream injection or cross-attention).

The Memory Problem: Why AI Needs a Better Notebook

Maglev: The Train That Remembers Without the Baggage

The Results: Speed Meets Smarts

Why This Matters

Technical Summary: Maglev – Sliding Recurrent Memory

More like this