LUT-KAN: Segment-wise LUT Quantization for Fast KAN Inference

TL;DR

LUT-KAN introduces a segment-wise LUT-based quantization pipeline for PyKAN-style KAN layers to enable fast CPU inference with explicit OOB semantics. By converting each edge function into per-segment LUTs with affine int8/uint8 quantization and linear interpolation, the approach achieves substantial speedups (honest baselines of roughly 10–14× on NumPy and 9–11× on Numba) while maintaining negligible accuracy loss (MAE ≈ 1.6×10^-4 at L=64) and a memory overhead around 10×. The framework defines a rigorous deployment contract with boundary conventions and OOB policies, and validates these through an honest baseline methodology, controlled sweeps over LUT resolution, and an end-to-end DoS detection case study on CICIDS2017 that preserves F1 and accuracy while reducing latency. The results demonstrate LUT-KAN’s practical impact for edge deployment, enabling reproducible artifacts, improved cache locality, and a clear path toward deployment-ready KAN inference on CPU systems. Limitations include the linear memory growth for dense layers and the need for careful calibration; future work points to sparse representations, fully integer pipelines, and adaptive quantization.

Abstract

Kolmogorov--Arnold Networks (KAN) replace scalar weights by learnable univariate functions, often implemented with B-splines. This design can be accurate and interpretable, but it makes inference expensive on CPU because each layer requires many spline evaluations. Standard quantization toolchains are also hard to apply because the main computation is not a matrix multiply but repeated spline basis evaluation. This paper introduces LUT-KAN, a segment-wise lookup-table (LUT) compilation and quantization method for PyKAN-style KAN layers. LUT-KAN converts each edge function into a per-segment LUT with affine int8/uint8 quantization and linear interpolation. The method provides an explicit and reproducible inference contract, including boundary conventions and out-of-bounds (OOB) policies. We propose an ``honest baseline'' methodology for speed evaluation: B-spline evaluation and LUT evaluation are compared under the same backend optimization (NumPy vs NumPy and Numba vs Numba), which separates representation gains from vectorization and JIT effects. Experiments include controlled sweeps over LUT resolution L in 16, 32, 64, 128 and two quantization schemes (symmetric int8 and asymmetric uint8). We report accuracy, speed, and memory metrics with mean and standard deviation across multiple seeds. A two-by-two OOB robustness matrix evaluates behavior under different boundary modes and OOB policies. In a case study, we compile a trained KAN model for DoS attack detection (CICIDS2017 pipeline) into LUT artifacts. The compiled model preserves classification quality (F1 drop below 0.0002) while reducing steady-state CPU inference latency by 12x under NumPy and 10x under Numba backends (honest baseline). The memory overhead is approximately 10x at L=64. All code and artifacts are publicly available with fixed release tags for reproducibility.

Figures (7)

• Figure 1: Quantization error versus LUT resolution $L$ on in-range inputs (mean absolute error, MAE). Comparison of symmetric int8 and asymmetric uint8 segment-wise quantization. Error bars show $\pm 1$ std across 5 seeds.
• Figure 2: Maximum absolute quantization error versus LUT resolution $L$ on in-range inputs. Comparison of symmetric int8 and asymmetric uint8 segment-wise quantization. Error bars show $\pm 1$ std across 5 seeds.
• Figure 3: Honest baseline speedup of LUT inference relative to B-spline evaluation as a function of $L$ (in-range inputs). (a) NumPy backend. (b) Numba backend. Comparison of symmetric int8 and asymmetric uint8 LUT artifacts.
• Figure 4: Absolute latency (ms/iteration) versus $L$ for the sanity-layer benchmark under the Numba backend. Curves show Numba B-spline evaluation and Numba LUT evaluation. LUT latency is nearly flat across $L$ because the kernel is memory-bound.
• Figure 5: Memory footprint versus $L$. (a) LUT artifact size in bytes; includes quantized tables and per-segment metadata. (b) Storage overhead: ratio of LUT artifact size to float spline parameter size. Symmetric int8 and asymmetric uint8 are reported.