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

Authors: Benjamin L. Badger, Ethan Roland

arXiv · PDF

Summary

The authors propose Toeplitz MLP Mixers (TMMs), a Transformer alternative that replaces attention with multiplication by Toeplitz matrices (structured matrices where diagonals are constant). This drops computational complexity from quadratic to O(n log n) during training and linear at inference, while using less memory. Despite lacking the sophisticated gating and state tracking of architectures like Mamba, TMMs retain more information from inputs, showing better copying ability and in-context learning. The authors argue this comes from fewer architectural biases, and they find that trained Toeplitz layers tend to be nearly invertible even though the overall model isn't required to be.

Main takeaways:

• Toeplitz matrix multiplication (where each diagonal has the same value) replaces quadratic-cost attention with O(n log n) training and O(n) inference complexity
• TMMs lack explicit gating or state maintenance but achieve better training efficiency (lower loss per compute and memory usage) than comparable architectures
• Higher information retention from inputs leads to improved copying, retrieval, and in-context learning compared to other sub-quadratic alternatives
• Counterintuitively, trained Toeplitz layers are more likely to be invertible (or nearly so) than in architectures designed to be invertible
• The authors argue fewer architectural biases allow the model to retain more input information, benefiting tasks requiring faithful recall

Relevance

Potentially relevant to my work on persona geometry and activation patterns. The finding that TMMs retain more input information and the analysis of layer invertibility could relate to how behavioral installations propagate through networks. My result #237 on persona geometry collapse during fine-tuning might connect to their observations about trained layer properties, though the architectures differ significantly.

Abstract

arXiv:2605.06683v1 Announce Type: new Abstract: Transformer-based large language models are in some respects limited by the quadratic time and space computational complexity of attention. We introduce the Toeplitz MLP Mixer (TMM), a transformer-like architecture that swaps attention for triangular-masked Toeplitz matrix multiplication over the sequence dimension resulting in $\mathcal{O} (dn \log n)$ time and $\mathcal O(dn)$ space complexity during training and $\mathcal O(dn)$ time and space at inference prefill. Despite the lack of sophisticated input modulation or state maintenance present in other sub-quadratic architectures, TMMs yield greater training efficiency in terms of loss achieved per compute and device memory. We demonstrate that TMMs are capable of retaining more input information resulting in improved copying ability, which we argue results from a lack of architectural biases. Consistent with higher input information retention, TMMs exhibit superior information retrieval and in-context learning benchmark accuracy compared to comparable architectures. We conclude with an analysis from the perspective of operator index theory and show that, counterintuitively, trained Toeplitz layers of causal non-invertible models are more likely to be invertible or nearly so than models that are actually invertible over their inputs.