Invertible Koopman neural operator for data-driven modeling of partial differential equations

TL;DR

This work introduces the Invertible Koopman Neural Operator (IKNO), a data-driven PDE surrogate that combines Koopman operator theory with an invertible neural network to learn a forward model via a globally linearized observable space while guaranteeing an exact inverse without a reconstruction loss. IKNO evolves low-frequency observables in the frequency domain using a structured Koopman matrix and recovers the physical state through the inverse observable, with high-frequency content filled by a learned $1\times1$ convolution and pre/post-processing enabling non-Cartesian domains. The authors validate IKNO on a suite of PDEs (1D Burgers, 2D shallow water, 2D Navier–Stokes, 2D Darcy, 2D Euler for airfoils) and a global weather dataset, showing superior accuracy and robust zero-shot super-resolution compared to Fourier Neural Operators and Koopman Neural Operators. The approach delivers resolution-invariant predictions and supports advanced tasks such as inverse design, highlighting potential for efficient, high-fidelity PDE surrogates in engineering and geoscience contexts.

Abstract

Koopman operator theory is a popular candidate for data-driven modeling because it provides a global linearization representation for nonlinear dynamical systems. However, existing Koopman operator-based methods suffer from shortcomings in constructing the well-behaved observable function and its inverse and are inefficient enough when dealing with partial differential equations (PDEs). To address these issues, this paper proposes the Invertible Koopman Neural Operator (IKNO), a novel data-driven modeling approach inspired by the Koopman operator theory and neural operator. IKNO leverages an Invertible Neural Network to parameterize observable function and its inverse simultaneously under the same learnable parameters, explicitly guaranteeing the reconstruction relation, thus eliminating the dependency on the reconstruction loss, which is an essential improvement over the original Koopman Neural Operator (KNO). The structured linear matrix inspired by the Koopman operator theory is parameterized to learn the evolution of observables' low-frequency modes in the frequency space rather than directly in the observable space, sustaining IKNO is resolution-invariant like other neural operators. Moreover, with preprocessing such as interpolation and dimension expansion, IKNO can be extended to operator learning tasks defined on non-Cartesian domains. We fully support the above claims based on rich numerical and real-world examples and demonstrate the effectiveness of IKNO and superiority over other neural operators.

Figures (10)

• Figure 1: Structure of the INN. a Forward process of the INN. b Inverse process of the INN. c Splitting operation $\tilde{\mathcal{S}}$ and its inverse $\tilde{\mathcal{S}}^{-1}$. d Zeros-concatenating operation $\mathcal{S}_{i}$ and its inverse $\mathcal{S}_{i}^{-1}$. e Merging operation $\tilde{\mathcal{M}}$ and its inverse $\tilde{\mathcal{M}}^{-1}$. (f) Nonlinear mapping operation $\mathcal{H}_{i}$ (irreversible).
• Figure 2: Architecture of the proposed Invertible Koopman Neural Operator (IKNO). a Core components of the IKNO. b Structured linear matrix for approximating the Koopman operator. c The proposed IKNO is resolution-invariant, i.e., the IKNO trained at low resolution can be directly used to predict super-resolution results.
• Figure 3: Results of the 1D Burgers equation. a Comparison of the loss function. The proposed IKNO shows better convergence during training, achieving the lowest loss. Note that the training process is performed at the resolution of 32. b Comparison of the resolution-invariance and super-resolution prediction performance for different methods. The proposed IKNO has the minimum MAE at all resolutions. c The ground truth, prediction, and relative error distribution of the IKNO at the resolution of 32. Three temporal snapshots are also portrayed in d1-d3. e The ground truth, prediction, and error distribution of the IKNO at the resolution of 1024, 32$\times$ the training set. Three temporal snapshots are also portrayed in f1-f3. Even under the zero-shot super-resolution condition, the proposed IKNO still captures sharp jumps and discontinuity in the velocity field with high accuracy.
• Figure 4: Results of the 2D shallow water equation. a Comparison of the loss function. The proposed IKNO shows better convergence during training, achieving the lowest loss. Note that the training process is performed at the resolution of $64 \times 64$. b Comparison of the MAE with different prediction steps. Results of resolution-invariance are also plotted, where the methods are evaluated at the resolution of $64 \times 64$ and $128 \times 128$ (2$\times$ the training set). c The ground truth, prediction, and relative error distribution of the IKNO at the resolution of $64 \times 64$. e The ground truth, prediction, and error distribution of the IKNO at the resolution of $128 \times 128$, 2$\times$ the training set.
• Figure 5: Results of the 2D incompressible Navier-Stokes equation. a Comparison of the loss function with viscosity $\nu = 10^{-3}$. The proposed IKNO shows better convergence during training, achieving the lowest loss. b Comparison of the MAE under different prediction steps with viscosity $\nu = 10^{-3}$. c Comparison of the loss function with viscosity $\nu = 10^{-4}$. Note that the training process is performed at the resolution of $64 \times 64$. d Comparison of the MAE under different prediction steps with viscosity $\nu = 10^{-4}$. Results of resolution-invariance are also plotted, where the methods are evaluated at the resolution of $64 \times 64$ and $256 \times 256$ (4$\times$ the training set). e The ground truth, prediction, and relative error distribution of the velocity obtained by the IKNO at the resolution of $64 \times 64$, with $\nu = 10^{-3}$. f and g The ground truth, prediction, and relative error distribution of the velocity obtained by the IKNO at the resolution of $64 \times 64$ and $256 \times 256$, respectively, with $\nu = 10^{-4}$.