mlx-KANs

mlx-KANs is a high-performance suite of Kolmogorov-Arnold Networks (KAN) built on top of Apple MLX, specifically optimized for Apple Silicon (M1 / M2 / M3 / M4 / Pro / Max / Ultra) with custom Metal Shading Language (MSL) kernels, mixed-precision FP16/BF16, zero-allocation GEMM projections, and unified memory acceleration.

🇷🇺 Русскоязычная версия документации доступна в README_RU.md.

---

Key Highlights

• 10 Modern KAN Architectures in One Package:
  • FastKAN (Gaussian RBF basis) — ~3x faster than B-splines with $C^\infty$ smoothness.
  • ReLUKAN (Piecewise-linear tent basis) — zero transcendental operations, >1.17M samples/s.
  • LowRankKAN (LoRA / Bottleneck factorization) — up to 10x parameter compression for deep & wide layers.
  • WavKAN (Continuous Wavelets: Mexican Hat, Morlet, DOG) — multiresolution frequency-time localization, lowest MSE.
  • ChebyKAN (Chebyshev Polynomials 1st kind) — minimax polynomial optimality on $[-1, 1]$ without grid updates.
  • FourierKAN (Harmonic Fourier series) — trigonometric basis for periodic signals, audio, and PINNs.
  • RationalKAN (Padé-Chebyshev Rational Functions) — adaptive pole and singularity modeling ($P(x)/Q(x)$) for boundary layers and stiff PDEs.
  • JacobiKAN (Jacobi / Legendre / Gegenbauer orthogonal polynomials) — parametrized by $(\alpha, \beta)$.
  • MultKAN (KAN 2.0 with Multiplication Nodes) — exact analytical product terms $u \cdot v$ for physical conservation laws.
  • KAN (Cubic B-splines / Efficient-KAN) — classical Cox-de Boor reformulation with adaptive grid updates.
• Hybrid Metal + MLX Engine (HybridKAN):
  • Dynamically routes execution between MLX JIT and fused Direct Metal Shading Language (MSL) shaders. Small batches ($B < 1024$) leverage low-overhead MLX graph dispatch; large batches ($B \ge 1024$) switch to fused Metal kernels with 0 intermediate basis allocations, reaching up to 10.58M samples/sec (3.36x speedup on Apple Silicon GPU).
• Hardware-Fused Metal Kernels (mx.fast.metal_kernel):
  • Handcrafted Metal Shading Language (MSL) kernels executing RBF, tent-ReLU, and Chebyshev recurrences directly in GPU thread registers without VRAM memory traffic.
• Unified Memory Stream Orchestration:
  • Zero-copy CPU/GPU streaming via macOS unified memory for least-squares pseudo-inverses (pinv).
• JIT Graph Compilation:
  • Integrated build_train_step compiling forward pass, autograd backpropagation, and optimizer updates into a single monolithic Metal execution graph.
• Mixed Precision:
  • Native to_fp16() and to_bf16() casting doubling GPU memory bandwidth and throughput.
• Native MLX Optimizer Suite:
  • Direct re-exports of AdamW, Muon (Newton-Schulz 5th order matrix orthogonalization), Lion (sign momentum), RMSprop, Adam, and SGD directly from mlx_kans (from mlx_kans import Muon, AdamW).

---

Installation

# Clone the repository
git clone https://github.com/Shuril/mlx-kans.git
cd mlx-kans

# Install in editable mode
pip install -e .

# Or using uv (recommended for ultra-fast setup):
uv pip install -e .

---

Quick Start

import mlx.core as mx
import mlx.optimizers as optim
import mlx_kans as kans

# 1. Instantiate any KAN model
model = kans.FastKAN(layers_hidden=[2, 16, 1], num_grids=8)

# 2. Setup optimizer directly from mlx_kans (Muon, AdamW, Lion, RMSprop, SGD)
optimizer = kans.Muon(learning_rate=0.02)

def loss_fn(m, x, y):
    pred = m(x)
    return mx.mean((pred - y) ** 2)

# 3. Compile training step directly into Metal GPU graph
train_step = kans.build_train_step(model, optimizer, loss_fn)

# 4. Generate data & train
x_train = mx.random.uniform(-1.0, 1.0, (1000, 2))
y_train = mx.sin(4.0 * mx.pi * x_train[:, 0:1]) + (x_train[:, 1:2] ** 2)

for epoch in range(100):
    loss = train_step(x_train, y_train)
    mx.eval(model.parameters(), optimizer.state)
    if epoch % 20 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss.item():.5f}")

---

Performance Benchmarks on Apple Silicon Metal GPU

Benchmarked on Apple Silicon GPU (Device(gpu, 0)), architecture [64, 64, 64], batch size 1024:

========================================================================================
Model              | Trainable Params | Forward (ms)   | Throughput (samples/s) | Train Step (ms)
----------------------------------------------------------------------------------------
MLP (Baseline)     | 8,320            |      0.339 ms  |        3,016,866       |       0.691 ms
LowRankKAN         | 15,360           |      0.828 ms  |        1,236,498       |       1.951 ms
ReLUKAN (Tent)     | 57,344           |      0.871 ms  |        1,175,727       |       2.072 ms
WavKAN (Wavelet)   | 58,880           |      0.983 ms  |        1,041,365       |       3.028 ms
JacobiKAN          | 40,960           |      0.992 ms  |        1,032,210       |       1.762 ms
MultKAN (2.0)      | 86,016           |      1.016 ms  |        1,008,141       |       2.290 ms
FastKAN (RBF)      | 57,344           |      1.142 ms  |          896,675       |       2.087 ms
FourierKAN         | 65,536           |      1.397 ms  |          733,221       |       2.365 ms
ChebyKAN           | 40,960           |      2.186 ms  |          468,365       |       1.799 ms
B-Spline KAN       | 81,920           |      2.392 ms  |          428,156       |       5.193 ms
========================================================================================

3-Way GPU Benchmark: MLX vs metal-KANs (v0.3.1) vs slang-KANs (v0.2.0)

Tested across all 10 architectures on Apple Silicon GPU, Layer 64 -> 64:

Architecture	Batch	MLX (ms)	metal-KANs (Pure Metal AMX)	slang-KANs (Shared-Memory GEMM)	Winner (Speedup)
ChebyKAN	128	0.387 ms	0.295 ms	0.226 ms	slang-KANs (1.30x)
	1024	0.526 ms	0.372 ms	0.324 ms	slang-KANs (1.15x)
	4096	1.380 ms	0.484 ms	0.885 ms	metal-KANs (1.83x)
BSplineKAN	128	0.354 ms	0.258 ms	0.136 ms	slang-KANs (1.90x)
	1024	0.835 ms	0.404 ms	0.541 ms	metal-KANs (1.34x)
	4096	2.190 ms	0.826 ms	1.400 ms	metal-KANs (1.70x)
FastKAN	128	0.270 ms	0.254 ms	0.692 ms	metal-KANs (1.06x)
	1024	0.398 ms	0.352 ms	1.439 ms	metal-KANs (1.13x)
	4096	1.475 ms	1.006 ms	4.010 ms	metal-KANs (1.47x)
WavKAN	128	0.354 ms	0.367 ms	0.869 ms	MLX (1.04x)
	1024	1.845 ms	0.541 ms	1.941 ms	metal-KANs (3.41x)
	4096	1.844 ms	1.189 ms	4.132 ms	metal-KANs (1.55x)
ReLUKAN	128	0.311 ms	0.364 ms	0.852 ms	MLX (1.17x)
	1024	0.542 ms	0.511 ms	1.888 ms	metal-KANs (1.06x)
	4096	1.254 ms	1.062 ms	4.134 ms	metal-KANs (1.18x)
FourierKAN	128	0.351 ms	0.384 ms	0.883 ms	MLX (1.09x)
	1024	0.975 ms	0.605 ms	2.098 ms	metal-KANs (1.61x)
	4096	2.385 ms	0.868 ms	2.868 ms	metal-KANs (2.75x)
JacobiKAN	128	0.405 ms	0.294 ms	0.144 ms	slang-KANs (2.04x)
	1024	0.632 ms	0.415 ms	0.325 ms	slang-KANs (1.28x)
	4096	1.413 ms	0.515 ms	0.722 ms	metal-KANs (1.40x)
RationalKAN	128	0.691 ms	0.385 ms	0.166 ms	slang-KANs (2.32x)
	1024	4.354 ms	0.822 ms	0.648 ms	slang-KANs (1.27x)
	4096	18.225 ms	2.300 ms	2.468 ms	metal-KANs (1.07x)
MultKAN	128	0.378 ms	0.278 ms	0.783 ms	metal-KANs (1.36x)
	1024	0.494 ms	0.415 ms	3.738 ms	metal-KANs (1.19x)
	4096	2.179 ms	1.526 ms	7.734 ms	metal-KANs (1.43x)
LowRankKAN	128	0.514 ms	0.546 ms	0.811 ms	MLX (1.06x)
	1024	0.866 ms	0.612 ms	8.050 ms	metal-KANs (1.41x)
	4096	1.722 ms	1.101 ms	27.890 ms	metal-KANs (1.56x)

Standalone repositories:

• metal-KANs (v0.3.1): Pure Metal Shading Language (MSL) with Apple AMX coprocessor matrix multiplication and SIMD basis kernels.
• slang-KANs (v0.2.0): Cross-platform Slang shaders with 16x16 shared-memory tiled GEMM and Vulkan/Metal/CUDA portability.

---

Hybrid Engine: MLX + Direct Metal Shading Language (MSL)

For production workloads and massive batch inference, mlx-KANs features a Hybrid Engine (HybridChebyKAN, HybridFastKAN, HybridReLUKAN, HybridKAN).

Why Hybrid?

Standard MLX graphs evaluate KAN layers in two distinct stages:

1. Basis Expansion: Generates basis representations (e.g. $[B, D_{\text{in}}, \text{degree}]$) and writes them to intermediate GPU memory buffers.
2. Linear Projection: Executes GEMM over materialized basis tensors.

While MLX JIT compiles this graph effectively, large batch sizes ($B \ge 1024$) cause basis expansion memory allocations to dominate total latency.

The Hybrid approach dynamically routes:

• Small Batches ($B < 1024$): Dispatched to MLX JIT graph, minimizing per-call overhead and maximizing Python responsiveness.
• Large Batches ($B \ge 1024$): Dispatched directly to Fused Metal Shading Language (MSL) compute shaders via macOS zero-copy Unified Memory pointers (newBufferWithBytesNoCopy).

The Direct Metal kernel computes the basis polynomial/RBF in hardware thread registers & threadgroup SRAM and fuses it directly into the output accumulator: 0 intermediate VRAM allocations.

M1 GPU Microbenchmark ($64 \to 64$, degree 4):

Batch Size ($B$)	MLX JIT Forward	Direct Fused Metal	Speedup	Intermediate Allocations
16	0.06 ms (16 µs graph)	0.08 ms	MLX faster	0 (registers)
256	0.28 ms	0.19 ms	1.47x	0 vs 131 KB
1,024	0.74 ms	0.32 ms	2.31x	0 vs 524 KB
4,096	2.12 ms	0.78 ms	2.72x	0 vs 2.1 MB
16,384	5.21 ms (3.14M samples/s)	1.54 ms (10.58M samples/s)	3.36x ⚡	0 MB vs 8.4 MB

Usage Example:

import mlx.core as mx
import mlx.optimizers as optim
import mlx_kans as kans

# Hybrid model automatically chooses between MLX JIT and Fused Metal
model = kans.HybridChebyKAN(
    in_features=64, 
    out_features=64, 
    degree=4, 
    adaptive_threshold=1024  # switch point
)

# Small batch uses MLX JIT
x_small = mx.random.normal((32, 64))
y_small = model(x_small)

# Large batch zero-copy routes to Direct Fused MSL GPU kernel (10.58M samples/sec)
x_large = mx.random.normal((4096, 64))
y_large = model(x_large)

# Full autograd compatibility
optimizer = optim.Adam(learning_rate=1e-3)
def loss_fn(m, x, y):
    return mx.mean((m(x) - y) ** 2)

step = kans.build_train_step(model, optimizer, loss_fn)
loss = step(x_small, mx.random.normal((32, 64)))

---

Native Metal INT8 & INT4 Quantization

mlx-KANs includes native hardware-accelerated INT8 and INT4 weight quantization on Apple Silicon GPU using Metal Performance Shaders (mx.quantize and mx.quantized_matmul).

How It Works:

In Kolmogorov-Arnold Networks, >95% of weights and FLOPs reside in the linear projection matrices of basis functions ($W_{\text{spline}}$ and $W_{\text{base}}$).

• Zero-Dequantization Overhead: mx.quantized_matmul performs matrix multiplication directly from packed uint32 vectors (4 INT8 values per word) with group scales and biases on Metal GPU without expanding weights back into FP32 VRAM.
• Auto-Padding: Non-standard layer dimensions (e.g. input size 2, 7, 13) are automatically zero-padded to group_size (32, 64, or 128) for seamless Metal execution.
• Standard MLX Compatibility: Use either kans.to_int8(model) or official mlx.nn.quantize(model, group_size=64, bits=8).

import mlx_kans as kans
import mlx.core as mx

# 1. Instantiate any KAN model
model = kans.FastKAN([128, 256, 128], num_grids=8)

# Check memory size before quantization
print("FP32 size:", kans.get_model_size(model)["summary"])
# -> 2.262 MB (2,371,584 bytes, 592,896 elements)

# 2. Quantize in-place to INT8 natively on Apple Silicon GPU
kans.to_int8(model, group_size=64)

# Check memory size after quantization
print("INT8 size:", kans.get_model_size(model)["summary"])
# -> 0.641 MB (671,744 bytes, 167,936 elements) -- 3.53x memory reduction!

# 3. Quantize to INT4 for extreme memory savings
kans.to_int4(model, group_size=64)
print("INT4 size:", kans.get_model_size(model)["summary"])
# -> 0.364 MB -- 6.23x memory reduction!

Hardware Support Matrix (INT8 vs INT4):

• INT8 / INT4 (affine): Supported natively on hardware across all Apple Silicon generations (M1, M2, M3, M4, M5+).
• Executes matrix multiplication directly on Apple Silicon GPU without dequantization to FP32.

---

Advanced Memory & Deployment Optimizations

1. Gradient / Activation Checkpointing (checkpoint_kan)

Saves ~57% peak VRAM during backpropagation of deep KANs on Apple Silicon Unified Memory. By discarding intermediate spline activations and recomputing them on-demand during the backward pass, it prevents memory bandwidth thrashing and out-of-core paging (yielding up to 7.4x faster training on M1):

import mlx_kans as kans

# Wrap any deep KAN with activation checkpointing
model = kans.FastKAN([128] * 9, num_grids=8)
checkpointed_model = kans.checkpoint_kan(model)

# Train normally with nn.value_and_grad
loss, grads = nn.value_and_grad(checkpointed_model, loss_fn)(checkpointed_model, x, y)

2. Structural Pruning & Node Compaction (prune)

KANs possess intrinsic node-level sparsity under L1 regularization. Unlike MLPs that require sparse indexing masks, KAN inactive neurons can be physically sliced from the weight matrices, shrinking layer dimensions and yielding up to 87% parameter reduction and >3x inference speedup:

# Compute importance and prune inactive neurons with < 5% of peak coupling
compact_model, stats = kans.prune(trained_model, threshold=0.05, min_active=2)

print("Original dims:", stats["orig_dims"])
print("Compacted dims:", stats["new_dims"])
print(f"Pruned {stats['pruned_neurons']} neurons ({stats['percent_neurons_pruned']:.1f}%)")

# compact_model is a real, smaller FastKAN with 0 sparse overhead!
y = compact_model(x_test)

3. Exact Symbolic Formula Extraction (to_symbolic)

KAN univariate edge curves can be matched against candidate analytical functions ($x^2, \sin, \cos, \exp$, polynomials) via least squares ($R^2 > 0.90$). Once converted to symbolic form, the network becomes a 100% white-box mathematical formula:

# Extract analytical symbolic formula
sym_kan = kans.to_symbolic(trained_model, r2_threshold=0.90)

# Print human-readable mathematical equation
print(sym_kan.formula())
# Output: y0 = 0.998*x0^2 + 1.001*sin(pi*x1)

# Print LaTeX representation
print(sym_kan.latex())
# Output: y_{0} = 0.998 x_{0}^2 + 1.001 \sin(\pi x_{1})

# Ultra-fast inference with 0 MB VRAM (pure math evaluation, up to 38x faster on CPU!)
y_pred = sym_kan(x_test)

4. Zero-Dependency ANSI C99 / C++ Header-Only Export (export_c)

Export trained and symbolically distilled KANs into clean, self-contained C99 headers (.h) with zero external dependencies (no MLX, no Python, no BLAS). Ideal for direct deployment into high-performance physics simulators (OpenFOAM, MODFLOW, SU2), real-time financial trading engines, or embedded edge microcontrollers (STM32, ESP32):

# Export directly to a C header file
sym_kan.export_c("kan_model.h", function_name="kan_evaluate")

# Or obtain the raw C code as a Python string
c_code = sym_kan.to_c_code(function_name="kan_evaluate")

• Throughput: $> 1,000,000,000$ points/sec ($10^7$ evaluations in $9.9\text{ ms}$ with clang -O3).
• Memory: $0\text{ MB}$ RAM, $0\text{ MB}$ VRAM. Pure mathematical registers.

Quantization Benchmark (Topology [128, 256, 128], Metal GPU):

Model	FP32 Mem	INT8 Mem	INT8 Compression	INT8 MAE	INT4 Mem	INT4 Compression
FastKAN (RBF)	2.26 MB	0.64 MB	3.51x	0.729	0.36 MB	6.23x
ReLUKAN (Tent)	2.26 MB	0.64 MB	3.51x	0.552	0.36 MB	6.23x
ChebyKAN	1.75 MB	0.49 MB	3.56x	1.459	0.27 MB	6.40x
WavKAN (Wavelet)	2.27 MB	0.66 MB	3.46x	0.706	0.38 MB	6.06x
FourierKAN	3.50 MB	0.98 MB	3.56x	1.883	0.55 MB	6.40x
JacobiKAN	1.75 MB	0.49 MB	3.56x	1.283	0.27 MB	6.40x
MultKAN (2.0)	3.39 MB	0.96 MB	3.52x	0.561	0.54 MB	6.28x
LowRankKAN	0.47 MB	0.16 MB	2.93x	0.115	0.09 MB	4.99x
B-Spline KAN	2.77 MB	0.72 MB	3.83x	0.118	0.41 MB	6.76x

---

Iso-Parameter Stress Test (Identical Parameter Budget $\pm 1%$)

To ensure strict scientific fairness, all models were calibrated to have the exact same parameter budget ($\pm 1%$) across 4 diverse stress scenarios:

• Task 1: High-Frequency Oscillations (2D): $f(x) = \sin(8\pi x_1)\cos(6\pi x_2) + 0.5\sin(16\pi x_1 x_2)$
• Task 2: Sharp Non-Smooth Transitions (2D): $f(x) = |x_1| - 2\max(0, x_2) + \text{sign}(x_1 x_2)|x_1 - x_2|^{0.7}$
• Task 3: Multiplicative Physical Law (4D): $f(x) = (x_1 x_2) e^{-x_3^2} + x_3 x_4^2$
• Task 4: High-Dimensional Nonlinear Target (8D): $f(x) = \exp(-\sum x^2 / 4) \sin(\pi \sum_{1}^4 x_i)$

Results (Test MSE, 150 epochs on Metal GPU):

Model	Budget	Task 1 (High-Freq, 2D)	Task 2 (Non-Smooth, 2D)	Task 3 (Physics, 4D)	Task 4 (Target, 8D)	Train Time / Epoch
MLP (Baseline)	$\approx$ const	0.42251	0.04385	0.00055 🥈	0.01713 🥇	0.72 – 0.93 ms
ReLUKAN (Tent)	$\approx$ const	0.18538 🏆	0.02744	0.00679	0.36655	0.79 – 0.88 ms ⚡
FastKAN (RBF)	$\approx$ const	0.30261 🥈	0.02387	0.02616	0.25177	0.84 – 0.93 ms ⚡
WavKAN (MexHat)	$\approx$ const	0.31261 🥉	0.02515	0.00222	0.23987	1.05 – 1.38 ms
LowRankKAN	$\approx$ const	0.35542	0.01555 🏆	0.00174 🥉	0.02772 🥈	1.05 – 1.49 ms
MultKAN (2.0)	$\approx$ const	0.37618	0.01930 🥈	0.01482	0.18123	0.90 – 0.97 ms
B-Spline KAN	$\approx$ const	0.43318	0.02300 🥉	0.00017 🏆	0.08623 🥉	2.07 – 2.91 ms
JacobiKAN	$\approx$ const	0.35679	0.03450	0.02681	0.24693	0.90 – 0.98 ms
ChebyKAN	$\approx$ const	0.35915	0.05112	0.04086	0.25154	0.84 – 1.02 ms
FourierKAN	$\approx$ const	0.92909	1.47926	0.25044	0.38723	0.89 – 0.90 ms

Key Takeaways:

1. High-Frequency Details: ReLUKAN outperforms standard MLP by over 2.2x (MSE 0.185 vs 0.422) by eliminating spectral bias.
2. Non-Smooth Boundaries: LowRankKAN and MultKAN outperform MLP by 2.3x–2.8x because univariate edge functions isolate kinks without global ringing artifacts.
3. High-Dimensional Scaling (8D): LowRankKAN achieves near-parity with MLP (MSE 0.027 vs 0.017) thanks to rank-4 factorization while standard full-grid KANs degrade under tight parameter limits.

---

Documentation

Detailed guides and references are available in the docs/ directory:

• API Reference: Full documentation for all classes, methods, and constructor arguments.
• Architecture & Mathematics: Mathematical formulations for each basis family and the Kolmogorov-Arnold representation theorem.
• Metal GPU Kernels Guide: Custom MSL shader implementation, register utilization, and threadgroup tiling.

---

Running Tests & Benchmarks

# Run the test suite (14 unit tests, ~0.1s on Metal GPU)
uv run --with mlx python -m unittest test_kan.py

# Run the throughput benchmark
uv run --with mlx python benchmark.py

# Run the iso-parameter stress test
uv run --with mlx python iso_param_stress_test.py

# Run the comparative nonlinear regression demo
uv run --with mlx python example.py

---

License

Released under the MIT License.