lmKANs: Efficient Spline-Based Neural Nets

• lmKANs are neural network architectures that use low-dimensional spline lookup tables based on the Kolmogorov-Arnold theorem for efficient function approximation.
• They achieve significant reductions in inference FLOPs and higher GPU throughput compared to traditional MLPs, CNNs, and Transformers.
• The design enables scalable implementations for tasks like image classification and sequence modeling, opening avenues for further research in grid and spline optimization.

Lookup Multivariate Kolmogorov-Arnold Networks (lmKANs) are designed as efficient alternatives to traditional high-dimensional linear mappings used in various neural network architectures, including multi-layer perceptrons (MLPs), Transformers, and convolutional neural networks (CNNs). This approach capitalizes on the theoretical insights from the Kolmogorov-Arnold representation theorem, which allows a high-dimensional function to be approximated by a sum of lower-dimensional functions. Implementing these low-dimensional functions with spline lookup tables, lmKANs are capable of providing substantial improvements in computational efficiency while maintaining flexibility in function approximation.

Theoretical Foundations of lmKANs

The lmKANs are inspired by the Kolmogorov-Arnold representation theorem which posits that any continuous multivariate function can be decomposed into a composition and summation of univariate functions. This decomposition expresses a high-dimensional function $f: \mathbb{R}^n \to \mathbb{R}$ as:

$$f(\mathbf{x}) = \sum_{k=1}^{2n+1} \phi_k \left( \sum_{j=1}^n \psi_{k,j}(x_j) \right)$$

where $\phi_k$ and $\psi_{k,j}$ are univariate functions parameterized as splines. This structured approach allows the implementation with trainable basis functions (e.g., B-splines) to manage complexity efficiently.

Construction and Technical Details

In lmKANs, high-dimensional mappings are expressed through collections of trainable, low-dimensional multivariate functions, each implemented as spline lookup tables. Specifically, a 2D lmKAN layer groups input features into pairs such that each output component results from evaluating a spline function based on these pairs. The mathematical representation of a two-dimensional function is:

$$f(x_1, x_2) = \sum_{i_1, i_2} p_{i_1,i_2} \cdot B_{i_1,i_2}(x_1, x_2)$$

where $B_{i_1,i_2}(x_1, x_2)$ are basis functions defined on a static "sigma grid" and $p_{i_1,i_2}$ are learned parameters. The lookup nature of the spline allows for efficient computation by reducing the operation count to a handful within each function evaluation.

Performance Metrics

Empirical evaluations have demonstrated that lmKANs can significantly reduce inference floating-point operations (FLOPs) while matching the accuracy of traditional MLPs in approximating complex functions. Key performance benefits observed include:

• Up to $6\times$ reduction in inference FLOPs compared to MLPs.
• $10\times$ higher H100 GPU throughput in specific benchmarks.
• Reductions of 1.6–2.1× for CNNs on CIFAR-10 and roughly 1.7× on ImageNet-1k datasets, maintaining accuracy while reducing computational demands.

Applications and Implications

The lmKANs serve as efficient replacements for linear layers within various neural network frameworks, opening new possibilities for enhancing model performance across applications such as image classification, feedforward networks, and sequence models. The ability to lessen computational cost while increasing the parameter count allows enhanced model expressivity without proportional increases in computational demand.

Code and Implementation

A comprehensive codebase, featuring optimized CUDA kernels tailored for modern GPUs (e.g., H100), is available online. These implementations utilize shared-memory tiling to achieve significant throughput optimization. The kernels support float32 computations, with potential future expansions to include lower precision computations like bfloat16.

Link to the lmKAN code repository: https://github.com/schwallergroup/lmkan.

Future Directions and Challenges

The lmKAN approach highlights several challenges and avenues for further research:

• Efficiency Scaling: While lmKANs show promise in reducing computational burden, further research into better grid designs and faster spline computations can expand their applicability.
• Hyperparameter Optimization: Automated tuning approaches like Bayesian optimization can enhance model performance predictably.
• Theory-Driven Improvements: Developing stronger theoretical foundations will help bridge existence claims to practical training dynamics, ensuring robust scalability.

In summary, Lookup Multivariate Kolmogorov-Arnold Networks (lmKANs) present a compelling strategy for optimizing deep learning models by leveraging efficient spline-based representations to balance inference cost and capacity, supporting their widespread adoption across modern neural network architectures.