🤖 AI

Linearized 2-Simplicial Attention

This paper introduces Linearized 2-Simplicial Attention, a novel architecture that combines random feature approximation for global context with explicit short-window processing to achieve linear computational cost and superior downstream performance without using any softmax attention.

Original authors: Aritra Das, Dhruman Gupta, Debayan Gupta

Published 2026-08-11

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Aritra Das, Dhruman Gupta, Debayan Gupta

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 in AI Brains

Imagine you are trying to teach a robot to write a story. To make the story make sense, the robot needs to remember what happened at the very beginning of the sentence while it is writing the very last word. In the world of artificial intelligence, this is called "attention." The most popular way to do this is like a super-organized librarian who, every time the robot asks for a word, instantly scans the entire library of everything written so far to find the perfect match. This works incredibly well, but it has a massive flaw: the more words the robot writes, the longer the librarian takes to scan the shelves. If the story gets too long, the librarian gets overwhelmed, and the robot runs out of memory or time.

Scientists have been trying to build a "fast librarian" who can remember the whole library without scanning every single book. Some have tried to compress the library into a tiny summary note, but this often means the robot forgets specific details. Others have tried to only look at the last few pages, but then the robot can't remember the plot twist from chapter one. The big question in this corner of computer science is: Can we build a brain that remembers everything from the start, understands complex relationships between three different ideas at once, and does it all without getting slower as the story gets longer? This paper dives into that exact puzzle, proposing a new way to organize the robot's memory that is both fast and surprisingly deep.

The Paper's Big Idea: A New Kind of Memory

The researchers, Aritra Das, Dhruman Gupta, and Debayan Gupta from Truth Audit Labs, have invented a new type of attention mechanism they call Linearized 2-Simplicial Attention (or "LinSimp" for short). To understand why this is special, we first need to look at how standard AI brains usually work.

Most AI models use "attention" to connect a current word (the query) to a past word (the key). It's a one-on-one relationship. But sometimes, to really understand a sentence, you need to connect three things at once. Imagine the sentence: "The cat sat on the mat because it was soft." To understand why the cat sat there, the AI needs to link "cat," "mat," and "soft" together simultaneously. This is called a "2-simplicial" interaction. Previous attempts to do this were like trying to find a specific triple of friends in a crowd of a million people by checking every single possible trio. It was accurate but incredibly slow and expensive, getting slower and slower as the crowd grew.

The authors' breakthrough is a clever mathematical trick. They realized they could rewrite this complex three-way connection so that one part of the search happens globally (looking at the whole history) while the other part stays local (looking only at recent words).

Here is the analogy: Imagine the AI's memory is a giant, infinite whiteboard.

The Old Way: To find a connection, the AI had to draw a line from the current word to every pair of past words on the whiteboard. If the whiteboard had 1,000 words, that's a million lines to draw.
The New Way (LinSimp): The authors suggest a different strategy. They take the "current word" and a "recent anchor word" (like the last few words spoken) and mix them together to create a special "composite key." They then use a magical "random feature" filter to project this key onto a summary of the entire past whiteboard. This allows the AI to instantly "feel" the connection to the whole history without drawing every single line.

By using this method, the cost of processing the story grows linearly. If the story doubles in length, the work only doubles, rather than quadrupling like the old methods. They store the entire past in a fixed-size "state" (like a compact summary note) and only keep the most recent 64 words in a detailed "anchor window" to fine-tune the search.

What They Found and How Sure They Are

The team built a model using this new LinSimp layer and combined it with another advanced technique called "Kimi Delta Attention" (KDA). They tested this "no-softmax" model (meaning it doesn't use the traditional, slow attention method at all) against standard models and other experimental ones.

Their results suggest that this new approach is highly effective. In tests involving a 16,000-word context (a very long story), their model improved the average accuracy on various reasoning tasks by 0.0079 compared to a hybrid model that still used some slow attention. Even more impressively, it lowered the "perplexity" (a measure of how confused the model is) on the LAMBADA test from 715.6 down to 602.6. This means the model was significantly better at predicting the next word in long, complex sentences.

However, the authors are careful about their claims. They note that their custom computer code (called "kernels") is still a "functional first implementation." While it runs fast, it isn't quite as speedy as the mature, standard attention code yet. They also mention that their experiments used a single "seed" (a starting point for randomness), so they cannot yet provide statistical confidence intervals. They suggest that while the results are promising and the model achieves the highest mean accuracy among the compared architectures in their specific tests, more work is needed to fully understand how it scales to even larger sizes.

The Verdict

This paper doesn't claim to have solved the memory problem forever, but it offers a very strong new tool. It suggests that by mixing a global summary of the past with a focused look at the recent present, we can build AI models that are both fast and capable of deep, three-way reasoning. The authors show that you don't have to choose between remembering the whole story and processing it quickly; with the right mathematical trick, you can have both. While the speed of their custom code still has room for improvement, the accuracy gains suggest this is a promising direction for the future of long-context AI.

Technical Summary: Linearized 2-Simplicial Attention

Problem Statement
Standard softmax attention provides content-based access to previous tokens but suffers from quadratic computational complexity ( $O(n^2)$ ) and a growing KV cache, limiting scalability. While linear attention mechanisms compress the past into a fixed-size state, they often struggle with exact retrieval and selective updates. Conversely, 2-simplicial attention introduces higher-order interactions by scoring a query against pairs of tokens (triplets), which theoretically improves performance on triple-dependent tasks. However, exact 2-simplicial attention is prohibitively expensive, requiring $O(n^3)$ computation. Recent attempts to make this practical rely on restricting both token axes to fixed windows, sacrificing global reach for efficiency.

Methodology
The authors propose LinSimp, a linearized form of 2-simplicial attention that achieves linear cost ( $O(n)$ ) while maintaining global reach on one axis and local precision on the other.

Linearization via Trilinear Rewriting:
The core innovation is rewriting the trilinear score $\langle q_i, k_j, r_c \rangle$ as an inner product between a composite query and a key:
$\langle q_i, k_j, r_c \rangle = (q_i \odot r_c)^\top k_j$
This reformulation allows the summation over the token index $j$ to be treated as an ordinary exponential dot-product kernel, similar to standard linear attention.
Positive Random Features:
To approximate the exponential kernel $\exp((q_i \odot r_c)^\top k_j)$ , the authors employ positive random features (specifically orthogonal random features). This allows the summation over the global token history ( $j$ ) to be compressed into a fixed-size state ( $M_i$ and $a_i$ ), independent of sequence length.
- Global Axis ( $j$ ): Handled via the random feature state, enabling access to the entire causal prefix.
- Local Axis ( $c$ ): The "anchor" dimension is kept explicit but restricted to a short rolling window (e.g., 64 tokens). This creates a hybrid mechanism where each recent anchor generates a unique composite query sent to the global state.
Architecture and Implementation:
- Model Composition: The authors integrate LinSimp layers into a model architecture that also utilizes Kimi Delta Attention (KDA), a gated delta-rule fast-weight model. They construct a "no-softmax" model containing 18 KDA layers and 6 LinSimp layers.
- Custom Kernels: The method is implemented using custom Triton kernels. The forward pass processes queries in chunks, reconstructing prefix states from boundary conditions to avoid materializing full token pairs. The backward pass uses a two-sweep design to accumulate gradients without saving large intermediate tensors ( $T \times m \times D_v$ ).

Key Contributions
The paper outlines four specific contributions:

Derivation: A one-mode kernelization of 2-simplicial attention, transforming the trilinear score into a form amenable to linear attention approximations.
Implementation: A linear-time implementation utilizing a positive orthogonal random-feature bank for the global mode and a short explicit anchor window for the local mode.
Kernel Engineering: Custom Triton forward and backward kernels, including numerical validation against an FP32 reference.
Empirical Comparison: A comprehensive evaluation against standard attention, exact windowed 2-simplicial baselines, and KDA hybrids under both iso-token and iso-FLOP (equal compute) settings.

Results
The authors evaluated their models on web data (FineWeb-Edu) and mathematical data (FineMath-4+) across 2k and 16k context lengths.

Iso-FLOP Performance: Under matched compute budgets, the "no-softmax" KDA + LinSimp model achieved the highest mean downstream accuracy among compared architectures in the 3B-token math run and the 16k context setting.
Long Context (16k): At 16k context, the LinSimp-based model improved mean accuracy by 0.0079 over the KDA hybrid (which uses standard softmax attention) and significantly reduced LAMBADA perplexity from 715.6 to 602.6.
Efficiency: The custom kernels achieved throughput close to softmax attention (1.03x the time of a softmax layer for a single LinSimp layer at 2k context) but did not yet surpass mature softmax kernels in raw speed.
Baseline Comparison: The exact windowed 2-simplicial baseline (Roy et al., 2025) showed higher validation loss and perplexity compared to the linearized approach when compute was matched.

Significance and Claims
The paper claims that LinSimp successfully bridges the gap between the theoretical benefits of higher-order (trilinear) interactions and the practical constraints of linear-time inference. By decomposing the 2-simplicial interaction into a global random-feature state and a local explicit window, the authors demonstrate that:

It is possible to build a high-performing language model with no softmax attention layers at all.
This architecture can outperform KDA hybrids (which rely on standard attention for some layers) in long-context scenarios under equal compute budgets.
The method offers a "global reach" that windowed 2-simplicial attention lacks, while maintaining linear complexity.

Limitations and Modesty
The authors explicitly note that their experiments are preliminary, using a single random seed, which prevents the estimation of variance or confidence intervals. They acknowledge that the current backward pass kernel is not highly optimized and that the iso-FLOP budget is an analytic accounting measure rather than a direct energy or latency measurement. Furthermore, the scaling behavior of the model beyond the tested parameter counts (approx. 330M) remains unknown.

The Memory Problem in AI Brains

The Paper's Big Idea: A New Kind of Memory

What They Found and How Sure They Are

The Verdict

Technical Summary: Linearized 2-Simplicial Attention

More like this