How a Little Known Math Paper Inspired New Attention Mechanisms
Attention is the defining feature of transformers, specifying how different words (or “tokens,” more generally) modify one another to influence meaning. But is standard attention really all you need?
Transformers cannot solve simple linear regression problems unless their attention heads are nearly as large as the input data — a surprising limitation that reveals a fundamental trade-off between efficiency and expressiveness in modern AI systems. New research by CDS Professor of Computer Science and Data Science Andrew Gordon Wilson, CDS PhD students Yilun Kuang, Sanae Lotfi, and Andres Potapczynski, and Courant PhD students Shikai Qiu and Noah Amsel, shows how structured matrices can break through this bottleneck while actually reducing computational costs.
The work, titled “Customizing the Inductive Biases of Softmax Attention using Structured Matrices,” accepted to ICML 2025, addresses two core problems with standard attention mechanisms that have gone largely unnoticed. First, when attention transforms inputs into queries and keys, it compresses information through what researchers call a “low-rank bottleneck” — losing critical details that prevent the model from learning even basic mathematical relationships. Second, attention treats all token pairs equally, missing opportunities to allocate more computational power to pairs of nearby elements that are typically more relevant to each other.
“It all started with the CoLA paper,” Wilson said, referring to a 2023 research program exploring how modeling assumptions translate into exploitable matrix structures. This line of inquiry led to a systematic approach for replacing dense neural network layers with structured alternatives that maintain expressiveness while dramatically reducing computational overhead. This program led to a sequence of papers, including “Compute Better Spent,” and “Searching for Efficient Linear Layers.”
In this most recent paper, on structured attention, a key development came from a recent applied mathematics paper describing Multi-Level Low Rank (MLR) matrices — a structure that had never been applied to machine learning. “I was just browsing applied math papers and casually ran into it,” Kuang said. The novel structured matrix was originally designed to accelerate some numerical linear algebra applications, but the mathematical properties perfectly matched what the team needed to help address key limitations for attention.
The researchers developed two main approaches using structured matrices. For the low-rank bottleneck, they replaced standard attention’s query-key multiplication with Block Tensor Train (BTT) and MLR matrices that achieve full rank while requiring only a fraction of the computational resources. For the distance bias problem, they created MLR attention, which hierarchically allocates more computational power to local interactions than global ones.
The results were striking. On in-context regression tasks where transformers are prone to fail due to small attention head dimensions, the structured attention variants succeeded with smaller model sizes and less compute. For language modeling on OpenWebText, MLR attention achieved better scaling laws across multiple model widths compared to both standard attention and sliding window variants. The method also showed improvements in time series forecasting as sequence lengths increased.
Perhaps most importantly, the work reveals attention mechanisms as part of a much larger design space. The researchers showed that both BTT and MLR matrices belong to a broader family they call Multi-Level Block Tensor Contraction (MLBTC), which unifies many existing structured matrix approaches including Monarch, Butterfly, and Kronecker matrices.
“Machine learning is basically an application of linear algebra,” Wilson said. “The reason GPUs have been so important to deep learning is because they allow you to really accelerate matrix multiplication.” By recognizing that modeling assumptions naturally create structure that can be exploited algorithmically, this research program opens new possibilities for building more efficient architectures.
The work extends a research philosophy that Wilson’s group has pursued for over a decade, dating back to observations about Gaussian processes where recognizing inherent structure led to dramatic computational savings without sacrificing quality. “Instead of asking what approximations we need to make for computational tractability, we realized we were ignoring structure that was already arising from our modeling assumptions,” he explained.
The implications extend far beyond attention mechanisms. Kuang believes the approach could prove especially valuable for scientific computing and partial differential equation data, where inputs are genuinely high-dimensional rather than having the low intrinsic dimensionality typical of language and vision tasks. The team also envisions applications to code generation and point cloud processing, where nested hierarchical structures naturally align with the multi-level approach.
The research represents a fundamental shift in how we think about the efficiency/expressiveness trade-off in neural networks. Rather than imposing approximations that we don’t believe to achieve for computational efficiency, leverage the salient mathematical structures hiding in plain sight.
By Stephen Thomas
