Cubit: 基于核岭回归的 Token 混合器
机器学习
2026-05-20 v2 计算与语言
摘要
自 2017 年提出以来，Transformer 已成为现代深度学习中最广泛采用的架构之一。尽管在改进位置编码、注意力机制与前馈网络方面付出了大量努力，Transformer 中的核心 token 混合机制仍是注意力。在本工作中，我们表明 Transformer 中的注意力模块可被解释为执行 Nadaraya-Watson 回归，即计算 token 间的相似度并据此聚合相应的值。受此视角启发，我们提出 Cubit，一种利用核岭回归 (KRR) 的潜在下一代架构，而原始 Transformer 依赖 Nadaraya-Watson 回归。具体而言，Cubit 通过引入 KRR 的闭式解来修改经典注意力计算，将通过核相似度的值聚合与通过核矩阵逆的归一化相结合。为提高训练稳定性，我们进一步提出有限范围重标度 (LRR)，其在受控范围内重标度值层。我们认为，作为基于 KRR 的架构，Cubit 比原始 Transformer 提供了更强的数学基础，后者的注意力机制对应于 Nadaraya-Watson 回归。我们通过全面实验验证了这一主张。实验结果表明，Cubit 可能表现出更强的长序列建模能力。特别是，其相对于 Transformer 的性能增益似乎随训练序列长度的增长而增加。
引用
@article{arxiv.2605.06501,
title = {Cubit: Token Mixer with Kernel Ridge Regression},
author = {Chuanyang Zheng and Jiankai Sun and Yihang Gao and Yuehao Wang and Liangchen Tan and Mac Schwager and Anderson Schneider and Yuriy Nevmyvaka and Xiaodong Liu},
journal= {arXiv preprint arXiv:2605.06501},
year = {2026}
}
备注
Tech Report