论文

域间注意力:超越 词元级 键值记忆

Interdomain Attention: Beyond Token-Level Key-Value Memory

模型架构混合架构

摘要

Transformer 和深度状态空间模型 (SSM) 位于基本设计选择的两端:注意力通过基于内容的匹配以二次成本将每个查询路由到不断增长的键值 (KV) 缓存,而深度 SSM 将上下文压缩为固定大小的循环状态,该状态不会通过查询键匹配直接寻址。我们提出了域间注意力机制,它通过核方法将 SSM 集成到注意力模块中:注意力核由有限特征图近似,得到的关键特征和值被投影到由单个 SSM 递归维护的共享基函数集上,每个查询通过自己的特征图关注压缩系数,在固定大小的状态上恢复查询条件注意力。可扩展层是这种推导的学习松弛,我们通过消融来验证其组件。在 FineWeb-Edu 上进行的 125M 到 1.3B 自回归语言建模研究中,在匹配的循环状态预算下,域间注意力在各个规模上都改进了 SSM 词元混合器,在验证困惑度和八任务常识套件上超过了 1.3B 的相同配方 Softmax 基线,并继承了其固定状态核心的长度平坦行为,达到训练上下文的 3.5 倍。消融表明查询条件投影是增益的主要来源。