kan-gpt-2's Introduction

Training small GPT-2 style models using KANs instead of MLPs in JAX

This repository compares transformers using multilayer perceptron (MLP) and Kolmogorov-Arnold networks (KAN) layers.

Key points:

Uses Kolmogorov-Arnold Networks but with Chebyshev polynomials as the basis (inspired by this repo).
The tanh function is used to keep the activation values within [-1, 1] rather than using grids that update during training.
Both models are trained on 134M tokens of TinyStories.
They both use standard GPT-2 architecture (other than the KAN part).
The MLP version has 3.3M non-embedding weights and the KAN model has 2.5M non-embedding weights (~25% fewer).

Results:

They both achieve a final loss of ~2.46 (despite the KAN model having 25% fewer parameters!).

Hyperparameters:

d_model: 128
d_mlp: 768 (when applicable)
n_heads: 8
n_layers: 16
learning_rate: 1e-5
batch_size: 16
weight_decay: 0.001
optimizer: adamw
seq_len: 64

Hardware: Single 1080ti GPU

Wandb: link.

kan-gpt-2's People

Contributors

Stargazers

Watchers

kan-gpt-2's Issues

Any suggestions for KAN-ViT?

Dear @CG80499 ,
Thank you for your contribution.

Using your implementation of ChebyKAN layer, I am currently training Vision Transformer tiny model on ImageNet1K. But it seems like it is underperforming the original implementation.

Any suggestions? Why did you set polynomial degree as 8 for GPT?

Here are some minor details for my implementation on ViT.

I used DeiT replication written in JAX/FLAX.
Learning rate, batch size, weight decay, random seed etcs are set to be the same.
I replaced feedforward MLP that are subsequent to the attention block only, while remaining attention layers to use nn.Dense which is similar to your KAN-GPT-2 implementation.

Thank you.