TorchKAN: Kolmogorov–Arnold Networks in PyTorch

An independent PyTorch implementation of Kolmogorov–Arnold Networks (KANs) and several polynomial variants — spline, Legendre (KAL-Net), Chebyshev (KAC-Net), and a convolutional variant (KANvolver) — with classification and regression experiments, CUDA execution, post-training quantisation, and Integrated-Gradients interpretability.

Contents

1. KANvolver — monomial-basis convolutional model for MNIST classification
2. KAL-Net — Legendre-polynomial KAN
3. KAC-Net — Chebyshev-polynomial KAN
4. nKAN — generalised KAN for curve fitting / regression (nKAN.py)

This repository trains, validates, and quantises KAN variants in PyTorch with CUDA acceleration, evaluated on MNIST.

Status

Stable and maintained on a best-effort basis; the experiments and quantisation results below reflect the code in this repository. The base KAN in torchkan.py reaches ~97% accuracy on MNIST (8 epochs, RTX 4090).

Current Understanding: While there is considerable hype around KANs, it's important to recognize that learning weights for activation functions (MLPs) and the activation functions themselves are established ideas. The extent of interpretability, scalability, quantizability, or efficiency they offer remains unclear. However, quantizability does not seem to be an issue, as the quantized evaluation on the base model leads to only ~0.6% drop in test performance.

Note: KANs are an active research area; contributions, questions, and critiques are welcome.

---

Introduction to the KANvolver Model

The KANvolver model is a specialized neural network designed for classifying images from the MNIST dataset. It reaches ~99.5% accuracy on MNIST, combining convolutional feature extraction with polynomial feature expansions.

I am conducting large-scale analysis to investigate how KANs can be made more interpretable.

Thanks to @cometscome for writing this version in Julia: https://github.com/cometscome/FluxKAN.jl

Model Architecture

Convolutional Feature Extraction: The model begins with two convolutional layers, each paired with ReLU activation and max-pooling. The first layer employs 16 filters of size 3x3, while the second increases the feature maps to 32 channels.

Polynomial Feature Transformation: After feature extraction, the model applies polynomial transformations up to the n-th order to the flattened convolutional outputs, enhancing its ability to discern non-linear relationships.

How Monomials Work: In this model, monomials are polynomial powers of the input features. By computing monomials up to a specified order, the model captures non-linear interactions between the features, potentially leading to richer and more informative representations for downstream tasks.

For a given input image, the monomials of its flattened pixel values are computed and then used to adjust the output of linear layers before activation. This approach introduces an additional dimension of feature interaction, allowing the network to learn more complex patterns in the data.

Forward Propagation Process

1. Input Reshaping: Images are reshaped from vectors of 784 elements to 1x28x28 tensors for CNN processing.
2. Feature Extraction: Spatial features are extracted and pooled through convolutional layers.
3. Polynomial Expansion: Features undergo polynomial expansion to capture higher-order interactions.
4. Linear Processing: The expanded features are processed by linear layers with normalization and activation.
5. Output Generation: The network produces logits for each digit class in MNIST.

Performance and Conclusion

The KANvolver reaches ~99.5% accuracy on MNIST using CNN features with polynomial expansions. Results:

Integrated Gradients for Primary Explainability

Note that KANvolver uses polynomials that are distinct from the original KANs[1].

---

KANs seem to handle noise better compared to MLPs for functional approximation. This requires further investigation.

To reproduce the results, use the nKAN.py script.

---

Introducing KAL_Net

The KAL_Net represents the Kolmogorov Arnold Legendre Network (KAL-Net), a GAM architecture using Legendre polynomials to surpass traditional polynomial approximations like splines in KANs.

Key Features

• Polynomial Order: Utilizes Legendre polynomials up to a specific order for each input normalization, capturing nonlinear relationships more efficiently than simpler polynomial approximations.
• Efficient Computations: By leveraging functools.lru_cache, the network avoids redundant computations, enhancing the forward pass's speed.
• Activation Function: Employs the SiLU (Sigmoid Linear Unit) for improved performance in deeper networks due to its non-monotonic nature.
• Layer Normalization: Stabilizes each layer's output using layer normalization, enhancing training stability and convergence speed.

Design and Initialization

1. Weight Initialization: Weights are initialized using the Kaiming uniform distribution, optimized for linear nonlinearity, ensuring a robust start for training.
2. Dynamic Weight and Normalization Management: Manages weights for base transformations and polynomial expansions dynamically, scaling with input features and polynomial order.

Motivation vs. Splines

• Flexibility in High-Dimensional Spaces: Legendre polynomials offer a more systematic approach to capturing interactions in high-dimensional data compared to splines, which often require manual knot placement and struggle with dimensionality issues.
• Analytical Efficiency: The caching and recurrence relations in Legendre polynomial computations minimize the computational overhead associated with spline evaluations, especially in high dimensions.
• Generalization: The orthogonal properties of Legendre polynomials typically lead to better generalization in machine learning model fitting, avoiding common overfitting issues with higher-degree splines.

Performance Metrics

• Accuracy: KAL_Net reaches ~97.8% accuracy on MNIST.
• Efficiency: Legendre-polynomial computations are cached (functools.lru_cache) to reduce redundant work in the forward pass.

Prerequisites

Ensure the following are installed on your system:

• Python (version 3.9 or higher)
• CUDA Toolkit (compatible with your PyTorch installation's CUDA version)
• cuDNN (compatible with your installed CUDA Toolkit)

Installation

Tested on MacOS and Linux.

1. Clone the Repository

Clone the torchkan repository and set up the project environment:

git clone https://github.com/1ssb/torchkan.git
cd torchkan
pip install -r requirements.txt

2. Configure CUDA Environment Variables if Unset:

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

3. Configure Weights & Biases (wandb)

To monitor experiments and model performance with wandb:

1. Set Up wandb Account:

• Sign up or log in at Weights & Biases.
• Locate your API key in your account settings.

2. Initialize wandb in Your Project:

Before running the training script, initialize wandb:

wandb login

Enter your API key when prompted to link your script executions to your wandb account.

3. Adjust the Entity Name in mnist.py to Your Username (default is 1ssb)

Usage

python mnist.py

This script trains the model, validates it, quantizes it, and logs performance metrics using wandb.

Contact

For inquiries or support, please contact: Subhransu.Bhattacharjee@anu.edu.au

Cite this Project

If this project is used in your research or referenced for baseline results, please use the following BibTeX entry.

@misc{torchkan,
  author = {Subhransu S. Bhattacharjee},
  title = {{TorchKAN}: Simplified {KAN} Model with Variations},
  year = {2024},
  howpublished = {\url{https://github.com/1ssb/torchkan/}}
}

Contributions

Contributions are welcome. Please raise issues as needed. Maintained solely by @1ssb.

References

• [0] Ziming Liu et al., "KAN: Kolmogorov-Arnold Networks", 2024, arXiv. https://arxiv.org/abs/2404.19756
• [1] https://github.com/KindXiaoming/pykan
• [2] https://github.com/Blealtan/efficient-kan

---