Cubit: Token Mixer with Kernel Ridge Regression

• URL: http://arxiv.org/abs/2605.06501v1
• Date: Thu, 07 May 2026 16:18:55 GMT
• Title: Cubit: Token Mixer with Kernel Ridge Regression
• Abstract summary: We show that the attention module in Transformers can be interpreted as performing Nadaraya-Watson regression.<n>We propose Cubit, a potential next-generation architecture that leverages Kernel Ridge Regression (KRR)<n>We argue that Cubit, as a KRR-based architecture, provides a stronger mathematical foundation than the vanilla Transformer.
• Score: 30.306738326416223
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression [13.818160005611752]
  We study in-context kernel ridge regression (KRR) with Gaussian kernels.<n>We show that a standard softmax-attention transformer can approximate the KRR predictor during its forward pass.<n>We find that its error profiles align most consistently with preconditioned Richardson iteration.
  arXiv  Detail & Related papers  (2026-05-08T20:50:58Z)
• Converting Transformers into DGNNs Form [3.7468283401703797]
  We introduce a synthetic unitary digraph convolution based on the digraph Fourier transform.<n>The resulting model, which we term Converter, effectively converts a Transformer into a Directed Graph Neural Network form.<n>We have tested Converter on Long-Range Arena benchmark, long document classification, and DNA sequence-based taxonomy classification.
  arXiv  Detail & Related papers  (2025-02-01T22:44:46Z)
• Autoregressive + Chain of Thought = Recurrent: Recurrence's Role in Language Models' Computability and a Revisit of Recurrent Transformer [29.970200877158764]
  We investigate the influence of recurrent structures in neural models on their reasoning abilities and computability. We shed light on how the CoT approach can mimic recurrent computation and act as a bridge between autoregression and recurrence.
  arXiv  Detail & Related papers  (2024-09-14T00:30:57Z)
• PRformer: Pyramidal Recurrent Transformer for Multivariate Time Series Forecasting [82.03373838627606]
  Self-attention mechanism in Transformer architecture requires positional embeddings to encode temporal order in time series prediction. We argue that this reliance on positional embeddings restricts the Transformer's ability to effectively represent temporal sequences. We present a model integrating PRE with a standard Transformer encoder, demonstrating state-of-the-art performance on various real-world datasets.
  arXiv  Detail & Related papers  (2024-08-20T01:56:07Z)
• On the Convergence of Encoder-only Shallow Transformers [62.639819460956176]
  We build the global convergence theory of encoder-only shallow Transformers under a realistic setting. Our results can pave the way for a better understanding of modern Transformers, particularly on training dynamics.
  arXiv  Detail & Related papers  (2023-11-02T20:03:05Z)
• RWKV: Reinventing RNNs for the Transformer Era [54.716108899349614]
  We propose a novel model architecture that combines the efficient parallelizable training of transformers with the efficient inference of RNNs. We scale our models as large as 14 billion parameters, by far the largest dense RNN ever trained, and find RWKV performs on par with similarly sized Transformers.
  arXiv  Detail & Related papers  (2023-05-22T13:57:41Z)
• Your Transformer May Not be as Powerful as You Expect [88.11364619182773]
  We mathematically analyze the power of RPE-based Transformers regarding whether the model is capable of approximating any continuous sequence-to-sequence functions. We present a negative result by showing there exist continuous sequence-to-sequence functions that RPE-based Transformers cannot approximate no matter how deep and wide the neural network is. We develop a novel attention module, called Universal RPE-based (URPE) Attention, which satisfies the conditions.
  arXiv  Detail & Related papers  (2022-05-26T14:51:30Z)
• nnFormer: Interleaved Transformer for Volumetric Segmentation [50.10441845967601]
  We introduce nnFormer, a powerful segmentation model with an interleaved architecture based on empirical combination of self-attention and convolution. nnFormer achieves tremendous improvements over previous transformer-based methods on two commonly used datasets Synapse and ACDC.
  arXiv  Detail & Related papers  (2021-09-07T17:08:24Z)
• Finetuning Pretrained Transformers into RNNs [81.72974646901136]
  Transformers have outperformed recurrent neural networks (RNNs) in natural language generation. A linear-complexity recurrent variant has proven well suited for autoregressive generation. This work aims to convert a pretrained transformer into its efficient recurrent counterpart.
  arXiv  Detail & Related papers  (2021-03-24T10:50:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.