Koopman Autoencoder & AIKAE

• Koopman autoencoder is a neural latent-variable model that uses Koopman operator theory to map nonlinear dynamics into a linearly evolving latent space.
• The AIKAE variant augments an invertible encoder with additional non-invertible components to boost expressiveness while maintaining exact state reconstruction.
• These models facilitate interpretable modal analysis, efficient surrogate modeling, and robust long-term predictions in applications such as weather forecasting and satellite data assimilation.

A Koopman autoencoder is a class of neural latent-variable models that enforces linear evolution in the learned latent space, in accordance with Koopman operator theory, for modeling nonlinear dynamical systems. These models aim to learn nonlinear state embeddings (“liftings”) in which time evolution is linear, enabling direct multi-step prediction, interpretable mode analysis, and efficient surrogate modeling for a wide range of physical, spatiotemporal, and data-driven systems. A recent major development within this class is the Augmented Invertible Koopman AutoEncoder (AIKAE), which resolves representational limitations of previous invertible Koopman autoencoders by augmenting their dimension with a non-invertible subspace, while maintaining exact state reconstruction and globally linear latent dynamics (Frion et al., 17 Mar 2025).

1. Theoretical Foundations

Koopman operator theory provides a framework to linearize nonlinear dynamical systems via observables. For a system of the form

$x_{t+1} = F(x_t),\quad x_t \in \mathbb{R}^n,$

there exists an infinite-dimensional linear operator, the Koopman operator $\mathcal{K}$, such that for any observable $g$: $\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$ In practice, the goal is to find a finite set of observables $g_1,\ldots,g_d$ (or, in modern terms, a neural encoder $\phi: \mathbb{R}^n \to \mathbb{R}^d$) that spans an approximately invariant subspace for $\mathcal{K}$, such that the finite-dimensional approximation

$$z_{t+1} \approx K z_t,\quad z_t = \phi(x_t),\quad K\in\mathbb{R}^{d\times d}$$

holds. A Koopman autoencoder further includes a learned decoder $\psi:\mathbb{R}^d \to \mathbb{R}^n$ to reconstruct the state, jointly optimizing $\phi$, $\mathcal{K}$0, and $\mathcal{K}$1 to minimize reconstruction and prediction losses: $\mathcal{K}$2

Invertible Koopman autoencoders (IKAEs) adopt a bijective normalizing flow for the encoder, with an analytically invertible decoder, achieving exact reconstruction and tractable latent density modeling (Frion et al., 17 Mar 2025, Tayal et al., 2023).

2. AIKAE Architecture

The AIKAE introduces an augmented latent space to improve expressiveness beyond the input dimension, circumventing the bottleneck induced by invertibility constraints in prior IKAEs. The architecture is defined as follows:

• Invertible encoder $\mathcal{K}$3: $\mathcal{K}$4 is composed of $\mathcal{K}$5 coupling layers (RealNVP/affine transformations), each operating as

$\mathcal{K}$6

or

$\mathcal{K}$7

where $\mathcal{K}$8 are MLPs. Invertibility is exact, with tractable computation of the Jacobian determinant.

• Non-invertible augmentation encoder $\mathcal{K}$9: $g$0 is an arbitrary feed-forward network, producing additional latent coordinates $g$1.
• Combined latent state: Form $g$2, with $g$3. Latent dynamics proceed via

$g$4

Only the invertible coordinates are decoded:

$g$5

(Frion et al., 17 Mar 2025).

3. Objective Functions and Optimization

AIKAE is optimized by joint minimization of loss terms:

• Prediction loss (long-term forecast):

$g$6

• Linearity loss (Koopman invariance):

$g$7

• Orthogonality/norm-preservation loss:

$g$8

or regularization directly on $g$9:

$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$0

• (Optionally) Flow log-likelihood (if modeling with stochastic flows):

$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$1

Total loss combines these terms, and end-to-end training is performed with Adam. Multi-step prediction horizons and curriculum strategies are used to ensure both short- and long-horizon forecast quality (Frion et al., 17 Mar 2025).

4. Comparison to Other Koopman Autoencoder Variants

The following table summarizes main Koopman AE variants and key distinguishing features:

Model	Encoder Type	Latent Dim	Decoder	Key Losses	Constraints
KAE	General NN (MLP)	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$2	MLP	Rec, Lin, Pred	None
IKAE	Invertible Flow (RealNVP, NICE)	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$3	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$4	Lin, Pred	Bijectivity, $\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$5
AIKAE	Invertible + Non-invertible	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$6	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$7	Pred, Lin, Orth	Partial bijectivity
CKAE	MLP / Linear	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$8	MLP	Fwd/Bwd, Consistency	$\mathcal{K}\,g(x_t) = g(F(x_t)) = g(x_{t+1}).$9
KIA	Invertible NN	$g_1,\ldots,g_d$0	Invertible NN	Fwd/Bwd, Rec	Invertibility

AIKAE is unique in disentangling invertible (for exact state recovery) and augmentation (for increased expressivity), allowing latent inflation and improved learning of nearly Koopman-invariant subspaces (Frion et al., 17 Mar 2025).

5. Applications and Empirical Evaluation

AIKAE and its predecessors have been applied to a diverse set of time-series and spatiotemporal prediction problems:

• Informer long-term forecasting benchmark: On ETTh1, ETTh2, ETTm1/2, ECL, Traffic, Weather, AIKAE outperforms previous IKAEs and competitive models (PatchTST, TimesNet, iTransformer) by 1–3% MSE/MAE, achieving monotonic performance gains with increased lookback window size. Unlike transformers, AIKAE does not saturate with longer input histories and is robust to the curse of horizon (Frion et al., 17 Mar 2025).
• Satellite data assimilation: On Sentinel-2 spectral image sequences, AIKAE achieves minimum MSE/MAE relative to KAEs, non-augmented IKAEs, or zero-padded IKAEs, particularly for forecasting and cross-region transfer. Exact invertibility renders it more stable to overfitting and allows variational assimilation via reverse-time latent optimization (Frion et al., 17 Mar 2025).

6. Interpretability and Operator Structure

AIKAE, by virtue of invertible latent flows and explicit linear Koopman evolution, admits interpretability in both reconstruction and spectral analysis. The eigen-spectrum of the Koopman matrix $g_1,\ldots,g_d$1 encodes dominant modes, frequencies, and decay rates. Disentanglement between invertible and augmentation coordinates allows separation of static/reversible and dynamic/non-reversible features, which is beneficial for both data assimilation and transfer tasks (Frion et al., 17 Mar 2025).

7. Extensions and Future Directions

Research opportunities include:

• Stochastic Koopman models: Leveraging the flow's tractable log-likelihood for uncertainty quantification and probabilistic forecasting.
• Trend/seasonality/frequency augmentation: Integrating Fourier blocks or trend decomposition in the encoder, analogous to architectures such as Fedformer, to enhance performance for highly structured signals.
• Hierarchical and multi-scale representations: Constructing layered or multi-resolution latent spaces to scale AIKAE for very high-dimensional or spatiotemporal systems (Frion et al., 17 Mar 2025).

AIKAE represents the current state-of-the-art for invertible Koopman autoencoding in time series forecasting and spatiotemporal assimilation, combining exact reconstruction, latent space inflation, and stable, interpretable linear prediction.