Vision transformers improved with new attention and feed-forward design
Rethinking Vision Architectures with Gated Linear Attention and KAN
Computer Vision and Pattern Recognition
Summary
Vision Transformer models usually use a complex attention mechanism that can be slow and rely on fixed ways of processing data. The authors designed a new transformer called LKAT that combines a simpler, faster type of attention called gated linear attention with a different way to mix information using learnable functions. They tested LKAT on image tasks and found it works better than several existing models of similar size. This suggests their new design choices can improve how computers understand pictures.
What this means in practice
- •For computer vision engineers: Build more efficient image classification models by replacing standard attention with gated linear attention and KAN-based feed-forward layers as shown effective in LKAT.
- •For mobile device developers: Develop lightweight vision models that maintain accuracy with reduced computation for deployment on resource-constrained hardware using the LKAT architecture.
Authors
Ali Mehizel, Oussama Khaldi
Abstract
Vision Transformers allocate most parameters to multi-layer perceptrons (MLPs) for channel mixing, while token interactions usually rely on quadratic multi-head self-attention (MHSA). Linear attention reduces sequence complexity to O(N), but remains coupled with the same fixed-activation MLP as softmax Transformers. Kolmogorov-Arnold Networks (KANs) instead place learnable univariate maps on edges, yet prior vision KANs keep MHSA or omit attention entirely. We introduce LKAT (Linear Kolmogorov-Arnold Transformer), an isotropic ViT encoder that couples chunkwise Gated Linear Attention (GLA) with a two-layer KAN feed-forward, and we provide an I/O-aware fused RBF-KAN kernel for the radial-basis grid maps. Under a shared DeiT-style recipe we compare LKAT with ViT, ViT-5, and MLP-Mixer. LKAT-B exceeds ViT-B/16, ViT-5-B, and Mixer-B/16 on ImageNet-100. Tiny/Small/Base LKAT variants scale consistently on CIFAR-10/100, and ImageNet-100 pretraining transfers to CIFAR fine-tuning. The results support gated linear attention and KAN-based radial-basis functions as complementary inductive biases for mid-scale visual representation learning.