kanditioned 1.0.3

Fast, Conditioned KAN

KANditioned: Fast, Conditioned Training of Kolmogorov-Arnold Networks (KANs) via Lookup Interpolation and Discrete Cosine Transform

Training is accelerated by orders of magnitude through exploiting the structure of the linear (C⁰) B-spline (see Fig. 1) with uniformly spaced control points. Because the intervals are uniform, evaluating spline(x) reduces to a constant-time index calculation, followed by looking up the two relevant control points and linearly interpolating between them. This contrasts with the typical summation over basis functions typically seen in splines, reducing the amount of computation required and enabling effective sublinear scaling across the control points dimension.

Install

pip install kanditioned

Usage

[!IMPORTANT]
It is highly recommended to use this layer with torch.compile, which will provide very significant speedups (Triton kernel is coming sometimes later, but I found torch.compile to provide very satisfactory performance), in addition to a normalization layer before each KANLayer.

from kanditioned.kan_layer import KANLayer

layer = KANLayer(in_features=3, out_features=3, init="random_normal", num_control_points=8)
layer.visualize_all_mappings(save_path="kan_mappings.png")

Arguments

• in_features (int)
  Size of each input sample.
• out_features (int)
  Size of each output sample.
• init (str) – initialization method:
  • "random_normal": Slopes drawn from a normal distribution, then normalized so each “neuron” has unit weight norm.
  • "identity": Identity mapping (requires in_features == out_features). Output initially equals input.
  • "zero": All splines initialized to zero.
• num_control_points (int, default = 32)
  Number of uniformly spaced control points per input feature.
• spline_width (float, default = 4.0)
  Domain of the spline: [-spline_width / 2, spline_width / 2]. Beyond that, the spline will linearly extrapolate.
• impl (str, default = embedding_bag) Note: F.embedding_bag implementation appears to be much faster when used for inference with torch.compile enabled or when used for inference/training without torch.compile. However, F.embedding appears to be somewhat faster than F.embedding_bag when used for training with torch.compile enabled. Experiment with either implementation as necessary to achieve peak performance.
  • "embedding_bag": Implement the layer using F.embedding_bag.
  • "embedding": Implement the layer using F.embedding.

Methods:

visualize_all_mappings(save_path=path[optional]) - this will plot out the shape of each spline and its corresponding input and output feature

Figure

Figure 1. Linear B-spline example (each triangle-like shape is a basis):

Roadmap

• Use F.embedding_bag
• Update doc for variant and other new parameters introduced
• Update package with cleaned up, efficient Discrete Cosine Transform (with rank-2 correction) and parallel scan (prefix sum) parameterizations.
  • Both provide isotropic O(1) condition scaling for the discrete second difference penalty, as opposed to O(N^4) conditioning for the naive B-spline parameterization. This only matters if you care about regularization.
  • May add linearDCT variant first. Although it's O(N^2), it's more parallelized and optimized on GPU for small N since it's essentially a matmul with weight being a DCT matrix
• Proper baselines against MLP and various other KAN implementations on backward and forward passes
• Add sorting on indices and unsorting as an option (potentially radix sort, which is common optimization on embedding) to improve computational time through global memory "coalesced" access
• Add in feature-major input variant
• May change to either unfold or as_strided (slight performance improvement)
• Run benchmarks and further optimize memory locality
  • Feature-major input variant versus batch-major input variant
  • Interleaved indices [l1, u1, l2, u2, ...] versus stacked indices [l1, l2, ..., u1, u2, ...]
• Add optimized Triton kernel
• Update visualize_all_mappings method to something like .plot with option for plotting everything
• Research adding Legendre polynomials parameterization
  • Preliminary: does not seem to offer much benefits or have isotropic penalty conditioning
• Polish writing

Open To Collaborators. Contributions Are Welcomed!

LICENSE

This project is licensed under the Apache License 2.0.