Cubit: Token Mixer with Kernel Ridge Regression
cs.LG, cs.CL
Submitted: 2026-05-07
Updated: 2026-08-26
Comments: Tech Report
License: http://creativecommons.org/licenses/by/4.0/
The gist: Since its introduction in 2017, the Transformer has become one of the most widely adopted architectures in modern deep learning.
Terminology
Abstract
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.
Sources
- Neural Machine Translation by Jointly Learning to Align and Translate
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
- Efficiently Modeling Long Sequences with Structured State Spaces
- Mixtral of Experts
- Multiplicative LSTM for sequence modelling
- Virtual Width Networks
- DeepSeek-V3 Technical Report
- Optimal Rates and Saturation for Noiseless Kernel Ridge Regression
- Regularized least squares learning with heavy-tailed noise is minimax optimal
- RWKV-7 "Goose" with Expressive Dynamic State Evolution
- HGRN2: Gated Linear RNNs with State Expansion
- Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention
- Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation
- Fast Transformer Decoding: One Write-Head is All You Need
- JetMoE: Reaching Llama2 Performance with 0.1M Dollars
- Simplified State Space Layers for Sequence Modeling
- Retentive Network: A Successor to Transformer for Large Language Models
- Kimi K2: Open Agentic Intelligence
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks