Kolmogorov-Arnold Networks

KAN layers replace each scalar edge by a small trainable one-dimensional function. TorchLean keeps that structure visible: an edge family first expands every scalar input into basis features, and the KAN layer learns one coefficient per (output, input, basis) edge.

The built-in family uses triangular piecewise-linear hats. Another family consists of a basis size and a TorchLean model from [inputWidth] to [inputWidth * basisSize].

References:

• Z. Liu et al., "KAN: Kolmogorov-Arnold Networks", arXiv:2404.19756.
• C. de Boor, "A Practical Guide to Splines", Springer, 1978/2001.

structure TorchLean.nn.models.KAN.EdgeFamily : Type 1

Backend-compatible KAN edge family.

An edge family turns each scalar input coordinate into basisSize features. A KAN layer then applies a learned linear map to all expanded features. The basis is a TorchLean model fragment, not an arbitrary Lean callback, so the resulting KAN can run in eager, typed graph, CPU, and CUDA training paths supported by the underlying operations.

• name : String

  Short label shown in model summaries and training metadata.

• basisSize : ℕ

  Number of basis features produced per scalar input coordinate. Must be positive.

• basis (inputWidth : ℕ) : Sequential [inputWidth] [inputWidth * self.basisSize]

  Basis model for an unbatched feature vector.

structure TorchLean.nn.models.KAN.PiecewiseLinear.Config : Type

Configuration for triangular piecewise-linear KAN edge bases.

The basis functions are hats centered at the integer knots $0,\ldots,\mathrm{gridSize}-1$. The input is multiplied by inputScale before the hats are evaluated. For normalized data in $[0,1]$, setting $\mathrm{inputScale}=\mathrm{gridSize}-1$ spreads the grid across the full interval. The built-in family uses a nonnegative integer scale so the same definition works for every TorchLean scalar backend through Context's natural-number conversion.

• gridSize : ℕ

  Number of knots, hence the number of basis functions per scalar coordinate. Must be positive.

• inputScale : ℕ

  Nonnegative integer scale applied before basis evaluation; use $\mathrm{gridSize}-1$ for normalized $[0,1]$ inputs.

def TorchLean.nn.models.KAN.PiecewiseLinear.instReprConfig.repr : Config → ℕ → Std.Format

@[instance_reducible]

instance TorchLean.nn.models.KAN.PiecewiseLinear.instReprConfig : Repr Config

def TorchLean.nn.models.KAN.PiecewiseLinear.layer (config : Config) (inputWidth : ℕ) : Sequential [inputWidth] [inputWidth * config.gridSize]

Expand a tensor of shape [inputWidth] to all triangular basis features.

The output is flattened row-major from a (gridSize × inputWidth) table: $[\operatorname{basis}_0(x_0),\ldots,\operatorname{basis}_0(x_n), \operatorname{basis}_1(x_0),\ldots]$.

Each basis value is $\operatorname{ReLU}(1-|\mathrm{inputScale}\,x_i-k|)$, expressed directly in the ordinary TorchLean op language rather than through an opaque spline evaluator.

def TorchLean.nn.models.KAN.PiecewiseLinear.edgeFamily (config : Config) : EdgeFamily

Turn piecewise-linear triangular bases into a general KAN edge family.

structure TorchLean.nn.models.KAN.Config : Type 1

Architecture of a Kolmogorov-Arnold network over feature vectors.

• inputWidth : ℕ

  Number of scalar input coordinates. Must be positive.

• hiddenWidths : List ℕ

  Hidden KAN widths. Each must be positive and creates one KAN layer followed by tanh.

• outputWidth : ℕ

  Number of output coordinates/classes. Must be positive.

• edge : EdgeFamily

  Edge basis family. The default is a compact triangular piecewise-linear basis.

def TorchLean.nn.models.KAN.Config.validate (config : Config) : Except String Unit

Validate every architecture width and the selected edge family before construction.

@[reducible, inline]

abbrev TorchLean.nn.models.KAN.Config.inputShape (config : Config) (batchShape : Shape := []) : Shape

Typed input boundary with arbitrary batch axes.

@[reducible, inline]

abbrev TorchLean.nn.models.KAN.Config.outputShape (config : Config) (batchShape : Shape := []) : Shape

Typed output boundary with the same batch axes as the input.

def TorchLean.nn.models.KAN.layer (inputWidth outputWidth : ℕ) (edge : EdgeFamily) : Builder (Sequential [inputWidth] [outputWidth])

One KAN layer over a feature vector.

The layer first applies the selected edge basis to every input coordinate, then learns coefficients with an ordinary linear map from the expanded features to outputWidth.

def TorchLean.nn.models.KAN.Internal.stack (edge : EdgeFamily) (inputWidth : ℕ) (hiddenWidths : List ℕ) (outputWidth : ℕ) : Builder (Sequential [inputWidth] [outputWidth])

Recursive KAN stack over one feature vector. Hidden layers use tanh.

def TorchLean.nn.models.kan (config : KAN.Config) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build a KAN over any batchShape.

Task semantics are deliberately not baked into the model name: use Trainer.new with objective := .meanSquaredError, .oneHotCrossEntropy axis, or .custom ... with the same KAN constructor.