arXiv:2605.06501cs.LGcs.CL2026-05

用核岭回归替代注意力机制,提升长序列建模能力。

Cubit: Token Mixer with Kernel Ridge Regression

  • 以核岭回归重构注意力,数学基础更坚实。
  • 在长序列任务中表现更好,训练长度越长优势越明显。
  • 适合需要强长程依赖建模的NLP与序列任务。

自2017年问世以来,Transformer已成为现代深度学习中最广泛采用的架构之一。尽管在位置编码、注意力机制和前馈网络方面已有大量改进,Transformer的核心令牌混合机制仍依赖于注意力。本文揭示,Transformer中的注意力模块可被解释为执行Nadaraya-Watson回归,通过计算令牌间的相似性并相应聚合对应值。受此启发,我们提出Cubit,一种基于核岭回归(KRR)的潜在下一代架构,而标准Transformer依赖的是Nadaraya-Watson回归。具体而言,Cubit通过引入KRR的闭式解,修改经典注意力计算,将值聚合与核相似性结合,并通过核矩阵逆进行归一化。为进一步提升训练稳定性,我们提出有限范围重缩放(LRR),对值层在可控范围内重缩放。我们认为,基于KRR的Cubit相比依赖Nadaraya-Watson回归的原始Transformer具有更强的数学基础。通过全面实验验证,结果表明Cubit可能具备更强的长序列建模能力,尤其在训练序列长度增加时,其性能优势愈发显著。

原文摘要 · Abstract (English)

Since its introduction in 2017, the Transformer has become one of the most widely adopted architectures in modern deep learning. Despite extensive efforts to improve positional encoding, attention mechanisms, and feed-forward networks, the core token-mixing mechanism in Transformers remains attention. In this work, we show that the attention module in Transformers can be interpreted as performing Nadaraya-Watson regression, where it computes similarities between tokens and aggregates the corresponding values accordingly. Motivated by this perspective, we propose Cubit, a potential next-generation architecture that leverages Kernel Ridge Regression (KRR), while the vanilla Transformer relies on Nadaraya-Watson regression. Specifically, Cubit modifies the classical attention computation by incorporating the closed-form solution of KRR, combining value aggregation through kernel similarities with normalization via the inverse of the kernel matrix. To improve the training stability, we further propose the Limited-Range Rescale (LRR), which rescales the value layer within a controlled range. We argue that Cubit, as a KRR-based architecture, provides a stronger mathematical foundation than the vanilla Transformer, whose attention mechanism corresponds to Nadaraya-Watson regression. We validate this claim through comprehensive experiments. The experimental results suggest that Cubit may exhibit stronger long-sequence modeling capability. In particular, its performance gain over the Transformer appears to increase as the training sequence length grows.

注意力机制核方法序列建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。