Swapping MLPs for rational KANs in Vision Transformers

What it does

KAT is a Vision Transformer that replaces the standard MLP blocks with Kolmogorov-Arnold Network (KAN) layers. The authors provide PyTorch training code, pre-trained ImageNet checkpoints ranging from tiny to base size, and both CUDA and Triton kernels for the underlying rational activation functions. It is a full research release aimed at making KANs practical in large-scale vision models.

The interesting bit

Instead of the B-spline basis functions common in KAN literature, the project uses CUDA-implemented rational functions as the base activation. It also introduces “Group KAN” to share weights across edge groups, which appears to be the mechanism that keeps parameter counts and training stability comparable to standard Vision Transformers on ImageNet.

Key highlights

• ICLR 2025 implementation with pre-trained checkpoints (KAT-T, KAT-S, KAT-B) and reported ImageNet top-1 accuracies.
• Replaces MLPs with KAN layers using rational activations rather than B-splines.
• Group KAN shares weights among edge groups to keep the model efficient.
• Ships with both a CUDA kernel and a more accessible Triton implementation.
• Built on timm and supports standard ImageNet training recipes.

Caveats

• Code for detection, segmentation, and NLP tasks is listed as not yet released.
• The README focuses heavily on ImageNet classification; broader vision tasks remain unchecked on the roadmap.

Verdict

Worth a look if you are experimenting with KAN architectures or hunting for an MLP replacement in Vision Transformers. Skip it if you need a general-purpose language model or unified multi-task framework today.