NeFut Logo NeFut
EN 管理员登录

[AI学术] SMat-Attention:结构化长上下文序列建模

发布于:2026-09-30 22:00 最后更新:2026-10-06 12:11
#algorithm #AI #Machine Learning

长上下文序列模型面临一个基本权衡:softmax attention 能在 token 级别实现灵活交互,但代价是 $O(T^2)$ 的计算;线性 attention 通过将历史压缩到固定大小的状态,实现 $O(T)$ 的训练和常数时间的解码。\ \ 本文提出 结构化矩阵注意力 (SMat-Attention),通过一族因果掩码实现可调结构。掩码的每行支持集合的 VC 维度记为 $d$,其中 $d=1$ 恢复标准因果掩码,增大 $d$ 可产生更丰富的子集路由模式。\ \ 为提升硬件效率,作者设计了块状(chunkwise)前向和后向算法。对于长度为 $T$ 的序列,硬路由构造的工作量为 $$O\bigl(T^{2-3/d}+T\bigr)$$,即使掩码是密集的,也保持次二次复杂度。\ \ 在固定窗口流式(fixed‑horizon streaming)设置下,解码远程前缀后每个 token 的时间为常数,所需缓存状态规模为 $$O\bigl(T^{1-1/d}\bigr)$$。因此,VC 维度 $d$ 成为显式的调节旋钮,控制访问模式复杂度、预填成本和解码内存。\ \ 实验包括子集路由和规则辅助的多键检索,展示了掩码路由的表达能力。将 SMat-Attention 扩展到 Mamba‑2 与 Gated DeltaNet,采用学习路由并结合 top‑$k$ 查询读取,保持次二次预填,显著提升召回准确率,并在若干设置下实现与原模型相当的小规模语言建模性能。\ \ 点评:SMat-Attention 通过引入可调的 VC 维度,将软注意力的灵活性与线性注意力的高效性桥接,为长上下文建模提供了新的设计空间,实验验证了其在效率和精度上的双重优势。

原文链接: https://arxiv.org/abs/2609.36062

[h] 返回首页