Uncertainty Quantification with Bayesian Higher Order ReLU KANs

TL;DR

This work introduces the first method for uncertainty quantification in Kolmogorov-Arnold Networks (KANs), focusing on (Higher Order) ReLU KANs to improve computational efficiency in Bayesian settings. It defines variational posteriors over basis-function parameters $e_i,s_i$ and weights, using a mixture Gaussian with inverse normalizing flows to capture complex dependencies, and employs a surrogate Bayesian KAN to model aleatoric uncertainty with Gaussian or Student‑t likelihoods. The approach is validated on simple 1D function fits and two-dimensional stochastic PDEs (Poisson and Helmholtz), demonstrating accurate mean predictions and meaningful separation of epistemic and aleatoric uncertainty, with robustness to outliers via the Student‑t likelihood. While Bayes-HR-KAN matches deterministic HR-KAN performance in MSE, it incurs higher training cost due to KL-term computations, underscoring the trade-off between uncertainty quantification and computational overhead. Overall, the method enables efficient, interpretable uncertainty-aware KAN surrogates for scientific computing and PDE applications, with potential generalization to other basis-function families.

Abstract

We introduce the first method of uncertainty quantification in the domain of Kolmogorov-Arnold Networks, specifically focusing on (Higher Order) ReLUKANs to enhance computational efficiency given the computational demands of Bayesian methods. The method we propose is general in nature, providing access to both epistemic and aleatoric uncertainties. It is also capable of generalization to other various basis functions. We validate our method through a series of closure tests, including simple one-dimensional functions and application to the domain of (Stochastic) Partial Differential Equations. Referring to the latter, we demonstrate the method's ability to correctly identify functional dependencies introduced through the inclusion of a stochastic term. The code supporting this work can be found at https://github.com/wmdataphys/Bayesian-HR-KAN

Figures (7)

• Figure 1: Bayesian (Higher Order) ReLU KAN: The method consists of two Bayesian (Higher Order) ReLU KANs, in which we define a surrogate model tasked with learning the aleatoric component in the likelihoods. This model is capable of inheriting the structure of the functional KAN, or with its own unique set of parameters.
• Figure 2: One Dimensional Fits: Underlying function fitted with Bayesian ReLUKAN (left column), noised function fit with a Gaussian likelihood (center column) and noise function fit with Student-t likelihood (right column). Noise is sampled from a Student-t distribution with $\nu=3$. Note in all cases, the epistemic uncertainty is very small and not visible. This is in agreement with the true inaccuracy, measured as the difference between truth and mean prediction, being almost overlapping. The aleatoric term is captured with high fidelity under both likelihood assumptions.
• Figure 3: Poisson's Equation with Gaussian Noise: Three-dimensional representations of the fitting capability of the ReLU-KAN (top row, second column), the HRKAN (top row, third column) and the Bayesian HRKAN (top row, fourth column) on the Poisson equation (top left). The residuals ($u - \hat{u}$) for each model are shown directly below and highly resemble the noise contribution, implying accurate fits towards the average values of the function.
• Figure 4: Validation of Uncertainty on Poisson's Equation: Epistemic uncertainty (A), aleatoric uncertainty (B), absolute error (C), and the true aleatoric component (absolute value) are shown (D). The epistemic uncertainty quantifies the deviation of the learned function from the true mean, with an average value of $\bar{\sigma}_{epi.} = 0.0035 \pm 0.0001$. The accuracy of the learned aleatoric component can be observed by comparing the second and fourth plots. Note that the two rightmost plots depict individual samples rather than average behavior.
• Figure 5: Helmholtz's Equation with Gaussian Noise: Three-dimensional representations of the fitting capability of the ReLU-KAN (top row, second column), the HRKAN (top row, third column) and the Bayesian HRKAN (top row, fourth column) on the Helmholtz equation (top left). The residuals ($u - \hat{u}$) for each model are shown directly below and highly resemble the noise contribution, implying accurate fits towards the average values of the function. The learned solution space is with $a_1 = 1, a_2 = 2, k=1$.