Adaptive Training of Grid-Dependent Physics-Informed Kolmogorov-Arnold Networks

TL;DR

This paper introduces jaxKAN, an open-source, JAX-based framework for grid-dependent Physics-Informed Kolmogorov-Arnold Networks (PIKANs) to solve PDEs. It combines grid extension, adaptive state transitions, residual-based loss weighting, and residual-based collocation point re-sampling to achieve up to two orders of magnitude faster training than prior KAN implementations while attaining accuracy competitive with, or superior to, larger architectures. The study demonstrates substantial reductions in relative L^2 error across diffusion, Helmholtz, Burgers, and Allen–Cahn equations and emphasizes the importance of grid-dependent basis functions and adaptive training for efficiency. The authors also compare static versus fully adaptive basis functions and show that non-static, fully adaptive bases can significantly improve performance, with the framework enabling flexible exploration of alternative basis designs for efficient PDE solving.

Abstract

Physics-Informed Neural Networks (PINNs) have emerged as a robust framework for solving Partial Differential Equations (PDEs) by approximating their solutions via neural networks and imposing physics-based constraints on the loss function. Traditionally, Multilayer Perceptrons (MLPs) have been the neural network of choice, with significant progress made in optimizing their training. Recently, Kolmogorov-Arnold Networks (KANs) were introduced as a viable alternative, with the potential of offering better interpretability and efficiency while requiring fewer parameters. In this paper, we present a fast JAX-based implementation of grid-dependent Physics-Informed Kolmogorov-Arnold Networks (PIKANs) for solving PDEs, achieving up to 84 times faster training times than the original KAN implementation. We propose an adaptive training scheme for PIKANs, introducing an adaptive state transition technique to avoid loss function peaks between grid extensions, and a methodology for designing PIKANs with alternative basis functions. Through comparative experiments, we demonstrate that the adaptive features significantly enhance solution accuracy, decreasing the L^2 error relative to the reference solution by up to 43.02%. For the studied PDEs, our methodology approaches or surpasses the results obtained from architectures that utilize up to 8.5 times more parameters, highlighting the potential of adaptive, grid-dependent PIKANs as a superior alternative in scientific and engineering applications.

Figures (12)

• Figure 1: Schematic representation of a PIKAN with an underlying [2,3,2,1] KAN architecture.
• Figure 2: PIKAN results for the Diffusion equation (first row), the Helmholtz equation (second row), Burgers' equation (third row), and the Allen–Cahn equation (fourth row). The first/third and second/fourth columns correspond to the solution obtained via pykan/jaxkan and its absolute error compared to the reference solution, respectively. In each row, the solutions/errors share the same colorbars.
• Figure 3: Introducing a grid update (top), grid adaptation (middle) and optimizer reset (bottom) process during the training of a PIKAN.
• Figure 4: Training a PIKAN without (top) and with (bottom) the adaptive state transition technique.
• Figure 5: Training loss of a PIKAN for the Allen-Cahn equation with (blue, solid line) and without (orange, dashed line) RBA.