SMat-Attention：结构化长上下文序列建模

长上下文序列模型面临一个基本的权衡：softmax 注意力以二次成本使用灵活的 token 级交互，而线性注意力通过将历史压缩到固定大小的状态来获得线性时间训练和恒定时间解码。在这项工作中，我们询问是否可以通过可调节的结构概念将这些机制联系起来。为此，我们通过一系列具有结构化远程路由的因果掩码引入了结构化矩阵注意力（SMat-Attention），其行支持具有 VC 维度 $d$。在我们的构造中，$d=1$ 恢复标准因果掩码，并且增加 $d$ 允许更丰富的子集路由模式。我们提供分块前向和后向算法以实现硬件效率。对于长度为 $T$ 的序列，硬路由构造需要 $O(T^{2-3/d}+T)$ 工作，尽管对于我们指定的系列来说，掩码很密集。在固定水平流中，使用 $O(T^{1-1/d})$ 缓存状态，在远程前缀之后进行解码需要每个词元恒定的时间。因此，SMat-Attention 使 VC 维度成为控制访问模式复杂性、预填充成本和解码内存的显式旋钮。根据经验，子集路由和规则辅助多键检索实验说明了掩码的路由表达能力。使用具有 top-$k$ 查询读取的学习路由对 Mamba-2 和门控 DeltaNet 进行扩展，保留次二次预填充，提高多种设置中骨干网的召回准确性，并实现可比的小规模语言建模性能。