We are experimenting with the possibility of KAN replacing MLP in Vision Transformer, this project may be delayed for a long time due to GPU resource constraints, if there are any new developments, we will show them here!
pip install VisionKAN
from VisionKAN import create_model, train_one_epoch, evaluate
KAN_model = create_model(
model_name='deit_tiny_patch16_224_KAN',
pretrained=False,
hdim_kan=192,
num_classes=100,
drop_rate=0.0,
drop_path_rate=0.05,
img_size=224,
batch_size=144
)
| Dataset | MLP hidden dim | model | date | epoch | top1 | top5 | Checkpoint |
|---|---|---|---|---|---|---|---|
| ImageNet 1k | 768 | DeiT-tiny(baseline) | - | 300 | 72.2 | 91.1 | |
| CIFAR-100 | 192 | DeiT-tiny(baseline) | 2024.5.25 | 300(stop) | 84.94 | 96.53 | Checkpoint |
| CIFAR-100 | 384 | DeiT-small(baseline) | 2024.5.25 | 300(stop) | 86.49 | 96.17 | Checkpoint |
| CIFAR-100 | 768 | DeiT-base(baseline) | 2024.5.25 | 300(stop) | 86.54 | 96.16 | Checkpoint |
| Dataset | KAN hidden dim | model | date | epoch | top1 | top5 | Checkpoint |
|---|---|---|---|---|---|---|---|
| ImageNet 1k | 20 | Vision-KAN | 2024.5.16 | 37(stop) | 36.34 | 61.48 | - |
| ImageNet 1k | 192 | Vision-KAN | 2024.5.25 | 346(stop) | 64.87 | 86.14 | Checkpoint |
| ImageNet 1k | 768 | Vision-KAN | 2024.6.2 | 154(training) | 62.90 | 85.03 | - |
| CIFAR-100 | 192 | Vision-KAN | 2024.5.25 | 300(stop) | 73.17 | 93.307 | Checkpoint |
| CIFAR-100 | 384 | Vision-KAN | 2024.5.25 | 300(stop) | 78.69 | 94.73 | Checkpoint |
| CIFAR-100 | 768 | Vision-KAN | 2024.5.29 | 300(stop) | 79.82 | 95.42 | Checkpoint |
We released our current Vision KAN code, we used efficient KAN to simply replace the MLP layer in the Transformer block and are pre-training the Tiny model on ImageNet 1k, subsequent results will be updated in the table.
The model has started to converge, we use [192, 20, 192] as input, hidden, and output dimensions, and we reshape the input dimensions in order to fit the processing dimensions of KAN.
we change efficient kan to faster kan to speed up to 2x in training process, and change base model from Deit iii to Deit, so that we can use pre-trained model for most layers except kan layer
The convergence of the model seems to be entering a bottleneck, and I'm guessing that kan's hidden layer setting of 20 is too small, so I'm going to adjust the hidden layer to 192 if it doesn't converge after a few more rounds of running.
Fix Timm version dependency bugs and remove extraneous code.
The decline in losses is starting to slow down and it looks like it's getting close to the final result.
The model with 192 hidden layers is close to convergence and we will next try a larger KAN hidden layer, the same as the MLP. We release the best checkpoint of VisionKAN with 192 hidden dim.
We used DeiT as a baseline for Vision KAN development, thanks to Meta and MIT for the amazing work!
@misc{VisionKAN2024,
author = {Ziwen Chen and Gundavarapu and WU DI},
title = {Vision-KAN: Exploring the Possibility of KAN Replacing MLP in Vision Transformer},
year = {2024},
howpublished = {\url{https://github.com/chenziwenhaoshuai/Vision-KAN.git}},
}
