Scalable Bayesian Physics-Informed Kolmogorov-Arnold Networks
Zhiwei Gao, George Em Karniadakis
TL;DR
This work advances scalable Bayesian UQ for PDE surrogates by integrating compact Chebyshev Kolmogorov-Arnol d networks (cKANs) with a gradient-free dropout Tikhonov Ensemble Kalman Inversion (DTEKI) framework. An active-subspace mechanism reduces parameter dimensionality, enabling efficient training and robust uncertainty estimates on large datasets, often achieving accuracy comparable to or better than HMC while dramatically reducing computation. Across inverse/forward tests (transport, diffusion, nonlinear, and Darcy flow), DTEKI and its subspace variant (SDTEKI) consistently deliver well-calibrated posterior uncertainty, accelerated convergence, and substantial speed-ups, highlighting their potential for scalable Bayesian PINNs in high-dimensional settings. The results suggest that cKANs combined with gradient-free, regularized ensemble methods offer a practical pathway to reliable, scalable UQ for complex PDE systems.
Abstract
Uncertainty quantification (UQ) plays a pivotal role in scientific machine learning, especially when surrogate models are used to approximate complex systems. Although multilayer perceptions (MLPs) are commonly employed as surrogates, they often suffer from overfitting due to their large number of parameters. Kolmogorov-Arnold networks (KANs) offer an alternative solution with fewer parameters. However, gradient-based inference methods, such as Hamiltonian Monte Carlo (HMC), may result in computational inefficiency when applied to KANs, especially for large-scale datasets, due to the high cost of back-propagation. To address these challenges, we propose a novel approach, combining the dropout Tikhonov ensemble Kalman inversion (DTEKI) with Chebyshev KANs. This gradient-free method effectively mitigates overfitting and enhances numerical stability. Additionally, we incorporate the active subspace method to reduce the parameter-space dimensionality, allowing us to improve the accuracy of predictions and obtain more reliable uncertainty estimates. Extensive experiments demonstrate the efficacy of our approach in various test cases, including scenarios with large datasets and high noise levels. Our results show that the new method achieves comparable or better accuracy, much higher efficiency as well as stability compared to HMC, in addition to scalability. Moreover, by leveraging the low-dimensional parameter subspace, our method preserves prediction accuracy while substantially reducing further the computational cost.