跳转至

线性化二维单纯形注意力机制

文章背景与核心概要

在大语言模型和长序列处理的研究中，传统的Transformer注意力机制面临着计算复杂度随序列长度呈二次方增长的瓶颈，而现有的窗口化二维单纯形注意力（windowed 2-simplicial attention）虽然在一定程度上缓解了局部建模的问题，却无法兼顾全局上下文。

本文介绍了一种名为“线性化二维单纯形注意力（Linearized 2-Simplicial Attention）”的新型架构。该方法通过将三线性得分改写为复合查询（query）与键（key）之间的内积，并利用存储在固定大小状态中的正随机特征（positive random features）对其中一个token轴的求和进行近似，从而在保持全局上下文的同时实现了序列长度的线性计算成本。结合使用自定义Triton内核的Kimi Delta Attention（KDA），该架构彻底消除了传统的Softmax注意力，在同等计算量下实现了更高的下游任务准确率，并在16k上下文长度下显著降低了困惑度。

摘要

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.

我们提出了一种线性化的二维单纯形注意力形式，它通过将三线性得分改写为复合查询与键之间的内积，使得在一个token轴上的求和呈现出与普通Softmax注意力相同的形式。接着，我们使用正随机特征来近似这个求和，并将整个历史信息存储在固定大小的状态中，而第二个轴在近期token的短窗口内保持显式。这使得我们在实现序列长度线性成本的同时，具备了窗口化二维单纯形注意力所缺乏的全局覆盖能力。我们通过自定义的Triton内核实现该机制，并将其与Kimi Delta Attention相结合，构建了一个完全不包含Softmax注意力的模型。在同等计算量下，该模型在所有对比架构中实现了最高的平均下游准确率；在16k上下文长度下，它不仅提升了相比KDA混合模型的平均准确率，还将LAMBADA困惑度从715.6降至602.6。

资源与链接