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

Spyros Rigas${}^{\href https://orcid.org/0009-0009-2352-8709}$, Michalis
Papachristou${}^{\href https://orcid.org/0000-0001-5650-4206}$, Theofilos
Papadopoulos${}^{\href https://orcid.org/0000-0002-8534-785X}$, Fotios
Anagnostopoulos${}^{\href https://orcid.org/0000-0003-0176-4144}$ and Georgios
Alexandridis${}^{\href https://orcid.org/0000-0002-3611-8292}$ Spyros Rigas
and Georgios Alexandridis are with the Department of Digital Industry
Technologies, School of Science, National and Kapodistrian University of
Athens, 34400 Greece (e-mail<EMAIL_ADDRESS>gealexandri@uoa.gr).Michalis
Papachristou is with the Department of Physics, School of Science, National
and Kapodistrian University of Athens, 15784 Greece (e-mail:
mixpap@phys.uoa.gr).Theofilos Papadopoulos is with the School of Electrical &
Computer Engineering, National Technical University of Athens, 15780 Greece
(e-mail: teopap@mail.ntua.gr).Fotios Anagnostopoulos is with the Department of
Informatics & Telecommunications, University of the Peloponnese, 22131 Greece
(e-mail: fotisanagn@uop.gr).

###### 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) are the neural
network of choice, and significant progress has been 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. We propose an adaptive training scheme for
PIKANs, incorporating known MLP-based PINN techniques, introducing an adaptive
state transition scheme to avoid loss function peaks between grid updates, and
proposing a methodology for designing PIKANs with alternative basis functions.
Through comparative experiments we demonstrate that these adaptive features
significantly enhance training efficiency and solution accuracy. Our results
illustrate the effectiveness of PIKANs in improving performance for PDE
solutions, highlighting their potential as a superior alternative in
scientific and engineering applications.

###### Index Terms:

Adaptive training, deep learning for science and engineering, grid-dependent
basis functions, JAX, Kolmogorov-Arnold networks, PDEs, physics-informed
neural networks.

##  I Introduction

The ever-present need to address problems related to differential equations
for modeling complex systems, combined with the exponential progress in Deep
Learning (DL), has led to the emergence of a new paradigm: Physics-Informed
Neural Networks (PINNs) [1, 2, 3]. PINNs provide a framework for solving
partial differential equations (PDEs) by transforming them into a loss
function optimization problem. Specifically, PINNs employ a deep neural
network to represent the solution of a PDE within a given domain.
Incorporating experimental data is essential for inverse problems, however in
the case of forward problems it is sufficient to sample points within the
domain, known as collocation points, and train the network in an unsupervised
manner by encoding the underlying equation and boundary conditions as terms of
the loss function. This approach offers several advantages over traditional
methods, including the ability to provide mesh-free solutions, generate
solutions with tractable analytical gradients and jointly solve forward and
inverse problems within the same framework.

Despite their relatively recent introduction to the mainstream DL literature,
PINNs have already demonstrated significant potential and versatility in
various scientific and engineering applications. Notable examples include the
simulation of fluid flows as described by the Euler [4] or Navier-Stokes
equations [5, 6], applications spanning from medicine [7, 8] to electric power
systems [9, 10], and even domain-specific problems in physics, such as
modeling astrophysical shocks [11] or calculating eigenfunctions and
eigenvalues for quantum mechanical systems [12]. Nonetheless, PINNs are not
without shortcomings, which can stem from the framework itself, such as
weighting imbalances in the loss function’s terms [13], or the underlying DL
architecture, for instance, the issue of spectral bias inherent in Multilayer
Perceptrons (MLPs) with Rectified Linear Unit (ReLU) activations [14]. To
mitigate such issues, numerous studies have focused on improving the
performance of PINNs by proposing alternative architectures beyond the simple,
yet prevalent MLP [13, 15, 16, 17], or by employing adaptive training
strategies. The latter include, but are not limited to, adaptive re-weighting
of the loss function’s terms [13, 18, 19, 20, 21, 22], embedding gradient
information of the PDE residuals in the loss function [23, 24], adaptive re-
sampling of collocation points [18, 25, 26], and utilizing adaptive activation
functions [27, 28].

Inspired by the Kolmogorov-Arnold representation theorem, a novel architecture
called Kolmogorov-Arnold Networks (KANs) was recently proposed as an
alternative to MLPs [29]. Unlike in the case of MLPs, the computational graph
of KANs involves learnable activation functions at its edges and a sum
operation on its nodes, aiming to achieve higher accuracy and
interpretability, while utilizing a considerably smaller number of parameters.
Due to their potential, KANs have already been successfully adopted for tasks
like time-series analysis [30, 31, 32], image recognition [33, 34], image
segmentation [35] and human activity recognition [36]. However, their most
promising application lies in addressing complex scientific and engineering
problems that can be formulated as (symbolic) regression or PDE solving [29].
In this respect, the term Physics-Informed Kolmogorov Arnold Networks (PIKANs)
has been coined to refer to PINNs that are based on KANS for their underlying
architecture [37].

Naturally, KANs are ideal candidates for the PINN framework, as they
correspond to an innovative architecture with inherent adaptability in its
activation functions. Preliminary studies on PIKANs have shown results
comparable [37] or superior [29, 38] to those of MLPs. Nevertheless, this
adaptability is also the cause of a significant drawback; the training of
KANs, and by extension PIKANs, is computationally expensive, as a consequence
of using learnable B-Splines as activation functions. To address this issue,
following studies have proposed more computationally efficient activation
functions, such as radial basis functions [39], Chebyshev polynomials [40],
wavelets [41], and ReLU-based functions [42]. For the latter, a speedup factor
between 5 and 20 in training time has been reported compared to the original
KAN implementation.

Albeit successful in increasing their computational efficiency, a common
pattern in all the aforementioned alternatives is that they remove the
dependence of PIKANs on the grid. The grid is the set of points used to
construct the basis functions, whose linear combinations correspond to the
learnable activations. Chebyshev (and other Jacobi) polynomials are inherently
grid-independent, while the implementations involving radial basis functions,
wavelets, and ReLU-based functions focus on uniform or static grids, which
lead to fixed basis functions. Herein, we argue that this may adversely affect
the training of PIKANs, resulting in slower convergence and potentially less
accurate solutions. Additionally, these implementations fail to leverage one
of the key features of KANs which has been reported to significantly enhance
their accuracy [29]: the ability to dynamically extend the grid during
training, thereby obtaining progressively more fine-grained basis functions.

The contributions of the current work can be summarized in the following
points:

  * •

The introduction of a new, publicly available computational framework for
KANs, developed using the JAX [43] and Flax [44] Python libraries.

  * •

The study of an issue involving abrupt jumps in the loss function’s values
after grid extensions and the introduction of an adaptive transition method to
address it.

  * •

The adaptation of loss re-weighting and collocation re-sampling schemes to a
grid-dependent framework for training PIKANs with relatively small
architectures.

  * •

The proposition of a general direction for designing (PI)KANs with alternative
basis functions, emphasizing the importance of preserving their dependency on
the grid to enable more adaptive training.

The remainder of this paper is structured as follows; in Section II we
formally review the PIKAN framework and introduce our JAX-based
implementation, including benchmarks compared to the original KAN
implementation. Section III presents the proposed adaptive training schemes,
aiming to form the foundation for future standardized end-to-end pipelines in
PIKAN training, inspired by [45]. Additionally, we provide a solution which
involves regular grid adaptations and a simple linear interpolation of the
optimizer’s parameters during training, to solve the issue of the loss
function’s sharp peaks after grid extensions. A thorough discussion on KAN
basis functions follows in Section IV, emphasizing the effect of grid
dependency on PIKAN accuracy. Using ReLU-KANs as a case study, we present a
method to make the basis functions fully adaptive to the grid. Finally, in
Section V, we summarize the main results of our work and propose avenues for
future research on the field of PIKANs.

##  II PIKAN Framework

Before discussing the technical details of the JAX-based implementation
developed to accelerate the training of PIKANs, we provide an overview of
their constituents.

###  II-A PINN Problem Formulation

Consider a differential equation of the generic form

$\displaystyle\mathcal{F}_{\lambda}\left(u\left(x\right)\right)$
$\displaystyle=f\left(x\right),~{}x\in\Omega,$
$\displaystyle\mathcal{B}_{\lambda}^{k}\left(u\left(x\right)\right)$
$\displaystyle=b^{k}\left(x\right),~{}x\in\Gamma^{k}\subseteq\partial\Omega,$
(1)

where $x$ is the coordinate vector, $u$ is the equation’s solution, $\lambda$
corresponds to a set of model parameters, $\mathcal{F}$ is a nonlinear
differential operator and
$\left\\{\mathcal{B}^{k}\right\\}_{k=1}^{n_{\Gamma}}$ is a set of $n_{\Gamma}$
boundary condition operators. For time-dependent problems, the temporal
coordinate is integrated in $x$, therefore the problem’s initial conditions
are considered as a special case of boundary conditions in $\Omega$. The main
intuition behind PINNs is to approximate the exact solution with a neural
network, $u\left(x;\theta\right)$, parameterized by a set of parameters,
$\theta$, and minimize the loss function, $\mathcal{L}\left(\theta\right)$,
which encapsulates the differential equation and boundary conditions as
distinct residual terms.

The training of PINNs generally involves three different types of datasets.
The first, $\mathcal{D}_{f}$, is a dataset consisting of $N_{f}$ collocation
points which lie in the interior of $\Omega$ and are used to enforce the
physics captured by the differential equation. The second,
$\left\\{\mathcal{D}_{b}^{k}\right\\}_{k=1}^{n_{\Gamma}}$, is a series of
datasets containing $N_{b}^{k}$ collocation points each, which are used to
enforce the boundary conditions on $\Gamma^{k}$. Finally, $\mathcal{D}_{d}$ is
a dataset comprising tuples of experimental data coordinates and measured
values which add a supervised component to the network’s training. The latter
is necessary only when $\lambda$ are unknown, corresponding to the so-called
inverse problem of differential equations. Herein, we only study forward
problems, meaning that all $\lambda$ are considered known and
$\mathcal{D}_{d}$ is null. Therefore the loss function, often referred to as
the physics-informed loss, assumes the form

$\mathcal{L}\left(\theta\right)=w_{f}\mathcal{L}_{f}\left(\theta\right)+\sum_{k=1}^{n_{\Gamma}}{w_{b}^{k}\mathcal{L}_{b}^{k}\left(\theta\right)},$
(2)

where $w_{f}$ and $\left\\{w_{b}^{k}\right\\}_{k=1}^{n_{\Gamma}}$ are
individual weights of the loss function’s terms and

$\displaystyle\mathcal{L}_{f}\left(\theta\right)$
$\displaystyle=\frac{1}{N_{f}}\sum_{i=1}^{N_{f}}{\lVert\mathcal{F}\left(u\left(x_{i};\theta\right)\right)-f\left(x_{i}\right)\rVert^{2}},$
(3) $\displaystyle\mathcal{L}_{b}^{k}\left(\theta\right)$
$\displaystyle=\frac{1}{N_{b}^{k}}\sum_{i=1}^{N_{b}^{k}}{\lVert\mathcal{B}^{k}\left(u\left(x_{i};\theta\right)\right)-b^{k}\left(x_{i}\right)\rVert^{2}},$
(4)

with $\lVert\cdot\rVert$ denoting the L2 norm. Derivatives of
$u\left(x;\theta\right)$ of arbitrary order and with respect to any coordinate
that may appear in (3) or (4) are computed using automatic differentiation
(AD) [46].

###  II-B Kolmogorov-Arnold Networks

In the vast majority of problems involving PINNs, $u\left(x;\theta\right)$
corresponds to a MLP with fixed number of linear layers, each followed by a
non-linear activation. Denoting
$\theta=\left\\{W^{(l)},b^{(l)}\right\\}_{l=1}^{L}$, where $W^{(l)}$ and
$b^{(l)}$ is the weight matrix and bias of the $l$-th layer, respectively, an
MLP with $L$ layers and activation functions $\sigma$ can be written as

$u\left(x;\theta\right)=\left[\Phi^{(L)}\circ\dots\circ\Phi^{(1)}\right]\left(x\right),$
(5)

where $\circ$ denotes successive application of $\Phi^{(l)}$, and

$\Phi^{(l)}\left(x^{(l)}\right)=\sigma\left(W^{(l)}x^{(l)}+b^{(l)}\right).$
(6)

In the case of KANs, (5) still holds, however the non-linear activation on the
linear combination of the layer’s inputs is replaced by

$\Phi^{(l)}\left(x^{(l)}\right)=\begin{pmatrix}\phi_{l,1,1}(\cdot)&\cdots&\phi_{l,1,n_{l}}(\cdot)\\\
\phi_{l,2,1}(\cdot)&\cdots&\phi_{l,2,n_{l}}(\cdot)\\\ \vdots&\ddots&\vdots\\\
\phi_{l,n_{l+1},1}(\cdot)&\cdots&\phi_{l,n_{l+1},n_{l}}(\cdot)\end{pmatrix}x^{(l)},$
(7)

where $n_{l}$ is the number of input nodes for the $l$-th layer and
$\phi_{l,i,j}$ is the $l$-th layer’s univariate activation function,
connecting its $i$-th input node to its $j$-th output node in the network’s
computational graph. Obviously, a layer’s output node coincides with the next
layer’s input node. Essentially, in KANs, the concept of the layer does not
revolve around a set of nodes, but instead refers to the edges, which is where
the activation functions lie. For a layer with $n_{l}$ input nodes and
$n_{l+1}$ output nodes, the number of univariate activation functions is the
product $n_{l}\cdot n_{l+1}$. To refer to the architecture of a KAN, it
suffices to refer to its shape as an array of integers, where two consecutive
numbers correspond to a layer’s input and output nodes. For instance, the
architecture of a $L$-layered KAN is written as
$\left[n_{0},n_{1},\dots,n_{L}\right]$.

Figure 1: Schematic representation of a PIKAN with an underlying [2,3,2,1] KAN
architecture.

In the original implementation of KANs [29], referred to as “vanilla” KANs in
[37], the univariate activation functions are given by:

$\phi\left(x\right)=c_{r}r\left(x\right)+c_{B}B\left(x\right),$ (8)

where

$r\left(x\right)=\frac{x}{1+\exp\left(-x\right)}$ (9)

is a residual-like activation and

$B\left(x\right)=\sum_{i=1}^{G+k}{c_{i}B_{i}\left(x\right)}$ (10)

is a B-Spline activation of order $k$, defined on a grid (knot vector) with
$G$ intervals. For a given grid and value of $k$, a set of spline basis
functions $\left\\{B_{i}\right\\}_{i=1}^{G+k}$ is uniquely defined,
constituting the activation functions of (8) grid-dependent by design. The
parameters $c_{r}$, $c_{B}$ and $\left\\{c_{i}\right\\}_{i=1}^{G+k}$ are
trainable, which is why in KANs the activation functions are not fixed, unlike
in MLPs. Based on these, extending the framework of PINNs to PIKANs is
straightforward: simply replace the expression for $\Phi^{(l)}$ from (7) in
(5) and use the obtained $u\left(x;\theta\right)$ in (3), (4). A schematic
representation of an example PIKAN along with its underlying KAN’s computation
graph can be seen in Fig. 1.

###  II-C JAX Implementation

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.

JAX is a numerical computing library that combines the flexibility, just-in-
time compilation and parallelization capabilities of high-performance
frameworks, with the GPU/TPU acceleration and AD features of DL frameworks
like PyTorch [47] and TensorFlow [48]. Thanks to these, JAX has become
increasingly popular for developing PINN-related code [17, 45], offering
significant performance improvements in both training speed and computational
efficiency. Leveraging these advantages, we developed a JAX-based
implementation of KANs called jaxKAN, with the aim of accelerating the
training process of PIKANs. For this reason, features relevant to symbolic
regression, available in the original pykan implementation [29], were not
included. Instead, we focused on methods and utilities relevant specifically
to PIKANs. Our implementation is available as a Python package in PyPI [49]
and has already been used for research on PIKANs [50].

The jaxKAN package is organized into three main modules. The models module
includes the KANLayer and KAN classes, which inherit from Flax’s linen module
and correspond to the vanilla KAN implementation. This module is designed to
also accommodate alternative implementations of KANs, such as ReLU-KANs. The
bases module includes functions related to the basis functions used to
construct the learnable activations. For vanilla KANs, these are B-Splines,
calculated using an iterative version of the Cox-de Boor recursion formula
[51, 52]. Finally, the utils module contains utility functions pertinent to
the functionality of KANs, such as a parallel implementation of the least
squares algorithm (used to re-initialize the
$\left\\{c_{i}\right\\}_{i=1}^{G+k}$ coefficients after grid extensions) based
on JAX’s vmap. Additionally, it includes methods required for the adaptive
training of PIKANs, as further explained in Section III.

In addition to using a different framework and including implementations
beyond the vanilla KAN, jaxKAN has two significant differences compared to
pykan. First, we developed the grid update as an internal model method that
not only adapts the grid to the data but also extends its size. As
demonstrated in subsequent experiments, the grid has a considerable impact on
the performance of KANs, therefore fine-tuning the hyperparameters related to
its update is crucial when training PIKANs. The second major difference is our
choice of optimizer. Instead of the resource-intensive LBFGS [53], we opted
for the Adam optimizer [54] implemented through the Optax optimization
framework [55]. Notably, while more recent implementations of KANs have aimed
to enhance training efficiency by proposing alternative KAN layer structures
[56], we chose not to use them due to their inability to produce the
regularization terms presented in the original paper.

TABLE I: Average GPU training time per epoch, measured on NVIDIA’s GeForce RTX-4070 12GB. PDE | pykan | jaxKAN
---|---|---
Diffusion Equation | 189 ms | 2.6 ms
Helmholtz Equation | 237 ms | 2.8 ms
Burgers’ Equation | 231 ms | 3.1 ms
Allen-Cahn Equation | 194 ms | 2.9 ms

Before reviewing our adaptive training framework, we provide in Fig. 2 the
results of vanilla PIKANs which were trained for $5\cdot 10^{4}$ epochs to
solve four different PDEs: the diffusion equation (first row), the Helmholtz
equation (second row), Burgers’ equation (third row), and the Allen–Cahn
equation (fourth row). These PDEs were selected because the former two have
analytical solutions, where KANs may have an advantage over other
architectures [29], while the latter two do not have analytical solutions.
Each of these PDEs is presented in Appendix A, including the reference
solutions required to calculate the absolute errors shown in the second and
fourth columns of Fig. 2, as well as details regarding the sampling of
collocation points. The PIKANs were implemented in both jaxKAN and pykan with
identical hyper-parameters and architectures: in every case, $k=3$ was chosen
for the spline basis functions’ order, the underlying KAN’s architecture was
$[2,6,6,1]$ and a grid of $G=3$ intervals was used. The weights of the loss
function’s terms were all set equal to $1.0$, and a constant value of
$10^{-3}$ was selected for the learning rate.

With the exception of the Helmholtz equation, where the absolute error for
jaxKAN is smaller than that for pykan by an order of magnitude, the two
implementations produce similar results, both quantitatively and
qualitatively: the diffusion equation’s solution is obtained with relatively
good accuracy in both implementations; Burgers’ equation appears to be well
approximated in all regions except for $x=0$, where the well-known
discontinuity exists; the approximated solution for the Allen-Cahn equation is
visibly different from the reference solution in both cases. Nonetheless, a
significant discrepancy between the two implementations arises when comparing
their training times. As shown in Table I, our framework is two orders of
magnitude faster than pykan. Importantly, the GPU times reported for jaxKAN
include the large latencies incurred by the compilation of some functions,
meaning that, if the same experiments were performed with pre-compiled
functions, the gap between the benchmarks for the two implementations would be
even wider. To avoid biased comparisons, the pykan code used to train PIKANs
was taken from its official repository and was even optimized to an extent,
for example by removing function declarations within the training loop, or by
calculating second-order derivatives without computing the corresponding
Hessian matrix.

##  III Adaptive PIKAN Training

In order to surpass the results shown in Fig. 2 and obtain more accurate
solutions for Burgers’ and the Allen-Cahn equations, one may increase the
number of training epochs and use a deeper and wider PIKAN architecture. While
this “brute-force” approach has been extensively used in the field of deep
learning with proven results, we advocate for utilizing adaptive techniques
that can lead to comparably good results with smaller architectures and
shorter computation times. A series of such techniques are presented in this
Section. It is noted that, due to their significantly faster training times,
PIKANs implemented using our framework were used for all subsequent
experiments. Additionally, for all trained PIKANs, the order of the spline
basis functions was set to $k=3$.

###  III-A State Transition after Extension

The first technique introduced is unique to PIKANs and pertains to the grid
extensions that can be performed during the training of KANs to make their
basis functions progressively more fine-grained. An issue with this grid
extension process, that is present in most published works and has been noted
by [41], is the fact that the loss function’s values tend to experience sharp
increases immediately after it is performed. In fact, this phenomenon is one
of the reasons why the paper introducing Wav-KANs argued against using
B-Splines as activation functions. We identify the cause of this effect as the
combination of the grid’s adaptation that occurs simultaneously with its
extension and the re-initialization of the optimizer’s state which follows the
grid’s update.

During the adaptation of the grid for each activation function in a KAN layer,
the activation function’s inputs are sorted, and the minimum and maximum
values are chosen as the grid’s ends. Then, a uniform grid, $\mathcal{G}_{u}$,
is constructed by placing equidistant points between the grid’s ends.
Simultaneously, a fully adaptive grid, $\mathcal{G}_{a}$, is constructed by
choosing the interior points based on the distribution of the activation
function’s inputs. The final input-adapted grid, $\mathcal{G}$, is then given
as a linear mixing of the two:

$\mathcal{G}=g_{e}\mathcal{G}_{u}+\left(1-g_{e}\right)\mathcal{G}_{a},$ (11)

where $0\leq g_{e}\leq 1$ is the hyper-parameter that controls whether the
grid tends to be more uniform ($g_{e}>0.5$) or more adaptive to the inputs
($g_{e}<0.5$). Whenever a grid extension occurs, it is followed by the
adaptation described by (11). However, for a network that has been trained for
a number of epochs before grid extension, this adaptation process induces a
discontinuous change in the grid’s basis functions, which in turns leads to
discontinuities in the loss function’s values. To remedy this, we suggest
performing regular grid adaptations between consecutive grid extensions.

When it comes to the re-initialization of the optimizer’s state, it should be
a contributing factor for the sharp loss increase for any optimizer with an
internal state that handles the scaling of the gradients involved in updating
the model’s parameters. Adam is a typical example, where gradient scaling is
based on estimates of the first and second-order moments of the gradients,
using exponential moving averages. Each time Adam is re-initialized, the
values for the gradients’ first and second moments are reset to zero,
essentially restarting the training process, albeit with better guesses for
the parameters’ values. In order to avoid this “hard reset” of the training
process, we propose an adaptive optimizer state transition described by
Algorithm 1, where $t$ corresponds to the training epoch and $\mu$, $\nu$ are
the first and second-order moments of the parameters’ gradients, respectively.

Algorithm 1 Adaptive Optimizer State Transition

0: Old Adam state $\left(t,\mu,\nu\right)$, new grid size
$G^{\prime}\in\mathbb{N}$

0: New Adam state $\left(t^{\prime},\mu^{\prime},\nu^{\prime}\right)$

1: Copy state step:

2: $t^{\prime}\leftarrow t$

3: Copy all moments for $c_{r}$ and $c_{B}$:

4: $\mu^{\prime}\left(c_{r},c_{B}\right)\leftarrow\mu\left(c_{r},c_{B}\right)$

5: $\nu^{\prime}\left(c_{r},c_{B}\right)\leftarrow\nu\left(c_{r},c_{B}\right)$

6: Interpolate new moments for $\left\\{c_{i}\right\\}_{i=1}^{G^{\prime}+k}$
per function $\phi$ per layer $l$:

7: for each layer $l=0,\dots,L-1$ do

8: for each function $\phi=1,\dots,n_{l}\cdot n_{l+1}$ do

9:
$\mu^{\prime}\left(\left\\{c_{i}^{(l,\phi)}\right\\}_{i=1}^{G^{\prime}+k}\right)\leftarrow\text{Interp}\left[\mu\left(\left\\{c_{i}^{(l,\phi)}\right\\}_{i=1}^{G+k}\right)\right]$

10:
$\nu^{\prime}\left(\left\\{c_{i}^{(l,\phi)}\right\\}_{i=1}^{G^{\prime}+k}\right)\leftarrow\text{Interp}\left[\nu\left(\left\\{c_{i}^{(l,\phi)}\right\\}_{i=1}^{G+k}\right)\right]$

11: end for

12: end for

13: return $\left(t^{\prime},\mu^{\prime},\nu^{\prime}\right)$

Instead of re-initializing the optimizer’s state, we keep the values for the
moments fixed for the parameters that are not affected by the extension of the
grid, namely $c_{r}$ and $c_{B}$. As far as
$\left\\{c_{i}\right\\}_{i=1}^{G+k}$ are concerned, their number increases due
to the grid’s extension, so their moments cannot be retained “as-is” in the
new state. However, rather than resetting them, we assign values to them by
performing a linear interpolation based on the optimizer’s state prior to the
grid’s extension. This process is analogous to using the least squares
algorithm after grid extensions to get good estimates for the values of the
$\left\\{c_{i}\right\\}_{i=1}^{G+k}$ parameters.

Figure 3: Training a PIKAN without (top) and with (bottom) regular grid
adaptations and optimizer state transition.

We implemented this algorithm in jaxKAN using jax.numpy.interp for the linear
interpolation and jax.vmap to parallelize the interpolation across functions
within each layer. During experimentation we found that incorporating Nesterov
momentum into Adam [57], as well as using a scheduler to progressively reduce
the learning rate, makes the state transition even smoother. Fine-tuning is
necessary to determine the optimal scheduler for each problem. In Fig. 3, we
demonstrate the effect of this technique on the training of PIKANs. For the
purposes of this example, we trained two identical $[2,8,8,1]$ PIKANs for
Burgers’ equation with $g_{e}=0.05$ for 3480 epochs with a learning rate of
$5\cdot 10^{-4}$. Starting from a grid with size $G=3$, we performed an
extension to $G^{\prime}=10$ during the 2500th training epoch. The loss values
of the first PIKAN, which was used as the baseline, experienced a small peak
during the grid extension. In contrast, the loss values of the second PIKAN,
for which the adaptive optimizer state transition technique was combined with
regular grid adaptations (every 275 epochs prior to extension), remained
smooth during the extension. Additionally, the loss value of the adaptive
PIKAN at the end of the experiment was slightly lower compared to that of the
baseline.

It should be noted that this adaptive state transition is neither exclusive to
Adam nor exclusive to spline basis functions. In fact, it can even be applied
to problems involving KANs outside the field of PDEs. In Appendix B, we
present one such example, where it is used for KANs that are trained to learn
functions. There, the technique’s effect on the loss function is much more
prominent than in the case of PIKANs, not merely smoothing the sharp increases
in the loss function’s value but turning them to sharp decreases instead.

###  III-B Loss Re-Weighting

The next method for adaptive training is related to the loss function’s terms
and has already been applied for PINNs with MLPs, as well as PIKANs [37]. This
method, known as residual-based attention (RBA), was introduced in [22].
Within the context of RBA, additional terms are introduced in (3) and (4),
transforming them to:

$\displaystyle\mathcal{L}_{f}\left(\theta\right)$
$\displaystyle=\frac{1}{N_{f}}\sum_{i=1}^{N_{f}}{\lVert\alpha_{0,i}\left(\mathcal{F}\left(u\left(x_{i};\theta\right)\right)-f\left(x_{i}\right)\right)\rVert^{2}},$
(12) $\displaystyle\mathcal{L}_{b}^{k}\left(\theta\right)$
$\displaystyle=\frac{1}{N_{b}^{k}}\sum_{i=1}^{N_{b}^{k}}{\lVert\alpha_{k,i}\left(\mathcal{B}^{k}\left(u\left(x_{i};\theta\right)\right)-b^{k}\left(x_{i}\right)\right)\rVert^{2}},$
(13)

where $\left\\{\alpha_{k,i}\right\\}_{k=0}^{n_{\Gamma}}$ are self-adaptive
weights that, unlike the “global” weights $w_{f}$ and
$\left\\{w_{b}^{k}\right\\}_{k=1}^{n_{\Gamma}}$, define local scalings unique
to each collocation point. Their initial value is $1.0$ and they are updated
based on the rule:

$\alpha_{k,i}^{\left(t+1\right)}\leftarrow\left(1-\eta\right)\alpha_{k,i}^{\left(t\right)}+\eta\frac{\left|r_{k}\left(x_{i}\right)\right|}{\max_{i}\left(\left\\{\left|r_{k}\left(x_{i}\right)\right|\right\\}\right)},$
(14)

where $x_{i}$ is the $i$-th collocation point’s coordinate vector, $t$ is the
training epoch’s index, $0\leq\eta\leq 1$ is a mixing factor and

$\displaystyle
r_{k}\left(x_{i}\right)=\begin{cases}\mathcal{F}\left(u\left(x_{i};\theta\right)\right)-f\left(x_{i}\right),~{}\text{
for }~{}k=0\\\
\mathcal{B}^{k}\left(u\left(x_{i};\theta\right)\right)-b^{k}\left(x_{i}\right),~{}\text{
for }~{}k>0\end{cases}$ (15)

is the $i$-th collocation point’s residual with respect to the PDE ($k=0$) or
the boundary conditions ($k>0$).

Figure 4: Training loss of a PIKAN for the Allen-Cahn equation with (blue,
solid line) and without (orange, dashed line) RBA.

RBA is an adaptive loss re-weighting scheme where the contribution of each
collocation point to the overall loss function is scaled based on the value of
its residual. Fig. 4 demonstrates an example of applying RBA with
$\eta=10^{-4}$, which is the value used for all subsequent experiments
presented herein, to train a $[2,8,8,1]$ PIKAN for the Allen-Cahn equation. An
identical PIKAN without RBA is used as the baseline. During the first 2000
epochs, the two PIKANs appear identical in terms of their training. However,
as the values of the $\alpha$-weights start adapting, it becomes evident that
the PIKAN with RBA trains faster. It is noted that a grid extension $G=3\to
G^{\prime}=8$ occurs for both PIKANs at the 8000th training epoch, although it
does not lead to a sharp loss increase since the technique presented in the
previous subsection is used.

###  III-C Collocation Points Re-sampling

The final adaptive training technique has also been used to train MLP-based
PINNs, however its application for PIKANs requires caution. It involves re-
sampling collocation points during training using a residual-based adaptive
distribution (RAD) [26], which may interact with other adaptive methods
employed simultaneously, including the grid itself. The main idea of RAD is to
dynamically change the set of collocation points by making their distribution
denser in regions of the domain where the absolute values of the PDE residuals
are high and sparser in regions where they are close to zero. To achieve this,
we define:

$p\left(x_{i}\right)=\frac{\left|r_{0}^{a}\left(x_{i}\right)\right|}{\frac{1}{N_{f}}\sum_{i=1}^{N_{f}}{\left|r_{0}^{a}\left(x_{i}\right)\right|}}+c,$
(16)

where $r_{0}\left(x_{i}\right)$ is the PDE residual given by (15) and $a\geq
0$, $c\geq 0$ are the RAD method’s hyper-parameters. Based on (16), the
adaptive resampling is performed by sampling a set of dense points,
$\mathcal{S}$, within the PDE’s domain (much denser than the set of
collocation points) and then computing $p\left(x_{i}\right)$ for every
$x_{i}\in\mathcal{S}$. The new set of collocation points corresponds to
$N_{f}$ points sampled from $\mathcal{S}$, according to $p\left(x\right)$. The
implementation of this procedure for jaxKAN is described by Algorithm 2 and an
example application for $a=3$ and $c=1$ can be seen in Fig. 5, from a snapshot
taken from the 1000th training epoch of a $[2,8,8,1]$ PIKAN trained to solve
the Helmholtz equation.

Algorithm 2 Adaptive Resampling of Collocation Points

0: $N_{f}$, parameters $a,c\geq 0$, weights $\left\\{\alpha_{0,i}\right\\}$,
grid size $G$

0: Set of new collocation points $\mathcal{T}$, new RBA weights
$\left\\{\alpha^{\prime}_{0,i}\right\\}$, new grid $\mathcal{G}^{\prime}$

1: Sample a set of dense points $\mathcal{S}\subseteq\Omega$ with
$\left|\mathcal{S}\right|\gg N_{f}$

2: for each point $x_{i}\in\mathcal{S}$ do

3: Calculate $p\left(x_{i}\right)$ using (15) and (16)

4: end for

5: $p\left(x\right)\leftarrow
p\left(x\right)/\sum_{x_{i}\in\mathcal{S}}{p\left(x_{i}\right)}$

6: Sample $N_{f}$ points from $\mathcal{S}$ according to $p\left(x\right)$:

7: $\left\\{i\right\\}_{i=1}^{N_{f}}\leftarrow$ jax.random.choice based on
$p\left(x\right)$

8:
$\mathcal{T}\leftarrow\mathcal{S}\left[\left\\{i\right\\}_{i=1}^{N_{f}}\right]$

9: Re-initialize RBA weights:

10: $\bar{\alpha}\leftarrow\frac{1}{N_{f}}\sum_{i=1}^{N_{f}}{\alpha_{0,i}}$

11: for $i=1,\dots,N_{f}$ do

12: $\alpha^{\prime}_{0,i}\leftarrow\bar{\alpha}$

13: end for

14: Perform grid adaptation using $\mathcal{T},G$:

15: $\mathcal{G}^{\prime}\leftarrow$ Equation (11)

16: return $\mathcal{T}$, $\left\\{\alpha^{\prime}_{0,i}\right\\}$,
$\mathcal{G}^{\prime}$

Figure 5: Example application of RAD during the training of a PIKAN for the
Helmholtz equation. Figure 6: 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 baseline/adaptive result and its absolute error
compared to the reference solution, respectively. In each row, the
solutions/errors share the same colorbars.

Clearly, the introduction of the RAD method directly impacts the RBA
technique, as well as the grid and therefore the basis functions themselves.
Since each RBA weight corresponds to a distinct collocation point, resampling
the collocation points necessitates re-initializing the RBA weights so their
training can begin anew using (14). Consequently, choosing the optimal epochs
to perform re-sampling is crucial; frequent applications of RAD may render RBA
ineffective, while infrequent applications may not be sufficient for the
method to work effectively. As far as the re-initialization of the RBA weights
is concerned, we found that starting from the mean value of the set’s trained
weights works well, although more sophisticated transition methods can be
devised. Regarding the grid, an adaptation is required after applying RAD to
ensure it is updated based on the new set of collocation points.

###  III-D Results with Adaptive Training

Using all the aforementioned adaptive methods or combinations thereof, we
revisited the four PDEs studied in Section II to acquire new benchmarks.
First, we trained $[2,8,8,1]$ PIKANs without adaptive features or grid updates
for $10^{5}$ epochs with a learning rate of $10^{-3}$ to establish baselines
for each PDE. For their grids, we chose $G=3$ and $g_{e}=0.05$ in all cases.
Next, we trained PIKANs adaptively for the same number of epochs, using
different configurations per PDE which are detailed in the following. The
results are presented in Fig. 6.

TABLE II: Comparison of PIKAN training parameters and performance metrics with [37] for the Helmholtz and Allen-Cahn equations. PDE | Implementation | Parameters Number | Epochs | Relative L2 Error | Time (ms/epoch)
---|---|---|---|---|---
Helmholtz | jaxKAN | 1848 | $10^{5}$ | 0.176% | 18.3
| [37] | 15840 | $2\cdot 10^{5}$ | 0.16% | 7.4
Allen-Cahn | jaxKAN | 1584 | $10^{5}$ | 1.414% | 12.1
| [37] | 6720 | $1.5\cdot 10^{5}$ | 5.15% | 39.31

  * •

Diffusion Equation: Starting from a learning rate of $10^{-3}$, we scaled it
by 0.7, 0.5, 0.5 and 0.5 during epochs No. $8\cdot 10^{3}$, $2\cdot 10^{4}$,
$5\cdot 10^{4}$ and $7.5\cdot 10^{4}$, respectively. A grid adaptation was
performed every 275 epochs until epoch No. $7\cdot 10^{4}$ was reached and
grid extensions were performed during epochs No. $8\cdot 10^{3}$ and $2\cdot
10^{4}$ for $G^{\prime}=8$ and $G^{\prime}=14$, respectively. RBA was used
throughout training, but RAD was not utilized. The relative L2 error for the
baseline is equal to 0.193%, while the relative L2 error for the adaptive
PIKAN is equal to 0.021%.

  * •

Helmholtz Equation: Starting from a learning rate of $10^{-3}$, we scaled it
by 0.2, 0.5, 0.5 and 0.5 during epochs No. $2\cdot 10^{4}$, $4\cdot 10^{4}$,
$6\cdot 10^{4}$ and $8\cdot 10^{4}$, respectively. A grid adaptation was
performed every 200 epochs until epoch No. $7\cdot 10^{4}$ was reached and
grid extensions were performed during epochs No. $2\cdot 10^{4}$ and $4\cdot
10^{4}$ for $G^{\prime}=7$ and $G^{\prime}=15$, respectively. RBA was used
throughout training, but RAD was not utilized. For both the baseline and the
adaptive PIKAN, a global weight $w_{f}=0.01$ was used, following [37]. The
relative L2 error for the baseline is equal to 2.069%, while the relative L2
error for the adaptive PIKAN is equal to 0.176%.

  * •

Burgers’ Equation: Starting from a learning rate of $10^{-3}$, we scaled it by
0.7, 0.7, 0.7, 0.8, 0.8 and 0.6 during epochs No. $8\cdot 10^{3}$, $2\cdot
10^{4}$, $2.5\cdot 10^{4}$, $5\cdot 10^{4}$, $7.5\cdot 10^{4}$ and $8.5\cdot
10^{4}$, respectively. A grid adaptation was performed every 300 epochs until
epoch No. $7.5\cdot 10^{4}$ was reached and grid extensions were performed
during epochs No. $8\cdot 10^{3}$ and $2\cdot 10^{4}$ for $G^{\prime}=8$ and
$G^{\prime}=14$, respectively. RBA was used throughout training, and RAD was
performed during epochs No. $5\cdot 10^{4}$ and $7.5\cdot 10^{4}$ with
$a=c=1$. The relative L2 error for the baseline is equal to 13.504%, while the
relative L2 error for the adaptive PIKAN is equal to 2.430%.

  * •

Allen-Cahn Equation: Starting from a learning rate of $10^{-3}$, we scaled it
by 0.6, 0.8, 0.7 and 0.7 during epochs No. $1.5\cdot 10^{4}$, $2.5\cdot
10^{4}$, $5\cdot 10^{4}$ and $7.5\cdot 10^{4}$, respectively. A grid
adaptation was performed every 275 epochs until epoch No. $7\cdot 10^{4}$ was
reached and grid extensions were performed during epochs No. $8\cdot 10^{3}$
and $2\cdot 10^{4}$ for $G^{\prime}=8$ and $G^{\prime}=12$, respectively. RBA
was used throughout training, and RAD was performed every $10^{4}$ epochs,
starting from epoch No. $2\cdot 10^{4}$ until epoch No. $6\cdot 10^{4}$ with
$a=c=1$. The relative L2 error for the baseline is equal to 44.430%, while the
relative L2 error for the adaptive PIKAN is equal to 1.414%.

It is noted that the aforementioned results are not necessarily the best-
obtainable using the presented adaptive training techniques; careful fine-
tuning and extensive testing may lead to even better benchmarks. Nonetheless,
the relative L2 errors achieved in our work are already comparable to or
better than these reported in [37], albeit with PIKANs that are trained for a
shorter number of epochs and with considerably fewer parameters. Table II
compares our results to the ones obtained in [37] for the Helmholtz and Allen-
Cahn equations, as the diffusion equation was not studied in [37] and Burgers’
equation was solved using a different approach. Importantly, even though a
better GPU was used in [37] and more collocation points were utilized, our
approach is comparable even in terms of training times. This is unprecedented
for PIKANs with grid extension during training, especially taking into account
that the latencies incurred by the compilation of some functions are included
in the reported times.

##  IV Grid-Dependent Basis Functions

The experiments presented thus far demonstrate how promising KANs are as
substitutes for MLPs in PINNs, particularly when trained in an adaptive
manner. Nevertheless, their training time remains a significant disadvantage
compared to MLPs. To address this, recent studies have focused on improving
KAN efficiency by replacing the computationally expensive calculation of
B-splines with other basis functions. While this is theoretically a sound
solution, a poor choice of alternative basis functions may indeed accelerate
the training time of KANs (and PIKANs by extension), but at a severe cost to
their performance. In the following, we focus on the distinction of candidate
replacements for spline basis functions based on their dependency on a grid.
In particular, we introduce the concepts of staticity and full grid adaptivity
for basis functions and demonstrate the advantages of opting for basis
functions that are both non-static and fully adaptive to grids when training
PIKANs.

###  IV-A Staticity

Jacobi polynomials, which have been extensively used as basis functions for
KANs, including for the training of efficient and accurate PIKANs [37, 58],
fall under the category of grid-independent basis functions. They are
univariate functions which can be constructed via a Gaussian hypergeometric
function or equivalent recurrence relations. Their inherent independence from
a grid classifies Jacobi-based KANs as “static”, because it prevents them from
exploiting the dynamic grid extension method, which has been demonstrated to
enhance efficiency by progressively fine-graining their basis functions during
the training process. Notably, in Section III, we showed that this effect can
be amplified via the adaptive optimizer state transition technique, where
changing the grid’s size leads to an immediate, sharp reduction in the loss
function’s values, even if they had reached a plateau prior to the grid’s
extension.

In contrast, KANs with non-static (i.e., grid-dependent) basis functions do
not have such restrictions. In FastKAN [39], which is an implementation of
KANs utilizing radial basis functions, the dependency on a grid arises from
the choice of centers for the set of radial basis functions. ReLU-KANs [42]
employ $R_{i}\left(x\right)$ as basis functions, with

$R_{i}\left(x\right)=\left[r_{i}\cdot\text{ReLU}\left(e_{i}-x\right)\cdot\text{ReLU}\left(x-s_{i}\right)\right]^{2},$
(17)

where $r_{i}=4\left(e_{i}-s_{i}\right)^{-2}$ is a normalization constant and
$\text{ReLU}\left(x\right)=\max\left(0,x\right)$. Evidently, their dependence
on $e_{i}$ and $s_{i}$ constitutes $R_{i}\left(x\right)$ as non-static basis
functions. Finally, in Wav-KANs [41], the grid dependency varies with
different types of wavelets. For instance, if one defines a basis of Ricker
(also known as “Mexican hat”) wavelets

$\psi_{i}\left(x\right)=2\cdot\frac{1-\left(\frac{x-\mu_{i}}{\sigma_{i}}\right)^{2}}{\sqrt{3\sigma_{i}}\pi^{1/4}}\exp\left(-\frac{\left(x-\mu_{i}\right)^{2}}{2\sigma_{i}^{2}}\right),$
(18)

then the grid dependency comes from the $\mu_{i}$, $\sigma_{i}$ parameters. It
is noted that, even though the studies that introduced them employed static
grids in all of their experiments, this staticity resulted from choosing not
to perform grid updates rather than an inherent limitation of the basis
functions.

###  IV-B Full Grid Adaptivity

Figure 7: Example of spline basis functions of order $k$ = 3 calculated on a
uniform (left) and a non-uniform (right) grid.

In addition to being non-static, it is important for candidate basis functions
to also be fully adaptive to the grid in order to constitute viable
alternatives to B-Splines. Full grid adaptivity means that uniform grids
should give rise to uniform basis functions, while non-uniform grids should
produce basis functions that are denser in more dense regions and sparser in
less dense regions of the grid. Simultaneously, it is essential to ensure that
there are no regions within the grid where all basis functions are equal to
zero. Fully adaptive basis functions are essential because, as discussed in
the context of adaptive sampling of collocation points, the grids in KANs can
be designed to adapt to the data and their distribution. Therefore, fully
adaptive basis functions are, by extension, adaptive to the problem’s data
themselves.

The concept of full grid adaptivity is demonstrated in Fig. 7 for spline basis
functions of order $3$, where
$\mathcal{G}_{1}=\left[0.0,1.0,2.0,3.0,4.0\right]$ is used as an example of a
uniform grid and $\mathcal{G}_{2}=\left[0.05,0.14,0.22,1.20,11.40\right]$ is
used as an example of a non-uniform grid. The left graph of Fig. 7 depicts the
uniform basis functions constructed on the uniform grid. Conversely, the basis
functions for the non-uniform grid shown in the right graph of Fig. 7 are
denser in the region around $0.0$, reflecting the higher concentration of grid
points in this area. Notably, there is no region inside the grid where all
spline basis functions are simultaneously zero. Note that, even though $G=4$
in both cases, the total number of spline basis functions is $7$ due to the
grid augmentation that occurs prior to their construction [29]. During this
augmentation process, additional support points are appended at the ends of
the grid, so that a grid with $G$ intervals corresponds to $G+k$ spline basis
functions of order $k$.

###  IV-C The case study of ReLU-KANs

Unlike spline basis functions, the $\left\\{R_{i}\right\\}$ basis introduced
for ReLU-KANs is a clear example of non-fully adaptive basis functions. In
[42], the $R$ basis functions were constructed to be equidistant on a uniform
grid by choosing $s_{i}=\left(i-k-1\right)/G$ and $e_{i}=i/G$, therefore no
experiments were performed on non-uniform grids. However, there are numerous
PDE problems where the corresponding grids are not uniform, such as when
adding more collocation points in a specific region to impose a boundary
condition. In such cases, a grid with many intervals (large $G$) would be
required for the uniform $R$ basis functions to provide the proper support in
denser regions of the grid, thus making the switch to $R$ basis functions for
higher efficiency counterproductive. For this reason, we propose a different
construction method to make the $R$ basis functions fully adaptive to a
generally non-uniform grid, described by Algorithm 3.

Algorithm 3 Grid-Adaptive $R$ Basis Functions Construction

0: Grid $\mathcal{G}\in\mathbb{R}^{G+1}$, parameters $p,k\in\mathbb{N}$

0: $R$ basis function parameters
$\left\\{s_{i},e_{i},r_{i}\right\\}_{i=1}^{G+1}$

1: Initialize $\tilde{\mathcal{G}}\leftarrow\mathcal{G}$

2: for $i=1,\dots,p$ do

3: Append a point at each end of $\tilde{\mathcal{G}}$:

4:
$p_{s}\leftarrow\tilde{\mathcal{G}}[0]-\frac{1}{k}\left(\tilde{\mathcal{G}}[k]-\tilde{\mathcal{G}}[0]\right)$

5:
$p_{e}\leftarrow\tilde{\mathcal{G}}[G]+\frac{1}{k}\left(\tilde{\mathcal{G}}[G]-\tilde{\mathcal{G}}[G-k]\right)$

6:
$\tilde{\mathcal{G}}\leftarrow\left\\{p_{s}\right\\}\cup\tilde{\mathcal{G}}\cup\left\\{p_{e}\right\\}$

7: end for

8: for $i=p,\dots,G+p$ do

9:
$s_{i}\leftarrow\tilde{\mathcal{G}}[i]-0.5\cdot\left(\tilde{\mathcal{G}}[i+p]-\tilde{\mathcal{G}}[i-p]\right)$

10: $e_{i}\leftarrow 2\cdot\tilde{\mathcal{G}}[i]-s_{i}$

11: $r_{i}\leftarrow 4\cdot\left(e_{i}-s_{i}\right)^{-2}$

12: end for

13: return $\left\\{s_{i},e_{i},r_{i}\right\\}_{i=1}^{G+1}$

Figure 8: Example of $R$ basis functions calculated on a uniform (left) and a
non-uniform (right) grid for $p=2$ and $k=3$. Figure 9: PIKAN results with $R$
basis functions for the Helmholtz equation with staticity (first column), non-
full adaptivity (second column) and full adaptivity (third column). The upper
row depicts each experiment’s result, while the lower row shows each
experiment’s absolute error compared to the reference solution.

First, we center each $R_{i}\left(x\right)$ around a grid point,
$\mathcal{G}[i]$, by imposing $s_{i}+e_{i}=2\mathcal{G}[i]$. This ensures that
all grid points have support from the basis functions, the density of which
follows the density of the grid. Next, to ensure that there are no regions
within the grid with zero support, we set
$s_{i}=\mathcal{G}[i]-0.5\left(\mathcal{G}[i+p]-\mathcal{G}[i-p]\right)$,
where $p$ is a parameter controlling how many neighboring grid points are
taken into account when defining the coverage (width) of the basis functions.
To apply this condition for points near the grid’s edges, we perform an
augmentation technique analogous to the one used for spline basis functions.
Specifically, we extend the grid by $p$ points at both ends, recursively
taking into account the distance between each edge point and its $k$-th
neighbor, normalized by $k$. This construction ensures that $G+1$ basis
functions are obtained for a grid with $G$ intervals. The $R$ basis functions
constructed using this process with $p=2$ and $k=3$ for $\mathcal{G}_{1}$ and
$\mathcal{G}_{2}$ can be seen in the left and right graphs of Fig. 8,
respectively.

It should be noted that this construction redefines the parameter $k$, which
had no physical meaning in the original implementation of ReLU-KANs other than
providing a way to directly compare $R$ basis functions to spline basis
functions on a uniform grid, both qualitatively and quantitatively.
Importantly, this construction is by no means the only method to make $R$
basis functions fully adaptive. Nevertheless, it is intuitive and can serve as
the basis for a broader methodology of constructing non-static, fully adaptive
basis functions. For instance, extending these ideas to radial basis functions
or Ricker wavelets is straightforward: one may choose centers for the former
and $\mu_{i}$ for the latter which lie on the grid points’ coordinates and
then adjust their spread to ensure full coverage within the grid.

We implemented ReLU-KANs in jaxKAN based on this construction, utilizing just-
in-time compilation and vectorized operations for improved computational
efficiency. To showcase the importance of non-staticity and full adaptivity,
we trained three $[2,8,8,1]$ PIKANs for $10^{5}$ epochs using the $R$ basis
functions with $k=p=2$ to solve the Helmholtz equation. For the first PIKAN,
we imposed staticity by selecting a fixed set of basis functions and choosing
not to update the grid during training ($G=3$, $g_{e}=1.0$). For the second
PIKAN, we performed grid updates ($G=3\to 6\to 12$) but only utilized uniform
grids to impose non-full adaptivity ($g_{e}=1.0$). For the third PIKAN, we
performed grid updates ($G=3\to 6\to 12$) and used grids adaptive to the data
($g_{e}=0.05$), in order to investigate their impact in our fully adaptive
implementation of $R$ basis functions. To ensure a fair comparison, all other
parameters of the PIKANs were kept constant across all three experiments, with
the exception of the learning rate: in the first case a constant learning rate
was used, since no extension was performed. For the second PIKAN, the learning
rate was scaled by 0.5, 0.5, 0.4, 0.7 and 0.7 during epochs No. $8\cdot
10^{3}$, $1.5\cdot 10^{4}$, $3\cdot 10^{4}$, $5\cdot 10^{4}$ and $7\cdot
10^{4}$, respectively. In the final case, the learning rate was scaled by 0.6,
0.8, 0.7 and 0.7 durings epochs No. $2\cdot 10^{4}$, $3.5\cdot 10^{4}$,
$5\cdot 10^{4}$ and $7\cdot 10^{4}$, respectively. The reason why a different
scheduler was used in these two cases is because the loss function of the
second PIKAN would produce NaN values if a scheduler identical to the one used
in the third PIKAN was selected.

The solutions for each PIKAN along with their absolute errors from the
reference solution are depicted in Fig. 9. The corresponding relative L2
errors are 24.191% for the static case, 6.761% for the non-fully adaptive one
and 5.552% for the fully adaptive PIKAN, proving that non staticity and full
grid adaptivity of their basis functions are important when training PIKANs.
Notably, none of these PIKANs seem to outperform even the baseline set in
Section III using spline basis functions. However, the average training time
per epoch was 3.4 ms for ReLUKAN-based PIKANs with grid extensions, indicating
a tradeoff between performance and computational efficiency. Ultimately, the
choice of basis functions for a PIKAN will depend on whether the user favors
speed or accuracy in their application.

##  V Summary & Future Work

In the present study, we attempted to set the foundations for an adaptive
training scheme for PIKANs with the purpose of improving their accuracy and
optimizing their performance. To this end, we introduced an adaptive technique
to smoothen the transition process after grid extensions, based on regular
grid adaptations and a reconfiguration of the optimizer’s state, without
completely resetting it. Additionally, we adopted the RBA technique for
adaptively re-weighting the loss function’s terms, as well as the RDA
technique to adaptively re-sample collocation points during training, both
based on the residuals of the loss function. For the latter, we proposed an
implementation that does not hinder the progress of RBA and conforms to the
grid’ adaptivity.

All of these methods were implemented in JAX within our newly introduced
jaxKAN framework. PIKANs trained within our framework achieved training times
which were 2 orders of magnitude faster than the PIKANs trained using pykan.
Moreover, the performance of PIKANs which we trained adaptively to solve four
different PDEs was on par with or superior to other types of PIKANs which were
trained for larger numbers of epochs, with deeper and/or wider architectures.
Importantly, the average training times were also comparable.

Finally, we proposed a general methodology for choosing alternative basis
functions for PIKANs in cases when splines are too computationally
inefficient. By defining the concept of staticity and full grid adaptivity, we
advocated for utilizing basis functions that are able to cover the entire grid
and provide support in specific regions based on the distribution of their
inputs. Using ReLU-KAN-based PIKANs as a case study, we demonstrated that
fully adaptive basis functions outperform static ones by comparing their
accuracies for the solution of the Helmholtz PDE.

As far as future research in PIKANs is concerned, the field is ripe with
numerous opportunities and new avenues to explore. Regarding the newly
introduced jaxKAN framework, we intend to enrich it with additional KAN models
and their corresponding basis functions, based on the attributes highlighted
in Section IV. Moreover, studies to make the basis functions’ grid dependency
as efficient as possible are necessary. Ensuring that the basis functions
cover the entire grid, with more support in denser areas, is a good starting
point; however, further research is required, for instance, regarding their
amplitudes or their desired overlap in dense/sparse regions. Another direction
for future work is utilizing optimizers other than Adam with the adaptive
state transition technique between grid updates. Finally, the overall adaptive
training of PIKANs requires additional benchmarks, ideally followed by a
strict mathematical formulation to understand how techniques like adaptive
sampling of collocation points affect KANs differently from MLPs due to their
inherent grid dependency. Generally, we believe that the end goal is to
establish a set of mathematically-backed best practices for training PIKANs,
ultimately creating a robust and unified pipeline for solving any type of
differential equation.

##  Appendix A Studied Partial Differential Equations

The purpose of this Appendix is to present the four PDEs used in the
experiments of this paper: the diffusion equation, the Helmholtz equation,
Burgers’ equation and the Allen-Cahn equation. Before presenting the PDEs with
their boundary conditions, we note that for all PIKANs trained in the
experiments presented in our work, a total of $N_{f}=2^{12}$ collocation
points were used for the PDE and $N_{b}=2^{6}$ collocation points were used to
impose each boundary condition. The reason why we opted for powers of two is
because we sampled the collocation points from the Sobol sequence, a process
mentioned in [26]. Additionally, the performance of the trained models was
evaluated in terms of the relative L2 error, defined as

$L^{2}_{\text{rel}}=\frac{\lVert
u_{r}\left(x\right)-u\left(x;\theta\right)\rVert}{\lVert
u_{r}\left(x\right)\rVert},$

where $u\left(x;\theta\right)$ is the PIKAN, $u_{r}\left(x\right)$ is the
reference solution and $\lVert\cdot\rVert$ denotes the L2 norm. The reference
solution for each PDE is depicted in Fig. 10.

Figure 10: Reference solutions for each of the studied PDEs.

###  A-A Diffusion equation

The 1+1-dimensional diffusion equation is:

$\frac{\partial u}{\partial t}-\frac{\partial^{2}u}{\partial
x^{2}}=f\left(t,x\right),$

where $f\left(t,x\right)$ is the source term which was chosen equal to

$f\left(t,x\right)=\left(\pi^{2}-1\right)\exp\left(-t\right)\sin\left(\pi
x\right)$

for the purposes of this paper. Additionally,
$\Omega=\left[0,1\right]\times\left[0,1\right]$ is the chosen domain, with the
following boundary conditions:

$\displaystyle u\left(t=0,x\right)$ $\displaystyle=\sin\left(\pi x\right),$
$\displaystyle u\left(t,x=0\right)$ $\displaystyle=u\left(t,x=1\right)=0.$

The exact solution of this diffusion equation is

$u\left(t,x\right)=\sin\left(\pi x\right)\exp\left(-t\right),$

and is depicted in the upper left plot of Fig. 10.

###  A-B Helmholtz equation

The 2-dimensional Helmholtz equation is defined as:

$\frac{\partial^{2}u}{\partial x^{2}}+\frac{\partial^{2}u}{\partial
y^{2}}+k^{2}u=f\left(x,y\right).$

For our experiments we chose $k=1$ and

$f\left(x,y\right)=\left(1-17\pi^{2}\right)\sin\left(\pi
x\right)\sin\left(4\pi y\right),$

with $\Omega=\left[-1,1\right]\times\left[-1,1\right]$ and boundary conditions

$\displaystyle u\left(x=-1,y\right)$ $\displaystyle=u\left(x=1,y\right)=0,$
$\displaystyle u\left(x,y=-1\right)$ $\displaystyle=u\left(x,y=1\right)=0.$

This PDE also has an exact solution, which is

$u\left(t,x\right)=\sin\left(\pi x\right)\sin\left(4\pi y\right)$

and is depicted in the upper right plot of Fig. 10.

###  A-C Burgers’ equation

As far as Burgers’ equation is concerned, it is defined as:

$\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial
x}-\nu\frac{\partial^{2}u}{\partial x^{2}}=0$

in the $\Omega=\left[-1,1\right]\times\left[0,1\right]$ domain. For the
boundary conditions

$\displaystyle u\left(t=0,x\right)$ $\displaystyle=-\sin\left(\pi x\right),$
$\displaystyle u\left(t,x=-1\right)$ $\displaystyle=u\left(t,x=1\right)=0,$

the solution of Burgers’ equation does not have an analytical expression. For
this reason, we used the reference solution which was also used in [26],
depicted in the lower left plot of Fig. 10 for $\nu=0.01/\pi$.

###  A-D Allen-Cahn equation

Finally, the 1+1-dimensional nonlinear Allen-Cahn equation is expressed as:

$\frac{\partial u}{\partial t}-D\frac{\partial^{2}u}{\partial
x^{2}}+5\left(u^{3}-u\right)=0,$

where $\Omega=\left[0,1\right]\times\left[-1,1\right]$. For our experiments we
chose $D=0.001$, with the following boundary conditions:

$\displaystyle u\left(t=0,x\right)$ $\displaystyle=x^{2}\cos\left(\pi
x\right),$ $\displaystyle u\left(t,x=-1\right)$
$\displaystyle=u\left(t,x=1\right)=-1.$

Similar to Burgers’ equation, this PDE does not have an analytical solution.
Again, we used the reference solution from [26], depicted in the lower right
plot of Fig. 10.

##  Appendix B Adaptive State Transition in Function Learning

In this Appendix, we present an application of the adaptive state transition
technique introduced in Section III, for KANs that are trained to learn
functions. This was also one of their main tasks in the work that introduced
them [29]. To this end, we trained two $[4,5,2,1]$ KANs with $g_{e}=0.02$ and
spline basis functions of order $k=3$ to learn the multivariate function

$\displaystyle f\left(x_{1},x_{2},x_{3},x_{4}\right)=\exp$
$\displaystyle\left[\frac{1}{2}\sin\left(\pi x_{1}^{2}+\pi
x_{2}^{2}\right)+\right.$ $\displaystyle\left.~{}~{}\frac{1}{2}\sin\left(\pi
x_{3}^{2}+\pi x_{4}^{2}\right)\right],$

by sampling 3000 points from it. The initial size of each KAN’s grid was $G=3$
and was extended to $G\to 6\to 10\to 24$ during epoch No. 200, 400 and 600,
respectively. The baseline KAN was trained for a total of 800 epochs using a
uniform learning rate of $0.02$ and no adaptive state transition for the
optimizer. The other KAN was also trained for a total of 800 epochs, however
with the adaptive state transition technique and a learning rate which started
from $0.02$ and was scaled by a factor of 0.5, 0.2 and 0.5, during epoch No.
200, 400 and 600, respectively. The results for both KANs are depicted through
the values of their loss functions during training in Fig. 11.

Figure 11: Training a KAN without (top) and with (bottom) the adaptive
optimizer state transition technique.

This example showcases how the proposed technique reverses the peaks shown in
the loss function’s values during grid extensions: the sharp increases in the
baseline’s training are turned into sharp decreases for the training of the
adaptive KAN. Additionally, the technique leads to overall better fits for the
trained network, with the adaptive KAN achieving losses that are lower than
those of the baseline KAN by approximately 1 order of magnitude.

## Reproducibility

The generated data and the code used to produce the results of all experiments
presented herein can be found in the jaxKAN GitHub repository in
https://github.com/srigas/jaxKAN under the “Paper Results” directory.

## Acknowledgment

S. R. thanks P. Angelopoulou for thoroughly reading the original draft to
point out typos and inconsistencies.

## References

  * [1] M. Raissi, P. Perdikaris, and G. E. Karniadakis, “Physics informed deep learning (part I): data-driven solutions of nonlinear partial differential equations,” _arXiv preprint_ , Nov. 2017. DOI: 10.48550/arXiv.1711.10561
  * [2] M. Raissi, P. Perdikaris, and G. E. Karniadakis, “Physics informed deep learning (part II): data-driven solutions of nonlinear partial differential equations,” _arXiv preprint_ , Nov. 2017. DOI: 10.48550/arXiv.1711.10566
  * [3] M. Raissi, P. Perdikaris, and G. E. Karniadakis, “Physics-informed neural networks: a deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations,” _J. Comput. Phys._ , vol. 378, pp. 686–707, Feb. 2019. DOI: 10.1016/j.jcp.2018.10.045
  * [4] Z. Mao, A. D. Jagtap, and G. E. Karniadakis, “Physics-informed neural networks for high-speed flows,” _Comput. Methods Appl. Mech. Eng._ , vol. 360, p. 112789, Mar. 2020. DOI: 10.1016/j.cma.2019.112789
  * [5] L. Sun, H. Gao, S. Pan, and J.-X. Wang, “Surrogate modeling for fluid flows based on physics-constrained deep learning without simulation data,” _Comput. Methods Appl. Mech. Eng._ , vol. 361, p. 112732, Apr. 2020. DOI: 10.1016/j.cma.2019.112732
  * [6] X. Jin, S. Cai, H. Li, and G. E. Karniadakis, “NSFnets (Navier-Stokes flow nets): Physics-informed neural networks for the incompressible Navier-Stokes equations,” _J. Comput. Phys._ , vol. 426, p. 109951, Feb. 2021. DOI: 10.1016/j.jcp.2020.109951
  * [7] F. Sahli Costabal, Y. Yang, P. Perdikaris, D. E. Hurtado, and E. Kuhl, “Physics-informed neural networks for cardiac activation mapping,” _Front. Phys._ , vol. 8, Feb. 2020. DOI: 10.3389/fphy.2020.00042
  * [8] K. Sel, A. Mohammadi, R. I. Pettigrew, and R. Jafari, “Physics-informed neural networks for modeling physiological time series for cuffless blood pressure estimation,” _npj Digit. Med._ , vol. 6, no. 1, pp. 1–15, Jun. 2023. DOI: 10.1038/s41746-023-00853-4
  * [9] G. S. Misyris, A. Venzke, and S. Chatzivasileiadis, “Physics-informed neural networks for power systems,” in _2020 IEEE Power & Energy Society General Meeting (PESGM)_, Aug. 2020, pp. 1–5. DOI: 10.1109/PESGM41954.2020.9282004
  * [10] R. Nellikkath and S. Chatzivasileiadis, “Physics-informed neural networks for AC optimal power flow,” _Electric Power Syst. Res._ , vol. 212, p. 108412, Nov. 2022. DOI: 10.1016/j.epsr.2022.108412
  * [11] S. P. Moschou, E. Hicks, R. Y. Parekh, D. Mathew, S. Majumdar, and N. Vlahakis, “Physics-informed neural networks for modeling astrophysical shocks,” _Mach. Learn.: Sci. Technol._ , vol. 4, no. 3, p. 035032, Sep. 2023. DOI: 10.1088/2632-2153/acf116
  * [12] H. Jin, M. Mattheakis, and P. Protopapas, “Physics-informed neural networks for quantum eigenvalue problems,” in _International Joint Conference on Neural Networks (IJCNN)_ , Jul. 2022, pp. 1–8. DOI: 10.1109/IJCNN55064.2022.9891944
  * [13] S. Wang, Y. Teng, and P. Perdikaris, “Understanding and mitigating gradient flow pathologies in physics-informed neural networks,” _SIAM J. Sci. Comput._ , vol. 43, no. 5, pp. A3055–A3081, Jan. 2021. DOI: 10.1137/20M1318043
  * [14] N. Rahaman  _et al._ , “On the spectral bias of neural networks,” in _Proceedings of the 36th International Conference on Machine Learning, PMLR_ , May 2019, pp. 5301–5310. [Online]. Available: https://proceedings.mlr.press/v97/rahaman19a.html
  * [15] Y. Zhu, N. Zabaras, P.-S. Koutsourelakis, and P. Perdikaris, “Physics-constrained deep learning for high-dimensional surrogate modeling and uncertainty quantification without labeled data,” _J. Comput. Phys._ , vol. 394, pp. 56–81, Oct. 2019. DOI: 10.1016/j.jcp.2019.05.024
  * [16] R. Zhang, Y. Liu, and H. Sun, “Physics-informed multi-LSTM networks for metamodeling of nonlinear structures,” _Comput. Methods Appl. Mech. Eng._ , vol. 369, p. 113226, Sep. 2020. DOI: 10.1016/j.cma.2020.113226
  * [17] B. Moseley, A. Markham, and T. Nissen-Meyer, “Finite basis physics-informed neural networks (FBPINNs): a scalable domain decomposition approach for solving differential equations,” _Adv. Comput. Math._ , vol. 49, no. 4, p. 62, Jul. 2023. DOI: 10.1007/s10444-023-10065-9
  * [18] J. Hou, Y. Li, and S. Ying, “Enhancing PINNs for solving PDEs via adaptive collocation point movement and adaptive loss weighting,” _Nonlinear Dyn._ , vol. 111, no. 16, pp. 15233–15261, Aug. 2023. DOI: 10.1007/s11071-023-08654-w
  * [19] S. Wang, X. Yu, and P. Perdikaris, “When and why PINNs fail to train: a neural tangent kernel perspective,” _J. Comput. Phys._ , vol. 449, p. 110768, Jan. 2022. DOI: 10.1016/j.jcp.2021.110768
  * [20] Z. Xiang, W. Peng, X. Liu, and W. Yao, “Self-adaptive loss balanced physics-informed neural networks,” _Neurocomputing_ , vol. 496, pp. 11–34, Jul. 2022. DOI: 10.1016/j.neucom.2022.05.015
  * [21] R. Bischof and M. Kraus, “Multi-objective loss balancing for physics-informed deep learning,” _arXiv preprint_ , Nov. 2022. DOI: 10.48550/arXiv.2110.09813
  * [22] S. J. Anagnostopoulos, J. D. Toscano, N. Stergiopulos, and G. E. Karniadakis, “Residual-based attention in physics-informed neural networks,” _Comput. Methods Appl. Mech. Eng._ , vol. 421, p. 116805, Mar. 2024. DOI: 10.1016/j.cma.2024.116805.
  * [23] J. Yu, L. Lu, X. Meng, and G. E. Karniadakis, “Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems,” _Comput. Methods Appl. Mech. Eng._ , vol. 393, p. 114823, Apr. 2022. DOI: 10.1016/j.cma.2022.114823
  * [24] H. Son, J. W. Jang, W. J. Han, and H. J. Hwang, “Sobolev training for physics-informed neural networks,” _Commun. Math. Sci._ , vol. 21, no. 6, pp. 1679–1705, 2023. DOI: 10.4310/CMS.2023.v21.n6.a11
  * [25] M. A. Nabian, R. J. Gladstone, and H. Meidani, “Efficient training of physics‐informed neural networks via importance sampling,” _Comput.-Aided Civ. Infrastruct. Eng._ , vol. 36, no. 8, pp. 962–977, Aug. 2021. DOI: 10.1111/mice.12685
  * [26] C. Wu, M. Zhu, Q. Tan, Y. Kartha, and L. Lu, “A comprehensive study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks,” _Comput. Methods Appl. Mech. Eng._ , vol. 403, p. 115671, Jan. 2023. DOI: 10.1016/j.cma.2022.115671
  * [27] A. D. Jagtap, K. Kawaguchi, and G. E. Karniadakis, “Adaptive activation functions accelerate convergence in deep and physics-informed neural networks,” _J. Comput. Phys._ , vol. 404, p. 109136, Mar. 2020. DOI: 10.1016/j.jcp.2019.109136
  * [28] R. Gnanasambandam, B. Shen, J. Chung, X. Yue, and Z. J. Kong, “Self-scalable tanh (Stan): faster convergence and better generalization in physics-informed neural networks.” _arXiv preprint_ , Apr. 2022. DOI: 10.48550/arXiv.2204.12589
  * [29] Z. Liu  _et al._ , “KAN: Kolmogorov-Arnold networks,” _arXiv preprint_ , May 2024. DOI: 10.48550/arXiv.2404.19756
  * [30] R. Genet and H. Inzirillo, “TKAN: temporal Kolmogorov-Arnold networks,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2405.07344
  * [31] C. J. Vaca-Rubio, L. Blanco, R. Pereira, and M. Caus, “Kolmogorov-Arnold networks (KANs) for time series analysis,” _arXiv preprint_ , May 2024. DOI: 10.48550/arXiv.2405.08790
  * [32] R. Genet and H. Inzirillo, “A temporal Kolmogorov-Arnold transformer for time series forecasting,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.02486
  * [33] M. Cheon, “Kolmogorov-Arnold network for satellite image classification in remote sensing,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.00600
  * [34] B. Azam and N. Akhtar, “Suitability of KANs for computer vision: a preliminary investigation,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.09087
  * [35] C. Li, X. Liu, W. Li, C. Wang, H. Liu, and Y. Yuan, “U-KAN makes strong backbone for medical image segmentation and generation,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.02918
  * [36] M. Liu, S. Bian, B. Zhou, and P. Lukowicz, “iKAN: global incremental learning with KAN for human activity recognition across heterogeneous datasets,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.01646
  * [37] K. Shukla, J. D. Toscano, Z. Wang, Z. Zou, and G. E. Karniadakis, “A comprehensive and FAIR comparison between MLP and KAN representations for differential equations and operator networks,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.02917
  * [38] H. Wu, H. Luo, Y. Ma, J. Wang, and M. Long, “RoPINN: region optimized physics-informed neural networks,” _arXiv preprint_ , May 2024. DOI: 10.48550/arXiv.2405.14369
  * [39] Z. Li, “Kolmogorov-Arnold networks are radial basis function networks,” _arXiv preprint_ , May 2024. DOI: 10.48550/arXiv.2405.06721
  * [40] S. S. Sidharth and R. Gokul, “Chebyshev polynomial-based Kolmogorov-Arnold networks: an efficient architecture for nonlinear function approximation,” _arXiv preprint_ , May 2024. DOI: 10.48550/arXiv.2405.07200
  * [41] Z. Bozorgasl and H. Chen, “Wav-KAN: wavelet Kolmogorov-Arnold networks,” _arXiv preprint_ , May 2024. DOI: 10.48550/arXiv.2405.12832
  * [42] Q. Qiu, T. Zhu, H. Gong, L. Chen, and H. Ning, “ReLU-KAN: new Kolmogorov-Arnold networks that only need matrix addition, dot multiplication, and ReLU,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.02075
  * [43] J. Bradbury  _et al._ , “JAX: composable transformations of Python+NumPy programs.” Accessed: Jun. 10 2024. [Online]. Available: https://github.com/google/jax
  * [44] J. Heek  _et al._ , “Flax: a neural network library and ecosystem for JAX.” Accessed: Jun. 10 2024. [Online]. Available: https://github.com/google/flax
  * [45] S. Wang, S. Sankaran, H. Wang, and P. Perdikaris, “An expert’s guide to training physics-informed neural networks,” _arXiv preprint_ , Aug. 2023. DOI: 10.48550/arXiv.2308.08468
  * [46] A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind, “Automatic differentiation in machine learning: a survey,” _J. Mach. Learn. Res._ , vol. 18, no. 153, pp. 1–43, 2018. [Online]. Available: http://jmlr.org/papers/v18/17-468.html
  * [47] A. Paszke  _et al._ , “Automatic differentiation in PyTorch,” Oct. 2017. [Online]. Available: https://openreview.net/pdf?id=BJJsrmfCZ
  * [48] M. Abadi  _et al._ , “TensorFlow: large-scale machine learning on heterogeneous systems,” 2015. [Online]. Available: https://www.tensorflow.org/
  * [49] S. Rigas and M. Papachristou, “jaxKAN: a JAX-based implementation of Kolmogorov-Arnold networks,” 2014. Accessed: Jun. 11 2024. [Online]. Available: https://pypi.org/project/jaxkan/
  * [50] A. A. Howard, B. Jacob, S. H. Murphy, A. Heinlein, and P. Stinis, “Finite basis Kolmogorov-Arnold networks: domain decomposition for data-driven and physics-informed problems,” _arXiv preprint_ , June 2024. DOI: 10.48550/arXiv.2312.08098
  * [51] C. de Boor, “On calculating with B-splines,” _J. Approx. Theory_ , vol. 6, no. 1, pp. 50–62, Jul. 1972. DOI: 10.1016/0021-9045(72)90080-9
  * [52] M. G. Cox, “The numerical evaluation of B-splines,” _IMA J. Appl. Math._ , vol. 10, no. 2, pp. 134–149, Oct. 1972. DOI: 10.1093/imamat/10.2.134
  * [53] D. C. Liu and J. Nocedal, “On the limited memory BFGS method for large scale optimization,” _Math. Program._ , vol. 45, no. 1, pp. 503–528, Aug. 1989. DOI: 10.1007/BF01589116
  * [54] D. Kingma and J. Ba, “Adam: A Method for stochastic optimization,” in _International Conference on Learning Representations (ICLR)_ , May 2015. [Online]. Available: https://arxiv.org/abs/1412.6980
  * [55] DeepMind, “The DeepMind JAX ecosystem” Accessed: Jun. 13 2024. [Online]. Available: https://github.com/google-deepmind
  * [56] H. Cao, “An efficient implementation of Kolmogorov-Arnold network.” Accessed: Jun. 13 2024. [Online]. Available: https://github.com/Blealtan/efficient-kan
  * [57] T. Dozat, “Incorporating Nesterov momentum into Adam,” 2016. [Online]. Available: https://openreview.net/pdf/OM0jvwB8jIp57ZJjtNEZ.pdf
  * [58] A. A. Aghaei, “fKAN: fractional Kolmogorov-Arnold networks with trainable Jacobi basis functions,” _arXiv preprint_ , Jun. 2024. DOI: 10.48550/arXiv.2406.07456