A comprehensive and FAIR comparison between MLP and KAN representations for differential equations and operator networks

Khemraj Shukla; Juan Diego Toscano; Zhicheng Wang; Zongren Zou; George Em Karniadakis

Source-linked AI summary

A comprehensive and FAIR comparison between MLP and KAN representations for differential equations and operator networks

Khemraj Shukla, Juan Diego Toscano, Zhicheng Wang, Zongren Zou, George Em Karniadakis

arXiv:2406.02917v1Jun 5, 2024cs.LGphysics.comp-ph

TL;DR

The paper asks whether KAN representations can provide effective alternatives to MLP-based PINNs and DeepONets for differential equations and operator learning. It develops and evaluates PIKANs and DeepOKANs, compares them with PINNs and DeepONets, and analyzes their loss landscapes and learning dynamics. Modified low-order polynomial variants achieve comparable accuracy in supported tasks, but robustness and training-cost limitations remain.

Problem
The study investigates whether KAN-based representations can effectively address scientific machine-learning problems alongside MLP-based PINNs and DeepONets.
Method
The paper develops PIKANs and DeepOKANs, benchmarks them against PINNs and DeepONets, and analyzes training dynamics using information bottleneck theory.
Results
Modified Chebyshev- and other low-order polynomial KAN variants achieve performance comparable to PINNs and DeepONets across the reported tasks.
Takeaways & Limitations
KAN representations can be competitive for scientific machine learning, while modified formulations are more practical than the original B-spline formulation in the reported comparisons.
Takeaways & Limitations
Higher polynomial orders, greater depth, high-Reynolds-number flows, and some random seeds can produce unstable or divergent training.

Abstract

from arXiv · show

Kolmogorov-Arnold Networks (KANs) were recently introduced as an alternative representation model to MLP. Herein, we employ KANs to construct physics-informed machine learning models (PIKANs) and deep operator models (DeepOKANs) for solving differential equations for forward and inverse problems. In particular, we compare them with physics-informed neural networks (PINNs) and deep operator networks (DeepONets), which are based on the standard MLP representation. We find that although the original KANs based on the B-splines parameterization lack accuracy and efficiency, modified versions based on low-order orthogonal polynomials have comparable performance to PINNs and DeepONet although they still lack robustness as they may diverge for different random seeds or higher order orthogonal polynomials. We visualize their corresponding loss landscapes and analyze their learning dynamics using information bottleneck theory. Our study follows the FAIR principles so that other researchers can use our benchmarks to further advance this emerging topic.

1. Introduction

The paper evaluates KANs as alternatives to MLP-based models for differential equations and operator learning, addressing accuracy, efficiency, stability, and learning dynamics. It develops PIKANs and DeepOKANs and compares them with PINNs and DeepONets across broad benchmarks.

MLPs are widely used universal approximators but can suffer from limited interpretability, overfitting, gradient problems, and scalability issues.
KANs were proposed as potentially more accurate and interpretable alternatives to MLPs, but their B-spline activations increase computational cost.
Earlier KAN studies reported strong results on differential equations and operator learning, but used shallow networks and simplified problems.
The study develops PIKANs and DeepOKANs and compares them with PINNs and DeepONets across six differential-equation benchmarks and two operator-learning tasks.
The comparisons incorporate residual-based attention and eddy-viscosity formulations to support fair evaluation on more complex problems.
The study analyzes PIKAN stability, sensitivity to polynomial order and depth, loss landscapes, and training dynamics through information bottleneck theory.
The paper examines three learning stages in PIKAN training and relates them to stages previously identified for PINNs.

2. Problem Formulation and representation models

The paper distinguishes neural differential equations from neural operators, then introduces MLP and KAN representations used to construct PINNs, DeepONets, and related KAN-based models.

2. Problem Formulation and representation models: A neural differential equation approximates one specific ODE/PDE solution uθ, while a neural operator learns a function-to-function mapping such as f to u.
2.1. Physics-informed neural networks (PINNs): PINNs encode governing differential and boundary operators in a physics-informed loss computed through automatic differentiation.
2.1. Physics-informed neural networks (PINNs): Residual-based attention scales local loss contributions so point-wise residuals are not overlooked in cumulative training losses.
2.2. Neural operators (NOs): Neural operators approximate solution operators from paired input and output function data and can serve as fast forward solvers or physics encoders.
2.3.1. Multilayer Perceptron (MLP): An MLP applies successive affine transformations and activation functions, with each layer output feeding the next until the final output y.
2.3.2. Kolmogorov-Arnold networks (KANs): KANs represent multivariate functions through compositions of univariate activation functions, whose specific forms define different KAN architectures.
2.3.2. Kolmogorov-Arnold networks (KANs): Vanilla PIKANs use weighted B-spline activations, whereas RBF, wavelet, and Jacobi variants replace or modify the univariate basis functions.
2.3.2. Kolmogorov-Arnold networks (KANs): cPIKAN combines physics-informed learning with Chebyshev KANs and avoids grid points, while KAN parameter counts depend on depth, width, grid size, and polynomial order.

3. Computational experiments

The computational experiments compare MLP- and KAN-based architectures across differential equations, operator regression, and noisy-data settings.

The experiments evaluate MLP and KAN architectures on steady and unsteady partial differential equations.
The study tests operator regression in both low- and high-dimensional regimes.
The experiments also assess Bayesian KAN-based approaches for solving PDEs with noisy data.

3.1. Approximation of a discontinuous and oscillatory function

The study tests KAN, Chebyshev-KAN, modified Chebyshev-KAN, and MLP representations on a discontinuous function containing low- and high-frequency modes.

The target function combines a discontinuity with low- and high-frequency modes to probe spectral bias and approximation robustness.
The comparison implements KAN, Chebyshev-KAN, modified Chebyshev-KAN, and MLP models using two hidden layers with 40 neurons each.
0.29%, 0.33%, 0.79%, and 0.81% are the relative l2-errors for KAN-I, KAN-II, modified Chebyshev-KAN, and MLP, respectively.
KAN-I and KAN-II achieved nearly identical accuracy despite KAN-II using an order of magnitude fewer parameters.
All four architectures captured the reference Fourier spectrum across the tested frequencies.
KAN-I and KAN-II converged more steeply than modified Chebyshev-KAN and MLP architectures.
The table reports network hyperparameters and GPU time per iteration for the approximation comparison.

3.2. Structure preserving Dynamical System: Hamiltonian neural network (HNN) vs Hamiltonian Chebyshev-KAN (HcKAN)

For the ideal mass-spring system, HcKAN replaces the HNN’s MLP layer but requires architectural modification for stability and generalizes less strongly than HNN. Increasing the training dataset substantially improves HcKAN’s test loss.

Architecture: HcKAN replaces the MLP layer in HNN with a Chebyshev-KAN layer for forecasting the dynamical system’s state.The models are trained by minimizing the stated loss function.
Stability: Chebyshev-KAN required a modified architecture and shallower network to achieve stable training.The unmodified setup was not stable enough for the experiment.
Generalization: HcKAN trained with 40 samples had weaker extrapolation and larger generalization error than HNN.Its test loss remained stagnant at a high value, indicating poor generalization to unseen data.
Data sensitivity: A slightly larger dataset reduced HcKAN’s test loss by several orders of magnitude.The improvement is shown in the convergence histories for the larger-data settings.

3.3. Helmholtz Equation

The Helmholtz experiments compare PINN and PIKAN variants across parameter, computation-time, and complexity settings. cPIKAN-based models can match or exceed PINN accuracy, but deeper or higher-order configurations become unstable and sensitive to initialization and precision.

Experimental design: The models are compared using Relative L2 error and training time, under parameter-matched, computation-time-matched, and higher-complexity settings.The higher-complexity case uses wave numbers a1 = a2 = 6.
Parameter-based analysis: 0.354% relative L2 error is achieved by cPIKAN+RBA for the (a1, a2) = (1, 4) case, with performance better than vanilla PINN and comparable final error to PINN+RBA.PIKAN is significantly slower because it does not benefit from GPU parallelization; multigrid PIKAN is faster than PIKAN but uses twice as many parameters in its final stage.
Computation time-based analysis: 0.160% relative L2 error is achieved by cPIKAN+RBA in the computation-time analysis, marginally outperforming PINN with and without RBA.The analysis uses deeper networks and 100 × 100 collocation points while roughly matching computational time.
Complexity-based analysis: 0.381% relative L2 error is achieved by cPIKAN+RBA for the higher-wave-number case, while vanilla cPIKAN does not converge.The higher wave number induces steeper PDE-residual gradients, making approximation more difficult.
Sensitivity analysis: Increasing polynomial order or network depth can improve cPIKAN performance but also cause instability, undefined losses, and sensitivity to initialization.For this example, k > 6 or nl > 5 causes the loss to become undefined after several iterations; single precision worsens instability at higher k.

3.4. Navier-Stokes equation

The study compares PINNs and polynomial-based PIKANs for steady lid-driven cavity flow at Re = 400 and Re = 2000 under matched training conditions. PIKANs can match PINN accuracy at Re = 400 with fewer parameters, but high-Re training is unstable without stabilization.

Experimental setup: The benchmark compares Chebyshev, Jacobi, Legendre, and Hermite PIKANs with PINNs at Re = 400 and Re = 2000 using matched residual points, epochs, and optimizer settings.The only intended model difference is the representation: MLP for PINN versus polynomial-based KAN for PIKAN.
Re = 400: At Re = 400, PIKANs achieve accuracy comparable to PINNs with substantially fewer trainable parameters, but each PIKAN iteration takes four times longer.Among the variants, cPIKAN is the most promising for inference accuracy and GPU training time.
Stabilization: RBA improves inference accuracy for both PINNs and PIKANs at the same computation time, while also accelerating their training.RBA dynamically and locally adjusts loss weights.
Re = 2000: At Re = 2000, vanilla PINN reaches l2-relative error below 20%, whereas vanilla cPIKAN exceeds 100% after 40,000 epochs.With EVM, both methods reduce relative error below 7% after the same number of epochs.
Flow structure: At Re = 2000, cPIKAN streamlines differ completely from the reference solution, whereas at Re = 400 both models reproduce the small corner eddies.At Re = 400, cPIKAN slightly improves the right-bottom streamline structure relative to PINN.
Conclusion: Overall, Jacobi-based PIKANs can match PINN accuracy but may cost several times more to train and suffer unstable high-Re training without EVM or RBA.These stabilization methods can return PIKANs to the correct training trajectory.

3.5. Allen-Cahn equation

The Allen-Cahn experiment evaluates PIKAN, cPIKAN, cPIKAN with RBA, and PINN with RBA against a reference solution. PIKAN performs poorly, while both cPIKAN variants achieve comparable errors under the reported training setup.

Setup: The experiment solves the nonlinear Allen-Cahn equation using PIKAN, cPIKAN, cPIKAN with RBA, and PINN with RBA.The solutions are obtained by minimizing a loss balancing initial-condition, boundary-condition, and PDE-residual terms.
Results: After 150,000 single-batch Adam iterations, PIKAN has 58.39% relative l2-error against the reference solution.The reported learning rate is 5×10^-4.
Results: cPIKAN and cPIKAN with RBA achieve relative l2-errors of 5.15% and 5.65%, respectively, indicating nearly similar accuracy.The reference solution is computed using a spectral element method.

3.6. Reaction-diffusion equation

The reaction-diffusion experiment applies Bayesian cPIKAN and Bayesian PINN to an inverse problem with noisy observations. Both methods infer u and f similarly, while B-cPIKAN provides slightly better inference of the unknown reaction parameter κ.

Method: B-cPIKAN equips cPIKAN with a Bayesian framework and uses Hamiltonian Monte Carlo to estimate parameter posteriors for uncertainty quantification.The method targets a noisy-data inverse problem for a 1D steady reaction-diffusion equation.
Results: B-cPIKAN and B-PINN perform similarly when inferring u and f from the noisy data.The comparison quantifies epistemic uncertainty for both Bayesian networks.
Results: B-cPIKAN infers κ slightly better than B-PINN: its error is lower and its uncertainty bounds the error.The exact reaction rate used for comparison is 0.7.
Problem: The inverse problem infers u and the unknown reaction rate κ from noisy data of u and f, with the exact κ equal to 0.7.Eight uniformly sampled measurements of u are corrupted with additive Gaussian noise.

3.7. 1D Burgers’ equation

The Burgers’ equation experiment compares DeepONet with two Chebyshev-KAN DeepOKANs for learning a solution operator and tests robustness to noisy inputs. DeepOKANs are reported to be more robust than DeepONet under noisy input functions.

Problem: The models learn the operator mapping arbitrary initial conditions u0 to the Burgers’ equation solution at t = 1.This is an operator-learning problem mapping functions to functions rather than coordinates to point values.
Data and training: The data use 128 uniform spatial grids, with 1,000 training functions and 200 testing functions, plus Fourier features and output normalization for DeepONet.The input functions are sampled from a Gaussian process.
Models: The comparison uses one DeepONet and two Chebyshev-KAN DeepOKANs with the architectures specified in Table 7.DeepOKANs use degree-three Chebyshev KANs, while DeepONet uses tanh activation.
Noise robustness: When trained on clean data and tested with 1%, 5%, and 10% noisy inputs, DeepOKANs are more robust than DeepONet.The noise is additive Gaussian noise whose standard deviation is proportional to the absolute input-function value.

3.8. 120-dimensional Darcy problem

The 120-dimensional Darcy benchmark compares DeepONet and DeepOKAN solution operators under clean training data and noisy testing inputs. Both models use matched branch and trunk architectures, while their representation functions differ.

Experimental setup: Table 9 evaluates DeepONet and DeepOKAN on the Darcy solution operator using identical branch and trunk architectures.Both networks use branch architecture [961, 100, 100, 100, 100] and trunk architecture [2, 100, 100, 100].
Experimental setup: DeepOKAN uses degree-three Chebyshev KANs, whereas DeepONet uses hyperbolic tangent activations.
Darcy problem: The Darcy problem imposes u(0, y) = 1, u(1, y) = 0 and zero normal derivatives on the remaining boundaries.
Darcy problem: The learned operator maps log(λ), the conductivity field representation, to u, the hydraulic head.
Data generation: The conductivity inputs are generated from a truncated Karhunen–Loève expansion of a Gaussian process with a specified kernel.
Data generation: Both log(λ) and u are represented on a 31 × 31 grid, retaining the first 120 expansion terms.

4. Learning in PIKANs

The paper interprets PINN and PIKAN optimization through information-bottleneck training stages identified by batch-wise gradient SNR. Both models progress from fitting through diffusion to total diffusion, with faster convergence after the final transition.

Information-bottleneck framework: The information bottleneck framework describes training as fitting, diffusion, and total diffusion stages that balance retained output information against discarded input details.
SNR diagnostic: Batch-wise SNR compares the L2 norm of mean loss gradients with the L2 norm of their standard deviation.
Fitting: During fitting, initially aligned gradients produce high SNR, while decreasing signal and increasing subdomain disagreement drive SNR lower.
Diffusion: During diffusion, the network weights explore directions that reduce errors across subdomains, producing a fluctuating low SNR.
Total diffusion: During total diffusion, SNR rises as the model exploits a consistent direction and simplifies representations by retaining important features.
Total diffusion: Relative L2 error decreases faster after total diffusion begins, and PINN+RBA and cPIKAN+RBA enter that stage first.

5. Summary

The study finds that modified Chebyshev-KAN representations can approach MLP-based performance across several scientific machine-learning tasks, but training stability, computational cost, and scalability remain important constraints.

Function approximation: MLPs and KANs both achieve high accuracy for discontinuous and oscillatory function approximation, but original KAN training is significantly slower.
Function approximation: Composing Chebyshev-KAN forward passes with tanh stabilizes training and makes approximation accuracy and runtime comparable to MLPs.
Dynamical systems: Chebyshev-KAN is less data-efficient than MLPs for structure-preserving and energy-conserving dynamical systems, but performs comparably with slightly more training data.
Noisy PDEs: Bayesian cPIKAN produces predicted means and uncertainties similar to Bayesian PINN at higher computational cost, while prior specification is less straightforward.
Operator learning: DeepOKAN shows competitive operator-learning performance relative to DeepONet and greater robustness to noisy input functions during testing after clean-data training.
Learning dynamics: PINNs and cPIKANs exhibit similar fitting, diffusion, and total-diffusion learning dynamics despite their architectural differences.
Future directions: Future work targets large-scale PDEs, time-dependent multidimensional problems, convergence theory, and industrial surrogate modeling.