Linearized 2-Simplicial Attention

2026-08-10Artificial Intelligence

Artificial Intelligence
AI summary

The authors developed a new way to speed up a kind of attention mechanism used in language models by rewriting it to use simpler math. They use random features to efficiently keep track of past information in a fixed-size memory while still focusing on recent tokens explicitly. This approach keeps computation costs low even for very long text inputs. When combined with another technique called Kimi Delta Attention, their model performs better on a range of tasks and improves prediction accuracy on a language benchmark.

2-simplicial attentiontrilinear scoresoftmax attentionrandom feature approximationsequence lengthcontext windowTriton kernelsKimi Delta AttentionLAMBADA benchmarkperplexity
Authors
Aritra Das, Dhruman Gupta, Debayan Gupta
Abstract
We present a linearized form of 2-simplicial attention by rewriting the trilinear score as an inner product between a composite query and a key, so that the sum over one token axis takes the same form as ordinary softmax attention. We then approximate this sum with positive random features and store the entire past in a fixed-size state, while the second axis stays explicit over a short window of recent tokens. This enables us to achieve linear cost in sequence length combined with a global reach that windowed 2-simplicial attention lacks. We implement it with custom Triton kernels and combine it with Kimi Delta Attention to build a model with no softmax attention at all. Under matched compute, this model achieves the highest mean downstream accuracy among the compared architectures, and at 16k context it improves mean accuracy over a KDA hybrid while lowering LAMBADA perplexity from 715.6 to 602.6.