Cubit: Token Mixer with Kernel Ridge Regression

TL;DR

Cubit introduces a novel token-mixing architecture based on Kernel Ridge Regression, offering a mathematically grounded alternative to Transformer attention with improved long-sequence modeling capabilities.

Contribution

The paper proposes Cubit, a new architecture replacing attention with Kernel Ridge Regression, and introduces the Limited-Range Rescale for training stability.

Findings

Cubit shows stronger long-sequence modeling performance.

Performance gain increases with training sequence length.

Provides a more solid mathematical foundation than traditional Transformer attention.

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.