Rethinking Vision Architectures with Gated Linear Attention and KAN

Ali Mehizel; Oussama Khaldi

cs.CVSep 18, 2026

Rethinking Vision Architectures with Gated Linear Attention and KAN

Authors: Ali Mehizel, Oussama Khaldi

Organizations: Department of Applied Statistics, ENSSEA, Kolea, Tipaza, Algeria

Abstract

Vision Transformers devote most of their parameters to MLPs for channel mixing, but still rely on quadratic multi-head self-attention for token interactions. While linear attention fixes the complexity problem, bringing it down to O(N), it is usually just paired with the same fixed-activation MLP as before. Kolmogorov-Arnold Networks take a different approach, placing learnable univariate functions on the edges instead. However, existing vision KANs either retain standard attention or remove attention entirely, so the two ideas have not been effectively combined. We introduce LKAT (Linear Kolmogorov-Arnold Transformer) to close this gap: an isotropic ViT-style encoder that couples chunk-wise Gated Linear Attention with a two-layer KAN feed-forward block, backed by an I/O-aware fused RBF-KAN kernel to make radial-basis grid functions efficient in practice. Under a shared DeiT-style training recipe, LKAT-B outperforms ViT-B/16, ViT-5-B, and Mixer-B/16 on ImageNet-100, while Tiny, Small, and Base variants scale consistently on CIFAR-10/100. ImageNet-100 pretraining also transfers effectively to CIFAR fine-tuning, suggesting that gated linear attention and KAN-based radial basis functions provide complementary inductive biases for mid-scale visual representation learning. Code: https://github.com/mehizelali/linear-kan-transformer

Figures & tables

Input: flattened $x\in\mathbb{R}^{N}$ in HBM
	centers $c\in\mathbb{R}^{G}$ , weights $w\in\mathbb{R}^{G}$ in HBM
	$\mathrm{inv}\leftarrow 1/d$ ; block size $\mathrm{BLOCK}$
Output: $y\in\mathbb{R}^{N}$ in HBM
1:	for each program $p$ in parallel do
2:		Load tile $x[i]$ , $i\in\mathrm{Block}(p)$ (HBM $\to$ on-chip)
3:		$\mathrm{acc}\leftarrow 0$ (registers)

Algorithm 1 Fused RBF–grid linear forward

Figure 1: Peak memory reduction of the fused Triton RBF–grid linear operator versus an unfused torch baseline for a joint forward+backward pass, as a function of channel width $D$ . Bars report $\mathrm{Mem}_{\mathrm{torch}}/\mathrm{Mem}_{\mathrm{Triton}}$ at grid sizes $G\in\{4,8,16\}$ (dotted line ${=}1$ ). The reduction is essentially independent of $D$ and scales nearly linearly with $G$ ( $\approx 5.5\times$ , $10.5\times$ , and $20.5\times$ ), consistent with avoiding materialization of the intermediate $[\ldots,D,G]$ expansion in HBM.

Figure 2: LKAT architecture. Patch embedding feeds $N$ pre-norm blocks (LayerNorm then Gated Linear Attention with per-head RMSNorm on Q/K; LayerNorm then KAN-based radial-basis functions), each with a residual skip. A linear head reads the final class token.

Panel of Figure 3 — Figure 3: Learned RBF–KAN edge functions $\phi_{p,q}(x)$ for ImageNet-100 LKAT-B (representative blocks; full set in Figs. 9 – 10 ).

Hyperparameter	Value
Resolution / grid / bases	$224^{2}$ / $G{=}4$ / RBF
Registers / RoPE	$4$ / 2D RoPE ( 2dv1 )
Optimizer	AdamW
Peak LR / weight decay	$3{\times}10^{-4}$ / $0$
Warmup / min. LR	$5$ epochs / $10^{-5}$
Global batch	$256$ / $512$

Table 1: Shared from-scratch experimental setup. All models (LKAT, rbfKAT, ViT, ViT-5, MLP-Mixer) use this recipe; only the block operators differ.

Model	Params	GFLOPs	CIFAR-10	CIFAR-100	ImageNet100
LKAT (our)
LKAT-T/16	6.25	2.40	82.24	62.00	69.74
LKAT-S/16	23.99	9.42	85.57	67.62	76.66
LKAT-B/16	93.99	37.34	88.43	70.25	80.14
LKAT-L/16	329.47	132.14	90.85	72.10	83.16
rbfKAT (our)

Table 2: From-scratch Top-1 Acc@1 (%). Params / GFLOPs at

224^{2}

Model	Params	GFLOPs	Dim	Heads	Depth
LKAT-T	6.25	2.40	192	3	12
LKAT-S	23.99	9.42	384	6	12
LKAT-B	93.99	37.34	768	12	12
LKAT-L	329.47	132.14	1024	16	24

Table 3: LKAT variants at

224^{2}

. Params (M) and GFLOPs for one forward pass.

Figure 4: ImageNet-100 Top-1 Acc@1 versus GFLOPs ( $224^{2}$ ).

Figure 5: Gated residual around the RBF–KAN block (Eqs. ( 9 )–( 12 )).

Gated residual
Variant	Params (M)	GFLOPs	Acc@1
LKAT-B	93.99	37.34	80.14
+ gate	149.96	60.10	82.02
RBF grid size $G$
$G{=}4$ (default)	93.99	37.34	80.14
$G{=}6$	93.29	37.34	81.08

Table 4: LKAT-B ImageNet-100 ablations (

224^{2}

, 300 epochs). Bold denotes the best Acc@1 in each block.

Setting	CIFAR-10	CIFAR-100
$\mathcal{R}_{\mathrm{A}}$ , $224^{2}$	93.01	78.20
$\mathcal{R}_{\mathrm{B}}$ , $224^{2}$	93.62	78.28
$\mathcal{R}_{\mathrm{B}}$ , $384^{2}$	95.45	78.63

Table 5: CIFAR fine-tuning of ImageNet-100 LKAT-B.

Shared
Setting	$\mathcal{R}_{\mathrm{A}}$	$\mathcal{R}_{\mathrm{B}}$
Init	ImageNet-100 pretrained encoder
Optimizer / peak LR	AdamW / $10^{-4}$
Resolution / aug.	$224^{2}$ / $384^{2}$ / DeiT
Mixup / CutMix	$0.8$ / $1.0$
Differing

Table 6: CIFAR fine-tuning recipe (

\mathcal{R}_{\mathrm{A}}

vs.

\mathcal{R}_{\mathrm{B}}

Require: $x,c,w,d$ as in Alg. 1 ;
	upstream $\mathrm{d}y\in\mathbb{R}^{N}$
Ensure: $\mathrm{d}x\in\mathbb{R}^{N}$ , $\mathrm{d}w\in\mathbb{R}^{G}$
1:	$\mathrm{inv}\leftarrow 1/d$ ; $\mathrm{inv}^{2}\leftarrow 1/d^{2}$ ; $\mathrm{d}w\leftarrow 0$
2:	for each program $p$ in parallel do
3:		Load tile $x[i]$ , $\mathrm{d}y[i]$
4:		$\mathrm{jac}\leftarrow 0$

Algorithm 2 Fused RBF–grid linear backward

Figure 8: End-to-end wall-clock speedup of the fused Triton RBF–grid linear operator versus an unfused torch baseline for a joint forward+backward pass, as a function of channel width $D$ . Bars report $\mathrm{Time}_{\mathrm{torch}}/\mathrm{Time}_{\mathrm{Triton}}$ at grid sizes $G\in\{4,8,16\}$ (dotted line ${=}1$ ). Unlike peak-memory reduction, which is nearly constant in $D$ , e2e speedup grows with both $D$ and $G$ : from about $1.4$ – $1.7\times$ at $D{=}128$ up to about $17\times$ / $29\times$ / $44\times$ at $D{=}4096$ for $G{=}4/8/16$ . Larger grids amplify the benefit of avoiding the materialised $[\ldots,D,G]$ intermediate, so the fused kernel’s latency stays nearly flat while torch scales with $G$ .

Panel of Figure 9 — Figure 9: Learned RBF–KAN edge functions $\phi_{p,q}(x)$ for ImageNet-100 LKAT-B, first KAN layer (fc1), blocks $0$ – $11$ (panel labels).

Panel of Figure 10 — Figure 10: Learned RBF–KAN edge functions $\phi_{p,q}(x)$ for ImageNet-100 LKAT-B, second KAN layer (fc2), blocks $0$ – $11$ (panel labels).

Explore similar work

Exact Linear Attention
May 13, 2026Weinuo OuKimi Delta Attention Transformer Attention
Wuyi University
Copy the Same, Distill the Difference: Initializing Linear Vision Transformers
Sep 28, 2026Huaiyuan Qin, Muli Yang, Gabriel James Goenawan +9Self-Supervised Vision Transformers Transformer Architectures
Institute of Advanced Intelligence and Computing (IAIC), A*STAR, Singapore · Nanyang Technological University · ST Engineering Geo-Insights, Singapore +3
Representative Attention For Vision Transformers
May 14, 2026Yuntong Li, Hainuo Wang, Hengxing Liu +2Self-Supervised Vision Transformers Vision Transformer
College of Intelligence and Computing, Tianjin University, Tianjin 300350, China

May 13, 2026, cs.LGJ/K move · Enter open · S save

Exact Linear Attention

Weinuo Ou

Wuyi University

Kimi Delta Attention Transformer Attention Transformer Architectures Kernel Method

Open paper AskPDF arXiv

This paper introduces Exact Linear Attention (ELA), a mechanism that achieves linear computational complexity for Transformer attention by exploiting the exact decomposition property of kernel functions, thereby eliminating approximation error. We identify and address two key limitations of prior linear attention -- gradient explosion and token attention dilution -- by imposing kernel constraints that ensure non-negativity, discriminability, and geometric interpretability. Several kernel functions are proposed, including the Hadamard Exp Kernel, Summation Squared Euclidean Distance Kernel, and Subtraction Squared Euclidean Distance Kernel, each tailored for specific attention behaviors. Beyond the core attention formulation, the paper presents three engineering innovations: (1) a Hyper-Link structure that replaces traditional residual connections to mitigate gradient degradation; (2) a Memory Lobe module based on bidirectional linear attention, which captures "transformation flow" across layers to implement qualitative memory and an implicit reinforcement learning paradigm; and (3) a routing-score-based bias mechanism for Mixture-of-Experts (MoE) to improve interpretability and semantic alignment. Experimental results demonstrate that ELA achieves up to 6x faster decoding speed and 75% reduction in KV cache memory usage compared to full attention, while maintaining comparable or superior training performance. The proposed memory module accelerates convergence and enhances generalization. Furthermore, we extend the linear attention principle to vision models, yielding YOLO-LAT, which attains up to 4.3x GPU inference speedup and 7.9x parameter reduction with competitive detection accuracy. These results underline the broad applicability of exact linear attention for scaling Transformer models to ultra-long sequences and efficient visual tasks.

Abstract

Explore similar work

Exact Linear Attention

Copy the Same, Distill the Difference: Initializing Linear Vision Transformers

Representative Attention For Vision Transformers