IvanDrokin/torch-conv-kan

TorchConv‑KAN – Convolutional Kolmogorov‑Arnold Networks in PyTorch

What it is

• A research‑oriented PyTorch library that implements convolutional layers based on the Kolmogorov‑Arnold representation (KAN). Instead of a fixed weight matrix, each kernel stores a set of learnable univariate functions (splines, polynomials, wavelets, etc.).
• Provides a growing collection of variants – KAN, Fast‑KAN, KALN, KAGN, ChebyKAN, WavKAN, JacobiKAN, Bernstein‑KAN, ReLU‑KAN, and bottleneck versions – for 1‑D, 2‑D and 3‑D convolutions.
• Packs these layers into familiar CNN backbones (ResNet‑like, DenseNet‑like, VGG‑like, U‑Net/U2‑Net) so you can drop a Conv‑KAN block into an existing architecture.

Key features

• Layer zoo – ready‑to‑use KANConv*, FastKANConv*, KALNConv*, KAGNConv*, WavKANConv*, JacobiKANConv*, BernsteinKANConv*, ReLUKANConv* and bottleneck equivalents.
• Model zoo – ResKANet, DenseKANet, VGG‑KAN, UKANet/U2KANet, and newer ConvNeXt‑style blocks built on these layers.
• Pre‑trained checkpoints – Imagenet‑1k weights for several VGG‑KAN variants (e.g., VGG‑KAGN‑11‑BN achieving 68.5 % top‑1 with only 7.25 M parameters) and recent ConvNeXt‑KAGN models.
• Training utilities – scripts for MNIST, CIFAR‑10/100, Tiny‑Imagenet, and Imagenet‑1k; integration with 🤗 Accelerate, Hydra configs, and Weights & Biases logging.
• Quantisation & PEFT – support for post‑training quantisation and parameter‑efficient fine‑tuning of KAN‑based models.
• Hyper‑parameter search – Ray‑Tune wrappers for automated tuning; optional LBFGS optimizer.

Typical use‑case

import torch, torch.nn as nn
from kan_convs import KANConv2DLayer

class SimpleConvKAN(nn.Module):
    def __init__(self, layer_sizes, num_classes=10, in_ch=1, spline_order=3):
        super().__init__()
        self.features = nn.Sequential(
            KANConv2DLayer(in_ch, layer_sizes[0], spline_order, kernel_size=3, padding=1),
            KANConv2DLayer(layer_sizes[0], layer_sizes[1], spline_order, kernel_size=3, stride=2, padding=1),
            KANConv2DLayer(layer_sizes[1], layer_sizes[2], spline_order, kernel_size=3, stride=2, padding=1),
            KANConv2DLayer(layer_sizes[2], layer_sizes[3], spline_order, kernel_size=3, padding=1),
            nn.AdaptiveAvgPool2d(1),
        )
        self.head = nn.Linear(layer_sizes[3], num_classes)
        self.drop = nn.Dropout(0.25)
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.drop(x)
        return self.head(x)

Run the provided training script (e.g., python mnist_conv.py) to train/evaluate on MNIST, CIFAR‑10/100, or switch to the Accelerate‑based scripts for larger datasets.

Installation & quick start

git clone https://github.com/IvanDrokin/torch-conv-kan.git
cd torch-conv-kan
pip install -r requirements.txt   # PyTorch + CUDA, accelerate, hydra, wandb, ray[tune]
# optional: wandb login   # for experiment tracking
accelerate launch cifar.py        # train a ResKANet on CIFAR‑10 with default config

Project status

• Actively developed (updates listed daily from May 2024 to July 2024, plus a July 2026 addition).
• Core research paper released: Kolmogorov‑Arnold Convolutions: Design Principles and Empirical Studies (arXiv 2407.01092).
• Benchmarks are still preliminary; the author notes that performance on CIFAR‑10/100 lags behind classic CNNs and that stability issues exist for some variants (e.g., ChebyKAN).

Who might find it useful

• Researchers exploring learnable activation functions inside convolution kernels.
• Practitioners wanting to experiment with KAN‑style layers in vision models without building everything from scratch.
• Anyone interested in quantisation or PEFT of non‑standard convolutional architectures.

Citations If you use the code or the baseline results, cite the accompanying arXiv paper (BibTeX provided in the repo).

---

All information above is taken directly from the repository’s README; no external assumptions have been added.