用核方法将SSM融入注意力,实现可扩展的长序列建模。
Interdomain Attention: Beyond Token-Level Key-Value Memory
- 通过核映射将键值压缩到固定大小状态,查询动态调用该状态。
- 在1.3B模型上超越softmax基线,验证困惑度和常识任务表现。
- 适合需要长上下文且资源受限的生成式AI场景。
Transformer与深度状态空间模型(SSM)代表了两种不同的设计范式:注意力机制通过内容匹配以二次代价遍历不断增长的键值缓存,而深度SSM则将上下文压缩为固定大小的递归状态,不依赖查询-键匹配。本文提出跨域注意力(Interdomain Attention),通过核方法将SSM集成到注意力模块中:将注意力核近似为有限特征映射,键特征与值投影至由单一SSM维护的共享基函数空间,每个查询通过其特征映射访问压缩后的系数,从而在固定大小状态上恢复查询条件化的注意力。可扩展层是此推导的可学习松弛形式,通过消融实验验证各组件有效性。在FineWeb-Edu数据集上,从125M到1.3B参数量的自回归语言建模实验中,跨域注意力在相同递归状态预算下优于所有规模的SSM混合器,在1.3B模型上超过同配方softmax基线,在验证困惑度及八项常识任务上表现更优,并保持其核心固定状态的长度无关特性,扩展至训练上下文3.5倍长度。消融分析表明,查询条件化投影是性能提升的主要来源。
原文摘要 · Abstract (English)
Transformers and deep state space models (SSMs) sit at opposite ends of a basic design choice: attention routes each query through a growing key-value (KV) cache by content-based matching at quadratic cost, while deep SSMs compress context into a fixed-size recurrent state that is not directly addressed by query-key matching. We propose Interdomain Attention, which integrates an SSM into an attention module through kernel methods: an attention kernel is approximated by a finite feature map, the resulting key features and values are projected onto a shared set of basis functions maintained by a single SSM recurrence, and each query attends to the compressed coefficients through its own feature map, recovering query-conditioned attention over a fixed-size state. The scalable layer is a learned relaxation of this derivation, and we validate its components through ablations. In a 125M to 1.3B autoregressive language-modeling study on FineWeb-Edu at matched recurrent-state budget, Interdomain Attention improves on an SSM token mixer at every scale, surpasses a same-recipe softmax baseline at 1.3B on validation perplexity and on the eight-task commonsense suite, and inherits the length-flat behavior of its fixed-state core out to 3.5x the training context. Ablations indicate that the query-conditioned projection is the main source of the gain.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。