Kolmogorov-Arnold Networks without the tensor shape explosion

What it does

Kolmogorov-Arnold Networks (KAN) replace fixed activation functions with learnable B-splines. The original PyTorch implementation materializes a giant (batch_size, out_features, in_features) tensor to apply those splines, which is exactly as memory-hungry as it sounds. This repo reformulates the math: it applies the fixed B-spline basis functions to the input once, then combines the results linearly. That turns the layer into a straightforward matrix multiplication, letting PyTorch’s autograd handle the rest without any custom expansion.

The interesting bit

The speedup comes at a cost to the paper’s sparsification recipe. The authors’ L1 regularization is defined per-input-sample and requires operating on that same bloated tensor, so the author swaps it for a conventional weight-level L1 penalty. It is a pragmatic compromise, though the README admits it is unclear whether this preserves the interpretability KAN is supposed to be famous for.

Key highlights

• Replaces per-activation tensor expansion with basis-function activation followed by linear combination
• Compatible with standard PyTorch autograd; no custom backward needed
• Includes an optional enable_standalone_scale_spline flag to match the original’s learnable per-activation scaling
• Switches base_weight and spline_scaler initialization to kaiming_uniform_, which reportedly jumps MNIST accuracy from roughly 20% to 97%
• Retains the original’s B-spline learnable activation functions

Caveats

• The sample-wise L1 sparsification used in the original paper is incompatible with the efficient formulation and has been replaced by weight-level L1 regularization; the README notes this may affect interpretability and needs more experiments
• The enable_standalone_scale_spline option defaults to True for fidelity to the original paper, but disabling it trades accuracy for further speed, and the README says it “needs more experiments”
• The switch to kaiming_uniform_ initialization was driven by a single MNIST observation; the author is unsure whether it generalizes

Verdict

Worth a look if you are curious about KANs but were put off by the original implementation’s memory footprint. Skip it if you need the exact sparsification and interpretability guarantees described in the paper, because this is explicitly a pragmatic, efficiency-first rewrite.