Fused truncated-power KAN activation
Implementation & benchmark of arXiv:2609.01956 — FlashKAN: B-Spline KANs via Truncated Power Form
Source paper: FlashKAN: B-Spline KANs via Truncated Power Form arXiv:2609.01956 ⓘ · analyzed Sep 3, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace Cox-de Boor evaluation of each cubic B-spline edge activation with its fixed truncated-power expansion. Normalize each scalar edge input to a bounded knot coordinate, evaluate the five shifted cubic positive-part terms in parallel, and contract them with the learned spline coefficients inside one fused kernel.
Formulas
Mathematical statement
The paper introduces the cubic Bernstein representation Q(t)=\sum_{i=0}^{3}B_{i,3}(t)P_i, with B_{i,3}(t)=\binom{3}{i}t^i(1-t)^{3-i}; this establishes the polynomial-basis viewpoint and the partition of unity \sum_i B_{i,3}(t)=1. For a uniform cubic cardinal B-spline, use the equivalent truncated-power form B_3(z)=\frac{1}{6}\sum_{r=0}^{4}a_r(z-r)_+^3, where (v)_+=\max(v,0), r is the shift index, and a=(1,-4,6,-4,1). If u=(x-x_{\min})/h is the normalized input coordinate, h is the uniform knot spacing, and c_i are learned coefficients for basis index i, the KAN edge activation is \phi(x)=\sum_{i=0}^{N-1}c_iB_3(u-i)=\frac{1}{6}\sum_{i=0}^{N-1}c_i\sum_{r=0}^{4}a_r(u-i-r)_+^3. Clamp u to the valid bounded coordinate interval before evaluation. The derivative needed by backpropagation is \partial\phi/\partial u=\frac{1}{2}\sum_i c_i\sum_{r=0}^{4}a_r(u-i-r)_+^2, away from knot points. The finite-difference coefficients cause cancellation of polynomial pieces outside spline support, while bounded u prevents very large shifted powers from worsening floating-point cancellation.
Implementation notes
(1) Integration point: replace the forward method of a cubic KAN edge layer, after the linear input has produced one scalar x per edge and before summing edge outputs into the next node. Store spline coefficients C with shape [out_features, in_features, N], and use a shared uniform knot spacing h per layer or per edge. Normalize x to u=(x-x_min)/h and clamp u to [0,N], or to the explicitly chosen valid support interval. Compute the activation in float32 under mixed precision, then cast the result back to the model dtype.
(2) Pseudocode:
u = ((x - xmin) / h).clamp(0.0, float(N))
I = torch.arange(N, device=x.device)[None, :, None]
R = torch.arange(5, device=x.device)[None, None, :]
a = torch.tensor([1, -4, 6, -4, 1], device=x.device, dtype=torch.float32)[None, None, :]
z = u.float()[..., None, None] - I[..., None] - R
B = (torch.relu(z) ** 3 * a).sum(dim=-1) / 6.0
y = (B * C).sum(dim=-1)
For a full KAN layer, broadcast x over input edges, evaluate this expression under torch.compile, and sum over input features. Fuse normalization, clamp, shifted powers, coefficient multiplication, and reductions; do not materialize B in a custom Triton or CUDA kernel. Autograd can differentiate the expression directly, or the backward kernel can use the displayed squared-positive-part derivative.
(3) Computed from the mathematics: the five shifts, coefficients [1,-4,6,-4,1], factor 1/6, bounded-coordinate rule, and analytic derivative. Estimated empirically: x_min, x_max or h if data-adaptive, kernel occupancy, register pressure, and whether bfloat16 or float16 accumulation is accurate enough. Verify numerical equivalence against a reference Cox-de Boor implementation on random values, especially near knots and boundaries.
(4) First cheap experiment: train a small KAN MLP on MNIST or Fashion-MNIST with two hidden layers, width 64, degree-3 splines, and N=16 or 32 knots. Compare the recursive layer, an unfused vectorized truncated-power layer, and the torch.compile or Triton-fused version with identical initialization, optimizer, batch size, and coefficients. Measure forward latency, full training-step latency, peak memory, output and gradient maximum error, and validation accuracy. Success means at least 2x lower forward latency and materially lower step time at equal predictions, with float32 spline-output error below 1e-5 and no accuracy degradation. A secondary test should evaluate inputs outside the knot range and check that clamping prevents NaNs or large numerical discrepancies.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built a PyTorch fused truncated-power cubic spline evaluator with bounded coordinate clamping, Cox–de Boor reference evaluation, autograd derivative checks, precision diagnostics, and CUDA timing. The formulas agree to approximately 1.7e-12 max error in float64, but float32 cancellation grows to 9.1e-4 output error and 1.17e-3 gradient error near the upper bounded coordinate range. The proposed evaluator was only 1.30x faster in forward and 1.26x faster for forward-plus-backward, so the claimed substantial speedup and float32 accuracy target were not observed.
- Agent confidence
- 9/10
- Baseline
- Cox–de Boor: 2.623 ms forward, 5.538 ms forward+backward; float32 comparison error 9.07e-4 max output and 1.17e-3 max gradient.
- Idea
- Truncated-power: 2.020 ms forward, 4.397 ms forward+backward; 1.30x and 1.26x speedups respectively. Float64 max error 1.70e-12, but float32 equivalence is insufficiently accurate.
Limitations: No end-to-end KAN training or validation-accuracy experiment was run. The implementation is vectorized PyTorch rather than a custom Triton/CUDA fused kernel or torch.compile graph, so these timings do not establish the performance of a true kernel fusion. The float32 cancellation issue was measured but not mitigated with compensated arithmetic or a numerically stable alternative.
How to run: python3 experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- experiment.py 3.5 KB View Raw
- report.md 1.6 KB View
- results.json 0.6 KB View Raw