# FlashKAT: Understanding and Addressing Performance Bottlenecks in the Kolmogorov-Arnold Transformer

arXiv: https://arxiv.org/abs/2505.13813 · HTML: https://mlsanity.com/paper/2505.13813/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 29 references that are in the map.
- First posted on arXiv 20 May 2025.

## Builds on (most cited first, 20 of 20)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [KAN: Kolmogorov-Arnold Networks](https://mlsanity.com/paper/2404.19756/) (2024)
- [U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation](https://mlsanity.com/paper/2406.02918/) (2024)
- [KAN or MLP: A Fairer Comparison](https://mlsanity.com/paper/2407.16674/) (2024)
- [Pad'e Activation Units: End-to-end Learning of Flexible Activation Functions in Deep Networks](https://mlsanity.com/paper/1907.06732/) (2019)
- [Kolmogorov Arnold Informed neural network: A physics-informed deep learning framework for solving forward and inverse problems based on Kolmogorov Arnold Networks](https://mlsanity.com/paper/2406.11045/) (2024)
- [KAN-ODEs: Kolmogorov-Arnold Network Ordinary Differential Equations for Learning Dynamical Systems and Hidden Physics](https://mlsanity.com/paper/2407.04192/) (2024)
- [Kolmogorov-Arnold Transformer](https://mlsanity.com/paper/2409.10594/) (2024)
- [Mimetic Initialization of Self-Attention Layers](https://mlsanity.com/paper/2305.09828/) (2023)
- [Benchmarking and Dissecting the Nvidia Hopper GPU Architecture](https://mlsanity.com/paper/2402.13499/) (2024)
- [Exploring the Limitations of Kolmogorov-Arnold Networks in Classification: Insights to Software Training and Hardware Implementation](https://mlsanity.com/paper/2407.17790/) (2024)
- [MatrixKAN: Parallelized Kolmogorov-Arnold Network](https://mlsanity.com/paper/2502.07176/) (2025)

## Built on by (most cited first, 3 of 3)

- [MetaCluster: Enabling Deep Compression of Kolmogorov-Arnold Network](https://mlsanity.com/paper/2510.19105/) (2025)
- [SparseKAN: Compressing Kolmogorov--Arnold Networks Across Basis Functions, Neurons, and Bits](https://mlsanity.com/paper/2608.00859/) (2026)
- [Rethinking Vision Architectures with Gated Linear Attention and KAN](https://mlsanity.com/paper/2609.22506/) (2026)

## Abstract (arXiv)

The Kolmogorov-Arnold Network (KAN) has been gaining popularity as an alternative to the multilayer perceptron (MLP) due to its greater expressiveness and interpretability. Even so, KAN suffers from training instability and being orders of magnitude slower due to its increased computational cost, limiting its applicability to large-scale tasks. Recently, the Kolmogorov-Arnold Transformer (KAT) has been proposed, achieving FLOPs comparable to traditional Transformer models with MLPs by leveraging Group-Rational KAN (GR-KAN). Unfortunately, despite the comparable FLOPs, our testing shows that KAT remains 123x slower during training, indicating that there are other performance bottlenecks beyond FLOPs. In this paper, we conduct a series of experiments to understand the root cause of the slowdown in KAT. We uncover that the slowdown can be isolated to memory stalls, linked more specifically to inefficient gradient accumulations in the backward pass of GR-KAN. To address this memory bottleneck, we propose FlashKAT, which minimizes accesses to slow memory and the usage of atomic adds through a restructured kernel. Evaluations show that FlashKAT achieves up to an 86.5x training speedup over state-of-the-art KAT while reducing rounding errors in gradient computation.