🤖 AI Summary
This work addresses the high computational complexity and limited global context modeling of conventional 2-simplex attention by presenting its first linearization. The authors reformulate the trilinear scoring function as an inner product between composite queries and keys, applying random feature approximation along one token axis while maintaining a fixed-size state, and preserving a local window along the other axis. This design achieves linear complexity with respect to sequence length while retaining a global receptive field. Integrating positive random features, fixed-state memory, custom Triton kernels, and Kimi Delta Attention fusion, the method establishes the first softmax-free attention architecture. Experiments demonstrate that, under identical computational budgets, the model achieves state-of-the-art average accuracy on downstream tasks; notably, it reduces LAMBADA perplexity from 715.6 to 602.6 at 16k context length, significantly outperforming KDA hybrid models.
📝 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.