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 first built-in family uses triangular piecewise-linear hats. Users can add another family by constructing KANEdgeFamily: provide a basis dimension and a TorchLean model that maps Vec inDim to Vec (inDim * basisDim).

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.KANEdgeFamily : Type 2

Backend-compatible KAN edge family.

An edge family turns each scalar input coordinate into basisDim 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, compiled, CPU, and CUDA training paths supported by the underlying operations.

• name : String

  Short label shown in model summaries and training metadata.

• basisDim : ℕ

  Number of basis features produced per scalar input coordinate.

• basis (inDim : ℕ) : Sequential (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim (inDim * self.basisDim) Spec.Shape.scalar)

  Basis expansion for an unbatched vector of length inDim.

structure TorchLean.nn.models.KANPiecewiseLinear : 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.

• gridSize : ℕ

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

• inputScale : ℕ

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

def TorchLean.nn.models.instReprKANPiecewiseLinear.repr : KANPiecewiseLinear → ℕ → Std.Format

@[implicit_reducible]

instance TorchLean.nn.models.instReprKANPiecewiseLinear : Repr KANPiecewiseLinear

def TorchLean.nn.models.KANPiecewiseLinear.basisLayer (cfg : KANPiecewiseLinear) (inDim : ℕ) : Sequential (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim (inDim * cfg.gridSize) Spec.Shape.scalar)

Expand x : Vec inDim to all triangular basis features.

The output is flattened row-major from a (gridSize × inDim) 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.KANPiecewiseLinear.edgeFamily (cfg : KANPiecewiseLinear) : KANEdgeFamily

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

structure TorchLean.nn.models.KANConfig : Type 2

Configuration for a KAN over batched row vectors.

• batch : ℕ

  Leading minibatch dimension.

• inDim : ℕ

  Number of scalar input coordinates.

• hidden : List ℕ

  Hidden KAN widths. Each entry creates one KAN layer followed by tanh.

• outDim : ℕ

  Number of output coordinates/classes.

• edge : KANEdgeFamily

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

@[reducible, inline]

abbrev TorchLean.nn.models.kanInShape (cfg : KANConfig) : Spec.Shape

Input shape (batch × inDim) for a KAN config.

@[reducible, inline]

abbrev TorchLean.nn.models.kanOutShape (cfg : KANConfig) : Spec.Shape

Output shape (batch × outDim) for a KAN config.

def TorchLean.nn.models.kanLayer (inDim outDim : ℕ) (edge : KANEdgeFamily) : M (Sequential (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar))

One unbatched KAN layer.

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 outDim.

def TorchLean.nn.models.kanGo (edge : KANEdgeFamily) (inDim : ℕ) (hidden : List ℕ) (outDim : ℕ) : M (Sequential (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar))

Recursive unbatched KAN stack. Hidden layers use tanh; the final layer is linear in bases.

def TorchLean.nn.models.KAN (cfg : KANConfig) : M (Sequential (kanInShape cfg) (kanOutShape cfg))

Build a batched KAN model.

Task semantics are deliberately not baked into the model name: use Trainer.new with task := .regression, .classification, .crossEntropy, or .custom ... with the same KAN constructor.