🤖 machine learning

Toeplitz MLP Mixers are Low Complexity, Information-Rich Sequence Models

The paper introduces the Toeplitz MLP Mixer (TMM), a transformer-like architecture that replaces attention with triangular-masked Toeplitz matrix multiplication to achieve sub-quadratic complexity while demonstrating superior training efficiency, information retention, and in-context learning performance compared to other sub-quadratic models.

Original authors: Benjamin L. Badger, Ethan Roland

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin L. Badger, Ethan Roland

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 "Library of Babel" Bottleneck

Imagine you are trying to write a story, but you have to read your entire story from the very first word every time you want to add a new sentence. If your story is short, this is easy. But if your story is a whole novel, reading it all over again for every single new word takes forever.

This is exactly the problem with current top-tier AI models (called Transformers). They use a mechanism called "attention" that lets them look at every previous word to decide the next one. While powerful, this gets incredibly slow and memory-hungry as the text gets longer. It's like trying to find a specific book in a library where you have to check every single book on every single shelf before you can pick one up.

The New Solution: The "Toeplitz Mixer"

The authors introduce a new model called the Toeplitz MLP Mixer (TMM). Think of this as a new way to organize that library.

Instead of checking every single book individually, the TMM uses a special, repeating pattern (mathematically called a Toeplitz matrix) to organize the information.

The Analogy: Imagine a conveyor belt in a factory. In a standard model, a worker has to walk down the entire line to grab a specific part. In the TMM, the parts are arranged in a repeating, predictable pattern. The worker can use a shortcut (a mathematical trick called the Fast Fourier Transform) to grab exactly what they need instantly, regardless of how long the line is.
The Result: This makes the model much faster and uses less memory, especially when it's reading a long prompt for the first time (like pre-filling a document).

The Surprise: Speed Doesn't Mean "Dumb"

Usually, when you make a computer model faster by simplifying how it looks at data, it gets "dumber." It starts forgetting things.

The "State Space" Models: Other fast models (like Mamba) try to be efficient by keeping a single "summary" of everything they've read so far. It's like trying to remember a 500-page book by only remembering the last sentence. This is fast, but the model often forgets details.
The TMM Advantage: The TMM doesn't rely on a single summary. It keeps the "pattern" of the whole text accessible.
- The Copy Test: The authors tested this by asking models to copy a long list of numbers or words.
- The Result: While the "summary" models (Mamba) struggled to remember the list, the TMM could copy it almost perfectly, just like the slow, heavy Transformers. It retained the information much better than other fast models.

Why Does This Work? (The "No Biases" Theory)

The paper suggests that other fast models have built-in "biases" or habits. For example, they might be programmed to pay extra attention to the most recent words and ignore the older ones.

The Analogy: Imagine a student who only studies the last chapter of a textbook because they think it's the most important. They might pass a quiz on the last chapter but fail if asked about the beginning.
The TMM: The TMM doesn't have this bias. It treats the pattern of the text evenly. Because it doesn't force itself to forget old information, it naturally keeps more details, leading to better performance on tasks like finding specific facts in a long document or following complex instructions.

The "Magic" of Invertibility

The authors did some deep math analysis (using something called "operator index theory") and found a counter-intuitive fact:

Even though the TMM is designed to process information in a way that should lose some details (because it's a "causal" model that only looks forward), the math shows that its internal layers are actually very close to being reversible.
The Analogy: Imagine a machine that shreds paper. Usually, you can't get the paper back. But the TMM is like a shredder that cuts the paper into strips that are still perfectly aligned. If you have the right tool, you can tape them back together almost perfectly. This suggests the model is holding onto the "shape" of the original information very well, even while it's processing it.

The Catch (Limitations)

The paper is honest about what the TMM can't do yet:

Scale: They tested models that are relatively small (20 to 300 million parameters). We don't know yet if this works for the massive models used by big tech companies today.
Generating One Word at a Time: While the TMM is super fast at reading a long prompt (the "prefill" phase), once it starts generating text one word at a time, it slows down to the same speed as the old Transformers. It's fast at the start, but not necessarily fast at the finish line.

Summary

The Toeplitz MLP Mixer is a new type of AI architecture that uses a clever mathematical shortcut to read long texts quickly without forgetting the details. It proves that you don't need to sacrifice "memory" to get "speed." It acts like a librarian who uses a smart filing system to find books instantly, rather than a librarian who tries to memorize the whole library in their head.

Technical Summary: Toeplitz MLP Mixers as Low Complexity, Information-Rich Sequence Models

1. Problem Statement

Transformer-based Large Language Models (LLMs) are fundamentally constrained by the quadratic time ( $O(n^2)$ ) and space complexity of self-attention mechanisms. While various attempts have been made to reduce this complexity to sub-quadratic or linear levels (e.g., linearizing attention, state-space models like Mamba), these approaches often introduce significant trade-offs. Specifically, linear-complexity models frequently suffer from reduced training efficiency, numerical instability, and a marked inability to retain input information. This information retention deficit manifests as poor performance in tasks requiring precise copying, in-context learning, and long-context retrieval, which are hallmarks of modern LLM capabilities. The core challenge is to reduce computational complexity without imposing severe restrictions on token information access that degrade these emergent capabilities.

2. Methodology: The Toeplitz MLP Mixer (TMM)

The authors introduce the Toeplitz MLP Mixer (TMM), a transformer-like architecture that replaces self-attention with triangular-masked Toeplitz matrix multiplication over the sequence dimension.

Core Architecture

Toeplitz Parameterization: Unlike standard MLP Mixers where token mixing weights are arbitrary matrices, TMMs parameterize these weights as Toeplitz matrices. In a Toeplitz matrix, all values on a given diagonal are identical, allowing the entire matrix to be represented by a single vector (the first row and column).
Causal Masking: To enforce causality for next-token prediction, the materialized weight matrix is triangular-masked. In the implementation, activations are left-multiplied by this masked Toeplitz matrix ( $T'$ ), which is upper triangular.
Computation via FFT: By embedding the Toeplitz matrix into a circulant matrix, the authors leverage the Fast Fourier Transform (FFT) to compute the matrix-vector product. This reduces the computational complexity from $O(n^2)$ to $O(dn \log n)$ during training and inference prefill, with $O(dn)$ space complexity.
Multi-Head and Multi-Kernel Variants: The architecture supports multi-headed configurations where independent Toeplitz matrices are applied to projected head dimensions. Additionally, "kernel" variants allow mixing along a limited number of hidden dimension elements, analogous to 1D convolutions.

Theoretical Underpinnings

The paper analyzes TMMs through the lens of operator index theory. Specifically, it examines the Fredholm index of the trained Toeplitz layers. The authors note that while TMMs are not strictly invertible over all possible inputs, trained causal TMM layers often exhibit indices (specifically -1) that resemble shift operators, suggesting a mechanism for sequential information shifting that differs from the biases found in other architectures.

3. Key Contributions and Claims

The paper makes three primary claims regarding the efficacy of TMMs compared to other sub-quadratic architectures (such as Hyena, Mamba, and windowed attention Transformers):

Training Efficiency: Under constant compute conditions, TMMs achieve similar or lower evaluation losses compared to other sub-quadratic models. They demonstrate superior convergence rates relative to linear-complexity baselines like Hyena and Mamba.
Information Preservation: TMMs retain significantly more input information than comparable sub-quadratic models. This is evidenced by:
- Information Extraction: Higher fidelity in regenerating input sequences from the final hidden state.
- Copy Ability: TMMs approach or exceed the copy-task training efficiency of quadratic-complexity Transformers, a task where Mamba and Hyena models typically struggle.
- Information Capacity: When explicitly trained to maximize information retention, TMMs exhibit a higher informational capacity than Hyena or Mamba models, even when requiring significantly less compute (1/2 to 1/8th) to train.
Benchmark Performance: Due to superior information retention, TMMs outperform other low-complexity models on benchmarks requiring precise input representation, such as Information Retrieval (IR) and In-Context Learning (ICL), while maintaining parity on general language understanding tasks.

4. Experimental Results

The authors trained models on a 10-billion-token subset of FineWeb-edu and evaluated them against Transformers, Mamba 2, and Hyena models.

Training Efficiency: TMMs converged faster than Hyena and Mamba models when matched for compute and memory. The authors note that linear attention implementations suffered from numerical instability and slow convergence, leading them to focus on TMMs.
Copy Task: In a task requiring models to copy 512 tokens from input to output, TMMs achieved training efficiency comparable to Transformers. Notably, a TMM with 1/32nd the model dimension ( $d_m$ ) of a Hyena model achieved nearly identical copy efficiency. The study found that trainable Toeplitz parameters were critical; freezing these matrices resulted in a total failure to copy, whereas trainable head projections alone were insufficient.
Benchmarks:
- IR/ICL: TMMs outperformed Hyena on SQuAD, SQuAD v2, LongBench, and IFEval.
- General Tasks: TMMs performed on par with Hyena and Transformers on tasks like ARC-Easy, GLUE, and HellaSwag.
Ablation Studies: The authors found that using two sequential Toeplitz layers separated by nonlinearity (standard MLP mixer style) was less efficient than a single layer. Furthermore, increasing the number of heads improved asymptotic convergence, while increasing the kernel size beyond a certain point reduced efficiency.

5. Significance and Limitations

Significance

The paper argues that the TMM offers a "simple, attention-free path" to competitive training efficiency and fast decoding. Its primary significance lies in demonstrating that low-complexity models do not inherently require severe architectural biases (like recency bias or state-space selectivity) that degrade information retention. By using explicit Toeplitz parameterization, the model avoids the "copying gap" seen in Mamba and Hyena, suggesting that structured access patterns (Toeplitz) can preserve the information-rich capabilities of Transformers without the quadratic cost during the prefill phase.

The authors also highlight a counterintuitive theoretical finding: trained causal TMM layers, which are non-invertible over their full input space, often possess Fredholm indices closer to zero (or specifically -1, acting as shift operators) than models explicitly trained for invertibility (like autoencoders). This suggests that the "invertibility" of a layer in a causal model is less about mathematical reversibility and more about the preservation of information flow through shifting.

Limitations

The authors explicitly acknowledge several limitations:

Scale: Experiments were conducted on models ranging from 20M to 300M parameters trained on ~13-26B tokens. There is no guarantee these findings scale to models trained on orders of magnitude more compute.
Autoregressive Inference: While TMMs are $O(dn \log n)$ during prefill, they revert to $O(n^2)$ complexity for autoregressive generation once the first token is computed (as rows of the input matrix must be cached and processed sequentially). Thus, TMMs offer faster time-to-first-token for long prompts but do not accelerate the generation of subsequent tokens compared to Transformers.
Configuration Space: The study tested a limited number of architectural configurations due to the exponential number of possible hyperparameter combinations.

In conclusion, the TMM represents a promising direction for sequence modeling that balances computational efficiency with the high information retention necessary for complex reasoning and retrieval tasks, provided the inference bottleneck of autoregressive generation is managed.