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

TL;DR

This paper introduces an adaptive, grid-dependent Physics-Informed Kolmogorov-Arnold Network (PIKAN) framework with a fast JAX implementation, significantly improving training speed and solution accuracy for PDEs over traditional methods.

Contribution

The paper presents a novel adaptive training scheme for PIKANs, enhancing efficiency and accuracy, and demonstrates their superiority over larger architectures in solving PDEs.

Findings

Training times up to 84x faster than original KAN

Reduced L^2 error by up to 43.02% with adaptive methods

Achieved comparable or better results with fewer parameters

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.