PaperlayerView exact revision on arXiv

Source-linked AI summary

KAN-ODEs: Kolmogorov-Arnold Network Ordinary Differential Equations for Learning Dynamical Systems and Hidden Physics

Benjamin C. Koenig, Suyong Kim, Sili Deng

arXiv:2407.04192v3cs.LG

TL;DR

Dynamical-system learning needs methods that reduce reliance on expert knowledge without sacrificing interpretability. This paper embeds KANs as gradient getters in Neural ODEs, and reports stronger quantitative performance than MLP-based Neural ODEs alongside data-lean modeling, symbolic inference, and interpretable learned representations.

  • Problem

    Data-driven dynamical modeling often lacks interpretability, while interpretable methods typically require substantial governing-equation knowledge or functional assumptions.

  • Method

    KAN-ODEs use KANs as learnable gradient getters inside Neural ODEs, combining KAN activation functions with differentiable ODE solving and adjoint-based training.

  • Results

    KAN-ODEs outperform comparable MLP-based Neural ODEs across the reported quantitative metrics and generalize in data-lean wave, shock, and complex-valued PDE examples.

  • Takeaways & Limitations

    The framework supports interpretable and modular dynamical-system models with zero prior physics knowledge, including activation visualization and symbolic-regression postprocessing.

Abstract

from arXiv · show

Kolmogorov-Arnold networks (KANs) as an alternative to multi-layer perceptrons (MLPs) are a recent development demonstrating strong potential for data-driven modeling. This work applies KANs as the backbone of a neural ordinary differential equation (ODE) framework, generalizing their use to the time-dependent and temporal grid-sensitive cases often seen in dynamical systems and scientific machine learning applications. The proposed KAN-ODEs retain the flexible dynamical system modeling framework of Neural ODEs while leveraging the many benefits of KANs compared to MLPs, including higher accuracy and faster neural scaling, stronger interpretability and generalizability, and lower parameter counts. First, we quantitatively demonstrated these improvements in a comprehensive study of the classical Lotka-Volterra predator-prey model. We then showcased the KAN-ODE framework's ability to learn symbolic source terms and complete solution profiles in higher-complexity and data-lean scenarios including wave propagation and shock formation, the complex Schrödinger equation, and the Allen-Cahn phase separation equation. The successful training of KAN-ODEs, and their improved performance compared to traditional Neural ODEs, implies significant potential in leveraging this novel network architecture in myriad scientific machine learning applications for discovering hidden physics and predicting dynamic evolution.

I. INTRODUCTION

The introduction motivates combining data-driven dynamical-system modeling with interpretability and presents KAN-ODEs as KAN-based gradient getters within Neural ODEs. This combination aims to retain grid-independent ODE modeling while enabling interpretable activations and symbolic relationships without prior knowledge or presumed functional forms.

  • Data-driven dynamical-system methods trade reduced expert-knowledge requirements against interpretability, whereas interpretable approaches often require governing-law knowledge or enforcement.
  • Neural ODEs learn continuous state evolution through neural gradient getters and ODE solvers, providing grid and timescale flexibility without requiring prior knowledge.
  • KANs learn activation functions rather than only weights and biases, supporting interpretability, faster convergence with fewer parameters, and symbolic-regression postprocessing.
  • KAN-ODEs use KANs as gradient getters to learn interpretable and modular ODE and PDE models without prior knowledge or functional-form assumptions.
  • KAN activations combine gridded Gaussian RBFs with a direct residual activation pathway, whose parameter requirements scale with input dimension, output dimension, and grid size.
  • A KAN is represented as a composition of layer mappings, with notation [n_l, n_l+1, N] encoding input width, output width, and grid size.

B. KAN-Ordinary Differential Equations

KAN-ODEs replace the neural gradient getter in a dynamical-system ODE with a KAN and train it through differentiable ODE solving and adjoint sensitivities. The framework maps state variables to their derivatives and updates trainable KAN parameters using observation-based loss.

  • KAN-ODEs approximate the system equation g with a KAN whose input and output dimensions equal the state-vector dimension.
  • The KAN gradient getter is integrated from initial time t0 and initial condition u0 to generate the state trajectory at a desired time.
  • The training cycle alternates forward ODE solution generation with adjoint-based loss-gradient computation and updates to KAN activation functions.
  • Training minimizes mean squared error between KAN predictions and observations over the time steps.
  • The adjoint sensitivity method backpropagates through the ODE integrator and is adopted for its scaling advantages with potentially large KANs.

III. EXPERIMENTS

The experiments assess KAN-ODEs across increasingly demanding dynamical-system settings. They include Lotka-Volterra benchmarking, symbolic source-term inference in a PDE simulation, and complete data-driven solution profiles for higher-complexity systems.

  • Five inference examples demonstrate KAN-ODE capability across benchmark ODEs, PDE submodels, and standalone higher-complexity dynamical systems.
  • Lotka-Volterra experiments compare KAN-ODEs with MLP-based Neural ODEs while examining neural scaling and representations obtained through sparsification and pruning.
  • A one-dimensional wave-propagation PDE uses KAN-ODEs to infer a symbolic source term and demonstrates their use as submodels in higher-complexity simulations.
  • Standalone KAN-ODE models are used to infer complete data-driven solution profiles in more complex systems.

A. KAN-ODEs vs Neural ODEs: Extensive Comparison via Lotka-Volterra Equations

The Lotka-Volterra study benchmarks KAN-ODEs against Neural ODEs under matched dynamical-system settings and evaluates accuracy, convergence, and generalization. The representative KAN-ODE converges to low loss and reconstructs both observed and held-out trajectory windows.

  • The benchmark uses Lotka-Volterra dynamics with α = 1.5, β = 1, γ = 1, δ = 3, initial condition u0 = [1, 1], and t ∈ [0, 14] s.
  • A 240-parameter KAN-ODE reaches strong convergence by 10^4 epochs and converges toward 10^-7 MSE by approximately 10^5 epochs.
  • The converged KAN-ODE is nearly indistinguishable between the training window t ∈ [0, 3.5] and testing window t ∈ [3.5, 14].
  • The comparison evaluates KAN-ODEs against standard MLP-based Neural ODEs on accuracy, convergence speed, neural scaling, interpretability, and generalization.

1. Benchmark Tests and Neural Scaling Behavior

In a controlled Lotka–Volterra benchmark, KAN-ODEs outperformed comparable Neural ODEs in accuracy and convergence efficiency. Across model sizes, KAN-ODEs also exhibited quartic-like scaling, while MLP-based Neural ODEs scaled more slowly.

  • Controlled comparison: Using identical training settings, a 240-parameter KAN-ODE was compared fairly against a slightly larger 252-parameter MLP-based Neural ODE.The same train/test windows, temporal grid, system parameters, and ODE solver were used.
  • Accuracy and convergence: 8.3 × 10^-7 training MSE was reached by the KAN-ODE after 10^5 epochs, compared with 3 × 10^-5 for the Neural ODE.Both loss profiles oscillated during training, but the KAN-ODE oscillations were smaller.
  • Generalization across time grids: KAN-ODE solutions could be generated at 0.01-second steps instead of the 0.1-second training grid without detectable testing changes.This indicates no observed overfitting to the specific training timesteps in this benchmark.
  • Computational efficiency: 3–4× effective speedup was achieved despite 2.5–3× slower iteration, because the KAN-ODE matched predictive performance with 10× fewer epochs.The KAN-ODE reached 2.6 × 10^-5 training loss after 10^4 epochs, compared with 3.0 × 10^-5 for the Neural ODE after 10^5 epochs.
  • Neural scaling behavior: N^-4 scaling was observed for KAN-ODEs, whereas two MLP depths scaled more slowly than N^-2 as model size increased.With 240 KAN parameters, convergence reached saturation; comparable MLP performance was projected to require thousands or tens of thousands of parameters if it did not plateau earlier.

2. Interpretation and Generalization of KAN-ODEs with Varying Sizes

KAN-ODEs become progressively smaller and more interpretable through sparsification, pruning, and symbolic regression, while retaining strong Lotka-Volterra generalization. Across tested models, KAN-based approaches outperform the MLP Neural ODE in extrapolating beyond limited training data.

  • Generalization comparison: The study compares dense, sparse, symbolic, and fully symbolic KAN-ODE variants against an MLP Neural ODE using unseen Lotka-Volterra state conditions.All models were trained from a single initial condition before heuristic testing on other (x, y) combinations.
  • Sparsification and pruning: Sparsification and pruning reduce the KAN-ODE from 240 to 72 parameters, producing the, architecture.Sparsification penalizes nonzero parameters, and pruning removes nodes whose inputs and outputs remain below the pruning threshold.
  • Symbolic regression: Symbolic regression replaces the sparse KAN’s twelve univariate activations with compact expressions without requiring knowledge beyond the training data.This yields a more human-readable representation while preserving the learned Lotka-Volterra relationships.
  • Generalization comparison: The MLP Neural ODE has low error near training points but significant gradient error across most regions outside them, indicating weak extrapolation.The error landscapes report this pattern for both dx/dt and dy/dt.
  • KAN size and generalization: Dense KAN-ODEs reduce gradient errors throughout the domain, while sparse KAN-ODEs add near-perfect reconstruction bands and retain relatively low error regions.Sparse models improve over dense KANs, although edge-of-domain performance is mixed for some derivatives.
  • Overall assessment: The authors report that KAN-ODEs beat comparable MLP-based Neural ODEs across accuracy, speed, model size, convergence rate, interpretability, and generalization.They identify sparse and symbolic KAN-ODEs as more broadly applicable than complete symbolic regression for realistic larger systems.

B. Modeling Hidden Physics in PDEs: Fisher-KPP PDE

The Fisher-KPP example uses KAN-ODEs to learn an unknown reaction source term while reconstructing the PDE solution and extracting an interpretable symbolic approximation.

  • KAN-ODEs are applied to learn hidden physics in a Fisher-KPP reaction-diffusion PDE, including symbolic functional relationships from training data.
  • The reaction term ru(1 −u) is treated as unknown and modeled with a KAN, while D denotes diffusion and r denotes local growth.
  • The problem is discretized for adjoint-based training, and the KAN uses a single layer with one node and ten gridded basis functions.
  • After 5,000 updates, the learned model accurately reconstructs the solution field, with one learned activation function representing the reaction source term.
  • Symbolic regression converts the learned activation function into an explicit expression, which approximates the true reaction source term closely when r = 1.0.

C. Data-Driven Solutions of PDEs

The Burgers’ equation example tests whether KAN-ODEs can infer complete PDE dynamics from sparse spatiotemporal observations, including behavior near shock formation and at unseen times.

  • KAN-ODEs infer the hidden states u(x, t) of Burgers’ convection-diffusion dynamics, representing the system’s spatiotemporal evolution.
  • Training uses five solution snapshots at t ∈ {0.1, 0.3, 0.5, 0.7, 0.9}, emphasizing a sparse dataset for testing extrapolation.
  • The KAN functions as a nonlinear operator for the discretized PDE solution, with input and output dimensions determined by the discretized state variables.
  • After 20,000 training epochs, the model predicts temporal evolution well and agrees strongly with training profiles, including near the shock wave.
  • The trained model extrapolates solution fields to unseen testing times, as shown by comparisons between inferred profiles and unseen ground-truth data.

2. Schr¨odinger Equation

The Schrödinger equation example evaluates KAN-ODE surrogate modeling for complex-valued wave dynamics using sparse observations of real and imaginary components.

  • The Schrödinger equation is used as a more involved surrogate-modeling problem with complex-valued states in quantum mechanics.
  • The model receives observations at eight selected times, with blue and red markers representing the real and imaginary parts of u, respectively.
  • The KAN-ODE doubles the discretized state dimension to represent both real and imaginary values of the complex field.
  • Despite eight training samples, the trained KAN-ODE reconstructs the complete spatiotemporal wave profile and agrees with the ground truth at all times.
  • Together with the shock-formation example, the result demonstrates matching and generalization across sparse temporal windows while retrieving spatial information without special KAN treatment.

IV. CONCLUSIONS

The conclusions present KAN-ODEs as a Neural ODE framework using KAN gradient-getters, combining data-driven flexibility with interpretability and parameter-efficient scaling.

  • KAN-ODEs replace MLP gradient-getters in Neural ODEs while retaining a black-box training approach requiring zero prior knowledge of dynamical-system physics.
  • The framework supports interpretable and modular models through activation-function visualization and symbolic regression.
  • KAN-ODEs outperform similar MLP-based Neural ODEs across training accuracy, testing accuracy, convergence speed, neural scaling rate, and generalizability.
  • RBF-based KAN-ODEs achieve better training and testing performance with less computational time and smaller network size.

APPENDIX A. SPARSE REGRESSION ON LOTKA-VOLTERRA KAN-ODE

The appendix describes symbolic KAN activations for Lotka–Volterra dynamics and evaluates KAN-ODE architectures on Fisher-KPP and Burgers’ equations. It also records the Allen-Cahn numerical setup used elsewhere in the appendix.

  • Symbolic activations: Symbolic KAN activations for Lotka–Volterra use x and y in the first layer, followed by n1, n2, and n3 in the second layer.
  • Symbolic activations: Each hidden node combines separate x- and y-dependent contributions, and the KAN derivative decomposes into corresponding node-wise terms.
  • Architecture evaluation: Fisher-KPP and Burgers’ equation tests varied KAN architectures to measure performance across different parameter counts.
  • Architecture evaluation: The Fisher-KPP hidden-physics model used a relatively small KAN with very few parameters and replaced only the source term.
  • Architecture evaluation: Training-loss scaling proportional to N^-4 was included as a reference line when comparing architectures by total parameter count.
  • Numerical setup: For the Allen-Cahn numerical setup, the PDE was discretized with Δx = 0.05 using central differences and solved with the Tsit5 ODE integrator.

C.1. Modeling Hidden Physics in PDEs

KAN-ODEs were used to learn the Allen-Cahn reaction source term and to reconstruct complete solution fields from sparse observations. The learned source closely matched the exact expression, while extrapolation beyond the training window required more data.

  • Modeling hidden physics: A single-node KAN encoded one activation function to model the unknown Allen-Cahn source term after converting the discretized PDE into an ODE system.
  • Modeling hidden physics: The candidate-expression table reported seven symbolic reaction-source expressions together with their complexities and losses.
  • Modeling hidden physics: After approximately 65,000 updates, the trained KAN-ODE accurately reconstructed the Allen-Cahn solution field and produced a symbolic activation function.
  • Modeling hidden physics: Equation C5, KAN(u) = (5.0015 − 5.0021u^2)u, was remarkably close to the exact source term 5u − 5u^3.
  • Complete solution reconstruction: For surrogate modeling, the entire right-hand side was replaced by a two-layer KAN trained on five profiles at t = 0.1, 0.3, 0.5, 0.7, and 0.9.
  • Complete solution reconstruction: Within the considered 1 s range, KAN-ODEs reconstructed physical phenomena between sparse training samples without prior physical knowledge in the model.
Loading 2407.04192v3…