ICLR 2025
Yes, I kan!
🎉 This is a PyTorch/GPU implementation of the paper Kolmogorov–Arnold Transformer (KAT), which replace the MLP layers in transformer with KAN layers.
For more technical details, please refer to our ICLR'25 paper.
Kolmogorov–Arnold Transformer
📝[Paper] </>[code] </>[Trition/CUDA kernel]
Xingyi Yang, Xinchao Wang
National University of Singapore
International Conference on Learning Representations (ICLR'25)
Vanilla ViT + KAN struggle to scale effectively. We introduce the KAT model, which integrates GR-KANs into transformers for large-scale training scenarios like ImageNet, achieving significant performance improvements.
- Base Function: Replace B-spline to CUDA-implemented Rational.
- Group KAN: Share weights among groups of edges for efficiency.
- Initialization: Maintain activation magnitudes across layers.
- Release the KAT paper, CUDA implementation and IN-1k training code.
- 🎉🎉🎉🎉 Triton Implementation, on 1D and 2D tasks. This is much easier to install than the CUDA version. Please See https://github.com/Adamdad/rational_kat_cu.
- KAT Detection and segmentation code.
- KAT on NLP tasks.
Please find our CUDA implementation in https://github.com/Adamdad/rational_kat_cu.git.
# install torch and other things pip install timm==1.0.3 pip install wandb # I personally use wandb for results visualizations git clone https://github.com/Adamdad/rational_kat_cu.git cd rational_kat_cu pip install -e .📦 Data preparation: ImageNet with the following folder structure, you can extract ImageNet by this script
│imagenet/ ├──train/ │ ├── n01440764 │ │ ├── n01440764_10026.JPEG │ │ ├── n01440764_10027.JPEG │ │ ├── ...... │ ├── ...... ├──val/ │ ├── n01440764 │ │ ├── ILSVRC2012_val_00000293.JPEG │ │ ├── ILSVRC2012_val_00002138.JPEG │ │ ├── ...... │ ├── ...... Refer to example.py for a detailed use case demonstrating how to use KAT with timm to classify an image.
Download pre-trained models or access training checkpoints:
| 🏷️ Model | ⚙️ Setup | 📦 Param | 📈 Top1 | 🔗 Link |
|---|---|---|---|---|
| KAT-T | From Scratch | 5.7M | 74.6 | link/huggingface |
| KAT-T | From ViT | 5.7M | 75.7 | link/huggingface |
| KAT-S | From Scratch | 22.1M | 81.2 | link/huggingface |
| KAT-S | From ViT | 22.1M | 82.0 | link/huggingface |
| KAT-B | From Scratch | 86.6M | 82.3 | link/huggingface |
| KAT-B | From ViT | 86.6M | 82.8 | link/huggingface |
All training scripts are under scripts/
bash scripts/train_kat_tiny_8x128.shIf you want to change the hyper-parameters, can edit
#!/bin/bash DATA_PATH=/local_home/dataset/imagenet/ bash ./dist_train.sh 8 $DATA_PATH \ --model kat_tiny_swish_patch16_224 \ # Rationals are initialized to be swish functions -b 128 \ --opt adamw \ --lr 1e-3 \ --weight-decay 0.05 \ --epochs 300 \ --mixup 0.8 \ --cutmix 1.0 \ --sched cosine \ --smoothing 0.1 \ --drop-path 0.1 \ --aa rand-m9-mstd0.5 \ --remode pixel --reprob 0.25 \ --amp \ --crop-pct 0.875 \ --mean 0.485 0.456 0.406 \ --std 0.229 0.224 0.225 \ --model-ema \ --model-ema-decay 0.9999 \ --output output/kat_tiny_swish_patch16_224 \ --log-wandbTo evaluate our kat_tiny_patch16_224 models, run:
DATA_PATH=/local_home/dataset/imagenet/ CHECKPOINT_PATH=kat_tiny_patch16_224_1f3ad3b2e69821f3d412f2924cf159a0e266f142d739cb68f68f796f5a0fe289.pth python validate.py $DATA_PATH --model kat_tiny_patch16_224 \ --checkpoint $CHECKPOINT_PATH -b 512 ################### Validating in float32. AMP not enabled. Loaded state_dict from checkpoint 'kat_tiny_patch16_224_1f3ad3b2e69821f3d412f2924cf159a0e266f142d739cb68f68f796f5a0fe289.pth' Model kat_tiny_patch16_224 created, param count: 5718328 Data processing configuration for current model + dataset: input_size: (3, 224, 224) interpolation: bicubic mean: (0.485, 0.456, 0.406) std: (0.229, 0.224, 0.225) crop_pct: 0.875 crop_mode: center Test: [ 0/98] Time: 3.453s (3.453s, 148.28/s) Loss: 0.6989 (0.6989) Acc@1: 84.375 ( 84.375) Acc@5: 96.875 ( 96.875) ....... Test: [ 90/98] Time: 0.212s (0.592s, 864.23/s) Loss: 1.1640 (1.1143) Acc@1: 71.875 ( 74.270) Acc@5: 93.750 ( 92.220) * Acc@1 74.558 (25.442) Acc@5 92.390 (7.610) --result { "model": "kat_tiny_patch16_224", "top1": 74.558, "top1_err": 25.442, "top5": 92.39, "top5_err": 7.61, "param_count": 5.72, "img_size": 224, "crop_pct": 0.875, "interpolation": "bicubic" }We extend our gratitude to the authors of rational_activations for their contributions to CUDA rational function implementations that inspired parts of this work. We thank @yuweihao, @florinshen, @Huage001 and @yu-rp for valuable discussions.
If you use this repository, please cite:
@inproceedings{ yang2025kolmogorovarnold, title={Kolmogorov-Arnold Transformer}, author={Xingyi Yang, Xinchao Wang}, booktitle={The Thirteenth International Conference on Learning Representations}, year={2025}, url={https://openreview.net/forum?id=BCeock53nt} }