🤖 machine learning

Structured Recurrent Mixers for Massively Parallelized Sequence Generation

This paper introduces the Structured Recurrent Mixer (SRM), a novel architecture that enables algebraic conversion between parallel training and recurrent inference to achieve superior training efficiency, information capacity, and inference throughput compared to existing linear-complexity models, while demonstrating significant performance gains in both standard benchmarks and reinforcement learning tasks.

Original authors: Benjamin L. Badger

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin L. Badger

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 Big Problem: The "Assembly Line" vs. The "Library"

Imagine you are trying to write a story, one word at a time.

Old School (Recurrent Models): Think of this like a single writer sitting at a desk. They write one word, remember it, write the next, and so on. They are very efficient at remembering the story as they go (low memory usage), but they can only write one word at a time. If you want 100 stories written, you have to wait for the writer to finish one before starting the next.
Modern Standard (Transformers): Think of this like a massive library where you can look up every word of a story at once. This is great for training because you can read the whole book in parallel. However, when it's time to write the story (inference), the library gets clogged. To write the next word, the library has to re-scan the entire book you've written so far to find the context. As the story gets longer, the library gets slower and slower, and it needs a huge amount of space (memory) to hold all those books.

The Dilemma: We want the speed of the library for training, but the efficiency of the single writer for generating text.

The Solution: The "Structured Recurrent Mixer" (SRM)

The author introduces a new architecture called the Structured Recurrent Mixer (SRM). You can think of the SRM as a chameleon or a transformer (in the superhero sense, not the AI sense) that can change its shape depending on what it's doing.

During Training (The Library Mode): When the model is learning, it acts like the massive library. It looks at the whole sequence of words at once. This makes learning fast and stable.
During Inference (The Writer Mode): When the model is actually generating text, it instantly switches to being the single writer. It doesn't need to re-scan the whole book; it just keeps a tiny, constant-sized "notebook" (cache) of what it just wrote. This makes it incredibly fast and allows it to handle many stories at once.

The magic trick is that the math behind the SRM allows it to switch between these two modes algebraically. It's like having a blueprint that says, "If we are building, we use these bricks; if we are living in it, we use these walls," without having to tear the building down and rebuild it.

Why This Matters: The "Batch" vs. The "Length"

The paper makes a crucial observation about how we should scale AI:

Scaling Length (The Story): The paper argues that trying to make these "single writer" models read infinitely long books is a bad idea. Eventually, the writer's memory (the notebook) gets too full, and they start forgetting details. It's like trying to remember a 1,000-page novel in a 1-page notebook; you will lose information.
Scaling Batch (The Crowd): However, these models are amazing at handling many different stories at the same time. Because they don't need a huge library for each story, you can have 100 writers working in parallel on 100 different stories without them tripping over each other.

The Analogy: Imagine a coffee shop.

Transformers are like a shop with one giant espresso machine that takes 10 minutes to brew a single latte, but it can brew 100 cups at once if you have a huge counter. As the order gets more complex, the machine slows down.
SRMs are like a shop with 100 simple, fast manual brewers. Each brewer makes one cup quickly and remembers the recipe in their head. You can't brew a 100-cup order on a single brewer (too much memory), but you can serve 1,000 customers simultaneously by using 1,000 different brewers.

The Results: Speed and Volume

The paper tested this new architecture and found some impressive numbers:

Speed: On standard hardware, the SRM was 12 times faster at generating text than a standard Transformer.
Concurrency: It could handle 170 times more simultaneous requests (users) than a Transformer.
Accuracy: Even though it was so fast and handled so many requests, it didn't lose its smarts. In math tests (GSM8k), it actually solved about 30% more problems than a Transformer when given the same amount of computer power.

The "Reinforcement Learning" Twist

The paper also looked at how this helps with Reinforcement Learning (teaching AI by trial and error).

Imagine you are teaching a dog to fetch.

Standard Approach: You send the dog out once. If it fails, you try again.
SRM Approach: Because the SRM is so fast, you can send 50 dogs out at the same time. You check which ones got the ball, and you only reward the winners.

The paper found that by generating many samples at once and using a special "resampling" trick (making sure you have enough "good" examples to learn from), the SRM learned much better than models that only tried a few times.

Summary

The Structured Recurrent Mixer is a new AI design that gets the best of both worlds:

It learns efficiently like a modern parallel computer.
It generates text efficiently like a simple, memory-light writer.
It is built to handle massive crowds of users simultaneously, which is becoming more important as AI moves from "one big answer" to "many quick answers."

The paper concludes that while we often try to make AI smarter by making it read longer books, we should instead focus on making it faster at handling many different tasks at once, and the SRM is the perfect tool for that job.

Technical Summary: Structured Recurrent Mixers for Massively Parallelized Sequence Generation

Problem Statement

The field of language modeling has shifted from recurrent architectures (RNNs, LSTMs) to non-recurrent, sequence-parallel models (Transformers) to achieve training efficiency and stability. However, this shift incurs a significant cost during inference: Transformers require $O(n)$ memory and sequential token generation, leading to lower inference throughput and concurrency compared to recurrent models. Conversely, while recurrent models offer $O(1)$ memory per sample and efficient sequential generation, they historically struggle with training efficiency and stability, and often fail to scale effectively to long context lengths due to information compression limits.

Hybrid approaches (e.g., Mamba, Gated Delta Nets) attempt to bridge this gap but often rely on specialized device kernels and still face cache scaling issues. Furthermore, for tasks where outputs can be rapidly verified (e.g., code generation, mathematical proofs), the paradigm of "accuracy per compute" suggests that generating many samples from smaller, faster models may be more efficient than generating fewer samples from larger, more accurate ones. The paper posits that recurrent models are ill-suited for scaling sequence length for information-rich inputs but are ideally suited for scaling in the batch (sample) dimension due to their constant memory footprint per sample.

Methodology: Structured Recurrent Mixers (SRMs)

The authors introduce the Structured Recurrent Mixer (SRM), an architecture designed to algebraically convert between a sequence-parallel representation (for training) and a recurrent representation (for inference) without requiring specialized kernels or device-specific memory management.

Core Architecture

SRMs are based on the Masked Mixer, a transformer-like architecture where attention is replaced by masked Multi-Layer Perceptrons (MLPs). The key innovation lies in the parameterization of the token mixing matrices:

Qualitative Feature Mapping: Analysis of trained Masked Mixers revealed that efficient token mixing matrices exhibit specific qualitative features: columns with near-identical values, rows with near-identical values, and weight decay increasing with distance from the main diagonal.
Structural Constraints: To enable algebraic conversion to a recurrent form, SRMs restrict token mixing matrices to row-repeat or column-repeat structures.
- Row-Repeat: Rows contain identical elements.
- Column-Repeat: Columns contain identical elements.
- Decay: A trainable decay constant $\lambda$ is applied to the diagonal or off-diagonal elements to manage information retention.
Algebraic Conversion: These constraints allow the matrix multiplication $Y = XM + B$ to be rewritten as a fixed-memory recurrent operation (Equation 2 and 3 in the paper). The recurrent hidden state (cache) becomes a simple sum term, requiring only constant memory ( $O(d)$ ) regardless of sequence length.
Hybrid Heads: The architecture employs "mixed" heads, where half the heads use row-repeat matrices and half use column-repeat matrices, combined via linear projections. This configuration was found to offer the best balance of training efficiency and information retention.

Training and Inference

Training: SRMs are trained using sequence parallelism (shifting tokens) similar to Transformers. The quadratic complexity at training time is effectively hidden by memory transfer overheads, making them comparable in throughput to Transformers for standard context lengths.
Inference: The model switches to a recurrent representation where tokens are generated one by one, updating a constant-sized cache. This eliminates the need for large KV caches, drastically reducing memory bandwidth requirements.

Key Contributions

Qualitative Feature Mapping: A method to map the token mixing operations of quadratic-complexity models to linear-complexity recurrent models by enforcing row/column repeat structures and decay.
Dual-Architecture Design: An architecture that supports sequence parallelism during training and batch parallelism during inference, allowing for algebraic conversion between the two without changing model parameters or cache structures.
Optimized Inference Engine: Implementation of SRMs using the Mojo/MAX framework, demonstrating significant performance gains over standard PyTorch implementations.
Balanced Resampling for RL: A resampling method for Group Relative Policy Optimization (GRPO) that maintains exploration in large-batch reinforcement learning by ensuring a mix of "good" and "bad" samples in the training batch.

Experimental Results

Training Efficiency

SRMs are slightly less efficient to train than $O(n^2)$ Transformers and Masked Mixers on a per-parameter basis but are more efficient than Mamba models across diverse datasets (FineWeb-edu, FineMath).
Unlike Mamba models, which suffer significant throughput drops on older hardware (e.g., V100s), SRMs perform consistently across different GPU architectures.

Benchmark Performance

Accuracy: On compute-equivalent benchmarks (IFEval, SQuAD, GSM8k, etc.), Transformers generally top the benchmarks. SRMs outperform Mamba models on most tasks but generally lag behind Transformers in single-sample accuracy.
Information Retention: SRMs exhibit higher information retention and capacity than Mamba models, comparable to quadratic complexity models. However, like all recurrent models, they face information loss when context lengths exceed their compression limits.

Inference Throughput and Concurrency

Throughput: SRMs demonstrate massive throughput improvements over Transformers due to reduced memory transfer.
- On PyTorch (V100): SRMs achieve 9.66x to 43.29x higher throughput than Transformers depending on context length.
- On Mojo/MAX (H100): SRMs achieve 12x the throughput of Transformers inferenced via vLLM.
Concurrency: SRMs support significantly larger batch sizes.
- On Mojo/MAX (H100): SRMs achieve 170x the concurrency of Transformers.
- Compared to Mamba: SRMs show >500x higher concurrency and ~7x higher throughput than equivalently powerful Mamba models.

Reinforcement Learning and Verifiable Outputs

Pass@k Scaling: For GSM8k, SRMs follow a similar Pass@k power law to Transformers.
Compute Efficiency: Due to the 10-30x throughput increase, SRMs can generate many more samples for the same compute budget. The paper reports a 30% increase in GSM8k Pass@k for SRMs compared to Transformers when compute is held constant.
RL Training: Using balanced resampling, SRMs trained with GRPO on large batches (50 samples) achieve higher single-sample accuracy and better exploration than Transformers trained on small batches (5 samples).

Significance and Claims

The paper concludes that while recurrent models are poorly suited for scaling sequence length for information-dense inputs, they are exceptionally well-suited for scaling in the batch dimension. The SRM architecture successfully bridges the gap between the training efficiency of quadratic models and the inference efficiency of recurrent models.

The authors claim that SRMs offer a viable path for "accuracy per compute" paradigms, particularly for tasks with rapid verification. By leveraging the constant memory footprint of recurrent models, SRMs can exploit the "memory wall" limitations of modern hardware more effectively than Transformers, achieving order-of-magnitude improvements in throughput and concurrency without the need for specialized kernels. The work suggests that for future multi-threaded hardware trends, architectures like SRMs that maximize arithmetic operations relative to memory accesses will become increasingly dominant.