Low-Rank Kolmogorov-Arnold Generator Flows for Interpretable Operator Learning

The gist: The Kolmogorov–Arnold flow (KAFlow) learns the generator of a physical operator in a low-rank symbol space, enabling exact integration and automatic distillation, which significantly lowers out-of-distribution loss and improves accuracy across various physical dynamics.

Low-Rank Structure of Position-Dependent Operators

The paper demonstrates that finite-order variable coefficient operators possess a resolution-independent separation rank, and the symbols learned by KANO for these benchmarks have rank two. Specifically, a rank-R symbol is applied exactly with O(Rm log m) operations instead of m2 symbol evaluations. This structure arises because position-dependent operators can be represented by a translation-invariant dispersion and dissipation as a spectral multiplier, position-dependent terms as a few separable products ar(x)mr(ξ), and the pointwise nonlinearity as a gauge-fixed Kolmogorov–Arnold function.

An Interpretable Generator Network

KAFlow is designed to learn the generator in this low-rank symbol space using univariate Kolmogorov–Arnold functions. The architecture enforces identifiability by gauge fixing, meaning every term has a unique place in the architecture, which makes automatic distillation possible. Furthermore, it integrates dynamics with exact exponential integrators and distills its closed form automatically without human-designed templates. The automatic distillation process involves regressing the learned components on monomials and dictionary functions to recover the exact closed form for all 18 trained models.

Validation on All KANO Benchmarks and Beyond

KAFlow shows significant performance gains over existing methods on synthetic benchmarks. For long-horizon quantum dynamics, KAFlow lowers the state infidelity by up to 6 orders of magnitude and by 30 times when trained only on position and momentum measurements. KAFlow also learns a variable-order fractional operator 15 times more accurately than existing operator networks and discovers nonlocal, fractional and non-polynomial dispersion relations where sparse regression fails.

Continuous-Time Flow and Integrators

For dynamics, KAFlow represents the target operator by its generator Gθu, which is integrated exactly in time. The flow integrates ∂tuˆ = l uˆ +N(uˆ) with the fourth-order Lawson integrating-factor Runge–Kutta scheme. For quantum dynamics, the generator is integrated by Strang splitting where the kinetic symbol T(ξ) is non-negative, monotone and zero at ξ = 0.

Nonlocal and Dispersive Dynamics

KAFlow attains the lowest 100-step loss on all five evolution equations with only 302 parameters, against models like FNO which require hundreds of thousands of parameters. The automatic distillation selects the exact closed form for equations like the fractional Burgers equation and Benjamin–Ono equation, recovering the fractional order and Hilbert-type symbol with coefficient errors of 0.02% to 6%.

Computational Cost

The computational cost of KAFlow grows with m log m for a training step, which is significantly better than the O(m 2) symbol evaluations required by standard KANO. When the resolution grows to m = 1,024, the cost difference becomes substantial.

Discussion

The paper concludes that the structure of physical operators themselves allows for a much more compact representation than what is typically achieved by standard neural operators. The key insight is that representing each component where it is both sparse and low-rank leads to superior generalization and interpretability. This approach extends to other scientific modeling problems by learning the generator instead of the propagator and fixing the gauge of components to turn symbolic regression into an automatic step. The learned structure describes many operators of mathematical physics, including variable-coefficient diffusion and transport.

Acknowledgments

M. Cheon conceived the study, developed the method, ran the experiments, analyzed the results and wrote the manuscript. The author acknowledges that they did not receive funding for this work. The data generators for the variable-order fractional operator and evolution equations, the trained KAFlow models and the code are available from the corresponding author upon reasonable request. The synthetic-operator and quantum datasets are available from the official KANO repository <ref:

1. Improved generalization of physical laws to unseen data by 3.4 to 11 times on synthetic benchmarks, specifically by learning a low-rank symbol built from univariate Kolmogorov–Arnold functions. This allows the system to recover every coefficient to within 10−4 without symbolic templates, enabling robust out-of-distribution generalization beyond the training distribution.

2. Enhanced accuracy in modeling complex physical dynamics by learning variable-order fractional operators 15 times more accurately than existing operator networks and discovering nonlocal, fractional and non-polynomial dispersion relations on which sparse regression fails. This capability allows the system to accurately represent equations like the Benjamin–Ono equation or Whitham’s waterwave model.

3. Superior long-horizon quantum dynamics simulation by lowering state infidelity by up to 6 orders of magnitude compared to Q-KANO, achieving 1.9 × 10−7 infidelity after 90 extrapolated steps with position and momentum supervision, which is 30 times lower than Q-KANO.

4. Automatic closed-form law extraction without human intervention by enforcing gauge fixing such that every term has a unique place in the architecture, enabling the system to distill its closed form automatically, as demonstrated by recovering coefficients to within 10−5 for specific operators.

5. More efficient computational scaling during training and inference by learning the generator in a low-rank symbol space, which reduces evaluation cost from m2d symbol evaluations and O(m3d) in the reference implementation of KANO to O(R + J + 1) md log md operations.

Neural operators that learn physical laws from data should generalize beyond their training distribution and expose the learned law in closed form. The Kolmogorov-Arnold neural operator achieves both for position-dependent dynamics by learning a dense pseudo-differential symbol over all positions and frequencies, at a cost that grows with the square of the resolution and with closed forms that must be read out by hand. This dense symbol is unnecessary: a finite-order variable-coefficient differential operator has a symbol whose separation rank is bounded by the number of its coefficient functions, independently of the resolution, and the symbols that this network learns on its own benchmarks have rank two. We introduce the Kolmogorov-Arnold flow (KAFlow), which learns the generator of a physical operator as a low-rank symbol built from univariate Kolmogorov-Arnold functions, fixes the gauge of these functions so that every term has a unique place, integrates the dynamics with exact exponential integrators, and distills its closed form automatically. On the synthetic benchmarks of the Kolmogorov-Arnold neural operator, KAFlow lowers the out-of-distribution loss by 3.4 to 11 times and recovers every coefficient to within 10-4 without symbolic templates, reaching the precision floor of the data. On long-horizon quantum dynamics, it lowers the state infidelity by up to 6 orders of magnitude, and by 30 times when trained only on position and momentum measurements. KAFlow also learns a variable-order fractional operator 15 times more accurately than existing operator networks and discovers nonlocal, fractional and non-polynomial dispersion relations on which sparse regression fails.