---
title: Koopman Autoencoder & AIKAE
url: https://www.emergentmind.com/topics/koopman-autoencoder
type: topic
---

# Koopman Autoencoder & AIKAE

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 [2503.12930].

## 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$, $\psi$, and $K$ to minimize reconstruction and prediction losses:
\[
\hat{x}_{t+\tau} = \psi(K^\tau \phi(x_t)).
\]

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 [2503.12930, 2309.10291].

## 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 $\phi$**: $\phi: \mathbb{R}^n \to \mathbb{R}^n$ is composed of $L$ coupling layers (RealNVP/affine transformations), each operating as
  \[
  \begin{align*}
  v_a &= u_a \\
  v_b &= u_b + m(u_a)
  \end{align*}
  \]
  or
  \[
  v_b = u_b \odot \exp(s(u_a)) + t(u_a)
  \]
  where $m, s, t$ are MLPs. Invertibility is exact, with tractable computation of the Jacobian determinant.

- **Non-invertible augmentation encoder $\chi$**: $\chi:\mathbb{R}^n \to \mathbb{R}^p$ is an arbitrary feed-forward network, producing additional latent coordinates $z^a_t$.

- **Combined latent state**: Form $z_t = [z^i_t; z^a_t] \in \mathbb{R}^d$, with $d = n + p$. Latent dynamics proceed via
  \[
  z_{t+1} = K z_t,\qquad K \in\mathbb{R}^{d \times d}
  \]
  Only the invertible coordinates are decoded:
  \[
  \hat{x}_{t+\tau} = \phi^{-1}(z^i_{t+\tau}),\quad [z^i_{t+\tau}; z^a_{t+\tau}] = K^\tau z_t
  \]
[2503.12930].

## 3. Objective Functions and Optimization

AIKAE is optimized by joint minimization of loss terms:

- **Prediction loss** (long-term forecast):
  \[
  L_\text{pred} = \sum_{t=0}^{T-\tau} \left\| x_{t+\tau} - \phi^{-1}([I_n, 0] K^\tau z_t) \right\|^2
  \]
- **Linearity loss** (Koopman invariance):
  \[
  L_\text{lin} = \sum_{t=0}^{T-1} \left\| z_{t+1} - K z_t \right\|^2
  \]
- **Orthogonality/norm-preservation loss**:
  \[
  L_\text{orth} = \sum_{t} \left( \|z_{t+1}\| - \|z_t\| \right)^2
  \]
  or regularization directly on $K$:
  \[
  L_{K\text{orth}} = \left\| K K^\top - I_d \right\|_F^2
  \]
- **(Optionally) Flow log-likelihood** (if modeling with stochastic flows):
  \[
  L_\text{flow} = \mathbb{E}_{x \sim \text{data}} \left[ -\log p_Y(\phi(x)) - \log |\det J_\phi(x)| \right]
  \]

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 [2503.12930].

## 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)               | $d \leq n$         | MLP            | Rec, Lin, Pred       | None                 |
| IKAE          | Invertible Flow (RealNVP, NICE)| $d = n$            | $\phi^{-1}$    | Lin, Pred            | Bijectivity, $\det$  |
| AIKAE         | Invertible + Non-invertible    | $d = n + p$        | $\phi^{-1}$    | Pred, Lin, Orth      | Partial bijectivity  |
| CKAE          | MLP / Linear                   | $d$                | MLP            | Fwd/Bwd, Consistency | $K_\text{fwd} K_\text{bwd} \approx I$   |
| KIA           | Invertible NN                  | $d$                | 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 [2503.12930].

## 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 [2503.12930].

- **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 [2503.12930].

## 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 $K$ 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 [2503.12930].

## 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 [2503.12930].

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.

Source: https://www.emergentmind.com/topics/koopman-autoencoder