KAN-SAs: Efficient Acceleration of Kolmogorov-Arnold Networks on Systolic Arrays

TL;DR

KAN-SAs tackle the end-to-end acceleration of Kolmogorov-Arnold Networks on systolic arrays by addressing the recursive B-spline evaluation bottleneck. The key idea is a Basis Function Unit that tabulates B-splines and an N:M sparsity-aware PE design that leverages the local support of B-splines to maximize SA utilization. Hardware synthesis on 28nm FD-SOI shows notable gains: higher PE utilization, reduced clock cycles, and up to tens-fold improvements in B-spline evaluation compared to prior methods. The design space across multiple KAN workloads confirms the practical efficiency and scalability of the approach for end-to-end KAN inference on spatial accelerators.

Abstract

Kolmogorov-Arnold Networks (KANs) have garnered significant attention for their promise of improved parameter efficiency and explainability compared to traditional Deep Neural Networks (DNNs). KANs' key innovation lies in the use of learnable non-linear activation functions, which are parametrized as splines. Splines are expressed as a linear combination of basis functions (B-splines). B-splines prove particularly challenging to accelerate due to their recursive definition. Systolic Array (SA)based architectures have shown great promise as DNN accelerators thanks to their energy efficiency and low latency. However, their suitability and efficiency in accelerating KANs have never been assessed. Thus, in this work, we explore the use of SA architecture to accelerate the KAN inference. We show that, while SAs can be used to accelerate part of the KAN inference, their utilization can be reduced to 30%. Hence, we propose KAN-SAs, a novel SA-based accelerator that leverages intrinsic properties of B-splines to enable efficient KAN inference. By including a nonrecursive B-spline implementation and leveraging the intrinsic KAN sparsity, KAN-SAs enhances conventional SAs, enabling efficient KAN inference, in addition to conventional DNNs. KAN-SAs achieves up to 100% SA utilization and up to 50% clock cycles reduction compared to conventional SAs of equivalent area, as shown by hardware synthesis results on a 28nm FD-SOI technology. We also evaluate different configurations of the accelerator on various KAN applications, confirming the improved efficiency of KAN inference provided by KAN-SAs.

Figures (8)

• Figure 1: A single layer of dimensions [2, 1] for MLP (left) and KAN (right).
• Figure 2: B-spline basis functions of degree $P=3$ on a uniform grid of size $G=3$ for the input domain $[t_3, t_6]$ that is extended on both ends to account for the contributions of B-spline functions that are non-zero within this input domain.
• Figure 3: GEMM weight stationary (WS) systolic array with an additional B-spline unit for on-the-fly B-spline activation computation
• Figure 4: Illustration of efficient tabular storing of half a cubic B-spline and computation of the others by aligning and shifting the input $x$. Only the values of $B_{0, 3}$ corresponding to the black sample points on the horizontal axis quantizing the interval $[0, 2]$ are stored. For a given input $x$, all the B-spline values can be obtained by retrieving $B_{0, 3}$ stored values at the aligned input $x_a$ and by shifting the evaluation point. E.g., $B_{t_{k - 2},3}(x) = B_{0,3}(2 - x_a)$.
• Figure 5: Simplified version of the B-Spline unit. The address is inverted, and the corresponding values are reverse-packed to obtain the remaining non-zero activations. The example values $0, 32$ at the output correspond to $x_{addr} = 0$. The unit also returns the values at $\sim x_{addr} = 3$ in reverse. i.e. $127, 32$.