[ARCH-PROP] Cubit: Token Mixer with Kernel Ridge Regression

[chuanyang-Zheng]

Architecture Name

Cubit: Use $Kernel(QK^T)(Kernel(KK^T))^{-1}V$ to replace $Kernel(QK^T)V$

Parent issue

#1

Motivations

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.

Proposed Architecture

class Cubit(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        attention_head: int,
        eps: float = 1e-10,
        upper: float = 2.0,
        lower: float = 0.5,
        share: bool = True,
        causal_mask: bool = True
    ):
        super().__init__()

        self.hidden_size = hidden_size
        self.attention_head = attention_head
        self.hidden_size_per_head = hidden_size // attention_head
        self.share = share
        self.causal_mask = causal_mask
        self.eps = eps

        # Linear projections for Q, K, V
        self.q = nn.Linear(hidden_size, hidden_size)
        self.k = nn.Linear(hidden_size, hidden_size)
        self.v = nn.Linear(hidden_size, hidden_size)

        # Optional separate R projection
        if not share:
            self.r = nn.Linear(hidden_size, hidden_size)
        

        # Learnable bounds for LR scale - shape: (1, attention_head, 1, 1)
        self.lower = nn.Parameter(torch.full((1, attention_head, 1, 1), lower))
        self.upper_scale = nn.Parameter(torch.full((1, attention_head, 1, 1), upper - lower))

        # LR scaling factor per head
        self.LRR = nn.Linear(hidden_size, attention_head)

        # Learnable scale for normalized R - shape: (1, attention_head, 1, 1)
        self.scale = nn.Parameter(torch.ones(1, attention_head, 1, 1))

        # Regularization strength lambda (log space for stability) - shape: (1, attention_head, 1, 1)
        self.log_lambda = nn.Parameter(
            torch.full((1, attention_head, 1, 1), math.log(eps)),
            requires_grad=True
        )



    def forward(self, x, pos_encoding_function,softmax, mask) -> torch.Tensor:
        b, t, d = x.shape
        device = x.device

        # Q, K, V, R projections
        q=self.q(x).reshape(b,t,self.attention_head,self.hidden_size_per_head).permute(0,2,1,3)
        k=self.k(x).reshape(b,t,self.attention_head,self.hidden_size_per_head).permute(0,2,1,3) 
        v=self.v(x).reshape(b,t,self.attention_head,self.hidden_size_per_head).permute(0,2,1,3) 
        # Determine R: separate if share=False, else reuse K
        if self.share:
             r=k
        else:
            r=self.r(x).reshape(b,t,self.attention_head,self.hidden_size_per_head).permute(0,2,1,3)
        # LR factor: sigmoid-scaled between lower and upper
        # lrr_logits shape: (B, n_head, T, 1)
        lrr_logits = self.lrr(x).reshape(b, t, self.attention_head, 1).permute(0, 2, 1, 3)
        # lrr shape: (B, n_head, T, 1)
        lrr = self.lower + self.upper_scale * torch.sigmoid(lrr_logits)




        # Normalize R with learnable scale
        # norm_r shape: (B, n_head, T, head_dim)
        norm_r = r / torch.norm(r, dim=-1, p=2, 
        keepdim=True) * self.scale
        
        # add position encoding
        q,k,r,norm_r=pos_encoding_function(q,k,r,norm_r)

        # Sigma inverse: softmax similarity + lambda * I
        # r @ norm_r.transpose(-2, -1) shape: (B, n_head, T, T)
        sigma_inv = softmax(r @ norm_r.transpose(-2, -1), mask)
        
        # I shape: (1, 1, T, T) -> broadcast to (B, n_head, T, T)
        I = torch.eye(t, device=device).unsqueeze(0).unsqueeze(0)
        lambda_reg = torch.exp(self.log_lambda)  # shape: (1, n_head, 1, 1)
        sigma_inv = sigma_inv + lambda_reg * I

        # Solve linear system for each head
        # rhs shape: (B, n_head, T, head_dim)
        rhs = lrr * v
        

        # solution shape: (B, n_head, T, head_dim)

        if self.casual_mask:
            solution = torch.linalg.solve_triangular(sigma_inv, rhs,upper=False)
        else:
            solution = torch.linalg.solve(sigma_inv, rhs)
        


        # Standard attention with Q, K
        A_weights = softmax(
            q @ k.transpose(-2, -1) / math.sqrt(self.hidden_size_per_head), 
            mask
        )  # shape: (B, n_head, T, T)

        # Final output
        output = A_weights @ solution  # shape: (B, n_head, T, head_dim)
        
        # Reshape back to (B, T, hidden_size)
        output = output.permute(0, 2, 1, 3).reshape(b, t, self.hidden_size)
        
        return output

Preliminary Results (if any)

https://arxiv.org/pdf/2605.06501

Experiments Plan

The primary results already give detailed evaluation. We hope that industry could further push it forward.

Ridge Regression: Linear Attention
Nadaraya-Watson Regression: Transformer
Kernel Ridge Regression (Minmax Optimal): Cubit

[JT-Ushio]

Hi, @chuanyang-Zheng, thanks for bringing Cubit to ArchSpace🌳 and for your interest in the project!

First, sorry for the late reply. ArchSpace already provides a ~6T-token pipeline for large-scale, reproducible, and comparable architecture experiments. Over the past few weeks, we have also been developing a lighter ~350B-token pipeline(covering pretraining, mid-training, and SFT) to make it easier and faster for the community to validate new proposals. The related infrastructure is being developed in PR #21.

Our thoughts on Cubit

We find Cubit an interesting direction. Its key idea is to reinterpret standard attention as Nadaraya–Watson regression and replace it with a Kernel Ridge Regression (KRR) formulation. In addition to the usual query-context interaction, Cubit introduces a context-context correction through the KRR kernel matrix.

From an architectural perspective, we believe Cubit is feasible to integrate into the OLMo-core baseline. The modification is largely localized to the attention/token-mixing module, while the rest of the Transformer and training pipeline can remain mostly unchanged.

The main concern is currently system efficiency. A straightforward implementation requires constructing an additional ($T\times T$) kernel matrix and performing a causal triangular solve. This introduces substantial overhead compared with highly optimized FlashAttention-based Transformers.

However, we do not view this as a fundamental blocker. Since the causal KRR matrix is lower triangular, there may be opportunities to implement the computation using a blockwise online-softmax + triangular-solve kernel, avoiding explicit materialization of the full ($T\times T$) matrix. Efficient kernel support would therefore be an important follow-up if the architecture shows promising results.

For the initial evaluation, we are particularly interested in:

• whether the reported improvement over standard attention is reproducible under our controlled training setup;
• whether the advantage becomes more significant with longer contexts (~64K) and larger models (1/3/7B @ 350B/6T);
• the numerical stability and actual wall-clock overhead of the KRR component.

Decision

Overall, we are very interested in Cubit's empirical performance and would like to move this proposal to In Progress.

For the next step, we see two possible collaboration modes:

1. You implement Cubit on top of our OLMo-core baseline.
   We can provide guidance on the baseline and expected interfaces, and then launch the experiments once the implementation is reviewed.

2. We implement Cubit, invite you to review the implementation and Cubit-specific design choices, and then launch the experiments.

Please let us know which option you prefer. Feel free to contact us at any point during the implementation or evaluation.

Thanks again for contributing Cubit to ArchSpace🌳 — we look forward to seeing how it performs under a controlled architecture comparison.

[chuanyang-Zheng]

Thank you for the support. We will further address the concerns below.

Q1: whether the reported improvement over standard attention is reproducible under our controlled training setup;

A1: We have run the experiment with 1.3B model size and 50 training tokens with length 1024, the Cubit achieves about lower loss gain and better downstream task performance.

Q2: Thether the advantage becomes more significant with longer contexts (~64K) and larger models (1/3/7B @ 350B/6T);

A2: We run 125M, 350M and 1.3B model size with training length 512, 1024, 2048, 4096 and 8192, with training token 1B. The performance gain gradually increases with longer training length. For here: https://arxiv.org/pdf/2605.06501 Figure 3

Q3: the numerical stability

A3: For numerical stability: Remember the equation is $Kernel(QK^T)Kernel(Norm(K)Norm(K)^T)^{-1}V$ or $Kernel(QK^T)Kernel(K Norm(K)^T)^{-1}V$. If not normalizing the key embedding, the diagonal value may be too small so that the condition number is too large, the forward substitution may cause NAN.

Another potential direction is using an RBF kernel to replace Softmax, but the performance may be worse. In this situation (with RBF), there is no numerical stability problem, but the performance may be worse. Currently, we are still unsure why this happens, because theoretically RBF is suitable KRR and NW Regression. The potential problem may be that RBF is not as well studied as Softmax, so we only have a few references to the implementation of RBF.

$RBF(x_i, x_j) = \exp(-\gamma |x_i - x_j|^2)$. With model size 125M and length 512

Kernel	Regression	Kernel(QK)	Kernel(KK)	5K	10K	15K	20K	25K	30K	35K	40K	45K	50K
RBF	Transformer (Nadaraya-Watson Regression)	RBF $( \gamma = 1/\sqrt{8D})$	-	4.4837	4.2662	4.1209	4.0741	3.9985	3.9430	3.8665	3.8371	3.7839	3.7508
RBF	Cubit (KRR)	RBF $( \gamma = 1)$	RBF $( \gamma = 1)$	5.9308	4.9975	4.5893	4.3778	4.2528	4.1703	4.0820	4.0459	3.9860	3.9499
RBF	Cubit (KRR)	RBF $( \gamma = 1/\sqrt{8D})$	RBF $( \gamma = 1/\sqrt{8D})$	4.4763	4.2527	4.1002	4.0508	3.9709	3.9132	3.8386	3.8095	3.7558	3.7242
Softmax	Transformer (Nadaraya-Watson Regression)	Softmax	-	4.4865	4.2676	4.1182	4.0698	3.9963	3.9339	3.8610	3.8313	3.7778	3.7443
Softmax	Cubit (KRR)	Softmax	Softmax	4.4295	4.2254	4.0799	4.0319	3.9554	3.9003	3.8251	3.7961	3.7432	3.7110

Q4: For long-context: experiment and theoretical

A4: The following is the results without flashattention and only pytorch implementation with GPT-NeoX codebae, with model size 125M.

Method	Cost	Length 512	Length 1024	Length 2048	Length 4096	Length 8192
Transformer	Time Cost (ms)	326.3	352.0	445.3	790.7	2078.3
	Memory Cost (MB)	7352.0	7552.0	8774.0	15980.0	39368.0
DeltaFormer	Time Cost (ms)	459.6	567.3	823.0	1766.0	5307.0
	Memory Cost (MB)	7352.0	7868.0	9542.0	19052.0	51704.0
Cubit	Time Cost (ms)	468.2	575.8	835.3	1761.7	5274.5
	Memory Cost (MB	7352.0	7448.0	9578.0	20180.0	51656.0

The text model erros are related to associate recall. When $q=k_i$, we may need the perfect $v_i$

$\text{NW}(q=k_t,k,v) = \frac{\text{kernel}(k_t, k_t)}{\text{kernel}(k_t, k_t) + \sum_{j \neq t} \text{kernel}(k_t, k_j)}v_t +\sum_{i \neq t} \frac{\text{kernel}(k_t, k_i)}{\text{kernel}(k_t, k_t) + \sum_{j \neq t} \text{kernel}(k_t, k_j)} v_i$

With the context length increase, the $\frac{\text{kernel}(k_t, k_t)}{\text{kernel}(k_t, k_t) + \sum_{j \neq t} \text{kernel}(k_t, k_j)}$ will decreases, suggests that the the needed $v_t$ will be difficult to be retrive. If the context is sufficient long, the retrived value is only noise

For KRR, we have th following

$\text{KRR}(q=k_t,k,v)=(qk)(kk^T)^-1v=(k_tk)(kk^T)v=v_t$.

Therefore, when $q=k_t$, we cold perfect have $v_t$, whatever the hidden size, hidden dimen and so on. The long context ability is only bounded by the computation precision

About Next Step:

I am still learning how to write cuda and use triton. Currently, I find that DeltaFormer fla implementation may be helpful. https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/deltaformer/parallel.py#L210

DeltaFormer
* Motivation: combine DeltaNet and Transformer
* Implementation: solution = torch.linalg.solve_triangular(sigma_inv, rhs,upper=False,unitrigular=True) to make the training stable; sigma_inv=Kernel(KW), while W is another projection like Q, K, V
Cubit

• Motivation: Transformer is NW Regression, so that we could use KRR to replace NW Regression. We do not have to care about the solution stability as long as we have sigma_inv=Kernel(KK) to have the Gram matrix.
• Implementation: solution = torch.linalg.solve_triangular(sigma_inv, rhs,upper=False); sigma_inv=Kernel(KK)

Note: The DeltaFormer fla implementation is chunk-wise (just like DeltaFormer) so that it may destroy the theoretical guarantee of KRR, but just an approximation.

I could directly give the pytorch implementation. As I am still learning Triton and CUDA, Step 2 may be better.

If there are any questions that I could help with, please let me know.

[JT-Ushio]

Looks good to me. We will proceed with a vibe-coded implementation of Cubit on the arch/cubit branch and invite you to review it once the initial version is ready. Looking forward to collaborating with you!

One small clarification regarding the following part:

For the initial evaluation, we are particularly interested in:

• whether the reported improvement over standard attention is reproducible under our controlled training setup;
• whether the advantage becomes more significant with longer contexts (~64K) and larger models (1/3/7B @ 350B/6T);
• the numerical stability and actual wall-clock overhead of the KRR component.

These are our intended evaluation targets using AILab resources (including access to thousands of Ascend NPUs and NVIDIA GPUs) under an OLMo3-like training pipeline, rather than requirements that need to be addressed at the proposal stage.

In other words, the proposal itself does not need to provide evidence at these scales in advance. We plan to investigate these questions during the ArchSpace implementation and evaluation process.

Looking forward to working together!

[chuanyang-Zheng]

Thank you for the detailed response. If there is anything that I could help, please let me know