InKAN

Fast B-spline Kolmogorov-Arnold Network layers for PyTorch.
Drop-in nn.Linear replacement with learnable spline activations.

pip install inkan

click to copy

What is a KAN?

In a standard MLP, each connection has a fixed weight and a shared activation function. In a Kolmogorov-Arnold Network, each edge has its own learnable univariate function φ(x) — a B-spline curve that transforms the scalar flowing along that edge.

MLP: y = σ(W·x + b) KAN: y_j = ∑_i φ_ji(x_i)

Each φ_ji is a learned cubic B-spline with local support: changing one coefficient affects the function over only ~4 knot intervals. This makes KANs inherently more interpretable and editable than MLPs.

Why InKAN?

2.8–5.4x Faster

Stable piecewise polynomial evaluation instead of recursive Cox-de Boor. Local 4-basis lookup instead of dense K-basis computation.

Drop-in Replacement

KANLayer(in, out) replaces nn.Linear(in, out). Same interface, same optimizer, same training loop.

Numerically Stable

Guaranteed non-negative basis output. Float32 accuracy ~1.2e-7. No catastrophic cancellation from alternating sums.

1D and 2D Splines

Univariate spline layers plus tensor-product B-spline surfaces for joint function approximation.

Quick Start

import torch
from inkan import KANLayer, KANNetwork

# Single layer — drop-in for nn.Linear
layer = KANLayer(8, 4)
x = torch.randn(32, 8)
y = layer(x)  # [32, 4]

# Multi-layer network
net = KANNetwork([784, 64, 10])
x = torch.randn(32, 784)
y = net(x)  # [32, 10]

# Standard PyTorch training
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
loss = y.sum()
loss.backward()
optimizer.step()

Performance

Method	Grid Size 8	Grid Size 20	Grid Size 50
Cox-de Boor (recursive)	1.0x	1.0x	1.0x
InKAN (piecewise + local)	2.8x faster	3.9x faster	5.4x faster

Used In

Attention Manifolds — Learned B-spline coordinate curves for AI safety and mechanistic interpretability in transformer attention. Per-coordinate spline curves on Q/K projections, trained on Llama-3.2-1B and 3B.

📦 PyPI 💻 GitHub 📄 Paper 📖 Docs