Information Shapes Koopman Representation

AI summary

Overview

Research area: Machine learning for dynamical systems — specifically Koopman operator theory, representation learning, and information-theoretic approaches to latent-variable models.

Technical level: Advanced. The paper assumes familiarity with the Koopman operator, mutual information, the Information Bottleneck principle, the von Neumann entropy, variational autoencoders, and spectral decompositions.

Scope: The paper develops an information-theoretic account of why learned Koopman representations succeed or fail, and translates that account into a practical loss function validated on physical simulations, visual-control tasks, and graph-structured dynamics.

What This Paper Is About

The Koopman operator turns nonlinear dynamics into linear evolution, but it lives in an infinite-dimensional function space, so practitioners must learn a finite-dimensional latent subspace. Deep models such as variational autoencoders have been used for this, yet the latent representations they produce are often unstable, collapse onto a few dominant modes, or fail to reproduce reliable long-horizon dynamics.

This paper argues that the root cause is suboptimal representation learning — latent variables that do not balance compactness against expressiveness. Framing Koopman learning as a dynamical Information Bottleneck problem, the authors ask whether structural simplicity (latent linearity) and expressiveness (sufficient mode diversity) can be jointly optimized under a single information-theoretic principle, and they build such a principle and an accompanying algorithm.

Key Contributions

1. A theoretical account of information loss in Koopman representations. The authors prove an information-loss chain, (I(x_{n-1};x_n) \geq I(z_{n-1};x_n) \geq I(z_{n-1};z_n)), showing that the Koopman operator sets an information limit on the latent space, and they convert this into a total-variation and mean-squared autoregressive error bound in which the error grows with the gap (I(x_{n-1};x_n) - I(z_{n-1};z_n)) summed over steps.

2. A disentanglement of latent information into three spectrally interpretable components. They decompose the latent mutual information into temporal-coherent information (associated with Koopman eigenvalues near the unit circle, (|\lambda| \approx 1)), fast-dissipating information (eigenvalues (|\lambda| < 1)), and residual information (no spectral counterpart), connecting each component to temporal coherence, predictive sufficiency, and structural consistency.

3. An information-theoretic Lagrangian and a tractable loss. They propose the objective (\max_z \alpha \log I(z_{t-n};z_t) - \beta I(z_t; x_t | z_{t-n}) + \gamma S(\mathcal{C}/\mathrm{tr}(\mathcal{C})) + \log p(x_t|z_t)), showing that maximizing latent mutual information yields a water-filling allocation that can degenerate into low-rank mode collapse, while penalizing the von Neumann entropy of the normalized latent covariance (measured within a minibatch of size (B)) acts as an anti-collapse regularizer. From this they derive an architecture-agnostic loss that works for both VAE and autoencoder structures, with implementations ranging from a closed-form expression for low-dimensional latents to InfoNCE for high-dimensional settings.

4. Empirical validation across three families of dynamical problems. They test on physical simulations (Lorenz 63, Kármán vortex street, Dam flow, ERA5 weather forecasting), visual-input control (Planar, Pendulum, Cartpole, 3-Link manipulator), and graph-structured dynamics (Rope, Soft Robotics), reporting improvements over VAE, KAE, KKR, PFNN, E2C, PCC, and CKO baselines.

Main Findings

• Latent mutual information governs the predictive limit. Consistent with Proposition 3.1, prediction error accumulates and is bounded by the information gap; regularizing with latent mutual information improved both short- and long-term prediction across the physical simulation tasks. On Lorenz 63, the method achieved 5-NRMSE of 0.003 (0.002) vs. 0.005 (0.002) for VAE, 0.006 (0.003) for KAE, 0.004 (0.002) for KKR, and 0.005 (0.003) for PFNN; at 50 steps it reached 0.013 (0.008) against 0.019 (0.011), 0.023 (0.013), 0.017 (0.009), and 0.017 (0.007) respectively, with a KLD of 0.285.

• Strong performance on chaotic and high-dimensional fluid tasks. On the Kármán vortex street ((n = 64 \times 64 \times 2)), the method recorded 5-NRMSE 0.068 (0.006), 50-NRMSE 0.138 (0.018), 5-SSIM 0.936 (0.025), and the lowest spectral distribution error (SDE) of 0.256; PFNN, the domain-specific chaos baseline, recorded 5-NRMSE 0.075 (0.007), 50-NRMSE 0.137 (0.015), 5-SSIM 0.920 (0.030), and SDE 0.278.

• Improvements on Dam flow and ERA5 weather. For Dam flow ((n = 64 \times 64 \times 2), PFNN not evaluated as it targets chaotic dynamics), the method achieved 5-NRMSE 0.018 (0.001), 50-NRMSE 0.026 (0.003), 5-SSIM 0.760 (0.012), 50-SSIM 0.577 (0.006), and SDE 0.244. On ERA5 weather (channel average), it reported 5-NRMSE 0.028, 10-NRMSE 0.035, 50-NRMSE 0.068, and SSIM of 0.867, 0.808, and 0.781 at 5, 10, and 50 steps.

• Spectral collapse explains baseline failure. In the Kármán vortex analysis (Figure 2), the authors observe that the method's Koopman eigenvalues cluster near the unit circle, matching the limit-cycle structure, whereas KAE exhibits eigenvalues collapsing toward zero, reducing effective latent dimension and producing drift in autoregressive prediction. KKR and PFNN capture limit-cycle structure via one-step reconstruction but deviate over long horizons. Penalizing von Neumann entropy is credited with preventing eigenvalue degeneration.

• Latent manifold quality. In the Planar visual task (Figure 3, manifolds visualized with locally linear embedding), the method's latent manifold aligned most closely with ground truth; E2C's latent geometry was heavily distorted, KAE's manifold collapsed to nearly one-dimensional, and PCC showed partial improvement but still lagged. The authors report the best average control performance in both noiseless and noisy environments (Tables 8 and 9 in Appendix G.5.2; specific numbers are not included in the provided text).

• Robustness claims. The paper states the method remains stable under noisy observations and extended prediction horizons, and demonstrates continuous humidity predictions from 2018-01-01-00:00 to 2018-01-08-08:00 (Figure 4) with error maps described as more stable and accurate than comparisons. The graph-structured dynamics comparison against CKO is listed as a task, but the corresponding result table is not included in the provided content, so those numbers cannot be reported here.

Methodology in Plain English

The authors begin from the observation that a Koopman model compresses a state (x_{n-1}) into a latent (z_{n-1}), evolves the latent linearly through the Koopman operator, and decodes back to (x_n). They write this pipeline as a probabilistic trajectory distribution and then use standard information-theoretic inequalities to show that information can only decrease along the pathway — each step imposes a ceiling on what the latent can carry about the future state.

Next, they ask which information is being preserved. By decomposing the latent mutual information, they show that the stable, slowly decaying part of the dynamics maps onto Koopman eigenvalues near the unit circle, transient short-term dependencies map onto eigenvalues inside the unit circle, and noise-like residuals have no spectral interpretation at all and are therefore safe to discard. This tells them what a good latent should keep and what it should throw away.

The resulting objective has three competing pieces: reward latent mutual information (keeps temporal coherence), penalize the conditional mutual information between the latent and the current state given the past (compresses the dissipative and noisy parts), and reward the von Neumann entropy of the normalized latent covariance (prevents the latent from collapsing onto a few directions). They show via a water-filling argument that the first term alone concentrates variance on the largest eigenvalues of a specific matrix derived from the latent covariance and the Koopman operator, and that the entropy term restores a nonzero allocation across all modes.

Finally, they convert this Lagrangian into a differentiable training loss with five recognizable pieces — a temporal-coherence term, a structural-consistency term, an encoder-entropy term, a reconstruction/decoder term, and a von Neumann entropy term — plus an ELBO for stability, and evaluate on physical, visual, and graph-based benchmarks.

Why This Matters

Impact on research. The paper reframes a long-standing practical problem — why deep Koopman models are unstable or collapse — as a tractable information-theoretic trade-off, and it links each abstract quantity (mutual information, von Neumann entropy) to concrete spectral properties of the learned operator. It also supplies a general principle that is architecture-agnostic, rather than relying on the domain-specific priors (symmetry, conservation laws, dissipation, ergodicity) that earlier work used. The code is released publicly at https://github.com/Wenxuan52/InformationKoopman.

Real-world applications implied by the tasks evaluated:

• Weather and climate forecasting, via the ERA5 global humidity prediction task, where long-horizon stability is essential.
• Fluid dynamics and flow modeling, via the Kármán vortex street and Dam flow simulations.
• Robot and controller learning from pixels, via image-based Planar, Pendulum, Cartpole, and 3-Link manipulator control.
• Soft robotics and deformable-object dynamics, via the graph-structured Rope and Soft Robotics tasks.

Industry relevance. Any sector that depends on surrogate models of nonlinear time-evolving systems — energy and utility forecasting through the EDF-affiliated collaboration, robotics, aerospace, and simulation-driven engineering — could use a more stable and interpretable latent linear model. The claim of improved latent-space control success rates under noise is particularly relevant for deployed controllers, where sensor noise is unavoidable.

Future Directions

• Scaling the mutual-information estimator. The closed-form latent mutual information is only practical for low-dimensional latents; high-dimensional settings currently require InfoNCE approximations. How well the theoretical guarantees survive that approximation is an open question.
• Extending the spectral analysis beyond the three identified components. The residual information has no spectral counterpart in the Koopman operator, and the paper treats it as compressible; whether structured non-spectral information ever needs to be retained is unaddressed.
• Graph-structured dynamics results. The comparison against the CKO baseline is stated as a core experimental setting, but the corresponding quantitative table is not included in the provided content; a complete reporting of those results would be a natural next step.
• Choosing the Lagrangian multipliers. The trade-off is governed by (\alpha), (\beta), and (\gamma), but the provided content does not describe a systematic procedure for selecting them across the very different system types tested, which leaves a practical tuning question open.

Target Audience

This paper is best suited to machine learning and applied mathematics researchers working on Koopman operator methods, latent-variable representation learning, or information-theoretic objectives for sequence models. It will also be useful to graduate students and practitioners in scientific machine learning, control theory, and computational physics who need stable, interpretable latent dynamics models for high-dimensional systems. Readers without a background in information theory or spectral operator theory will find the theoretical sections demanding, though the experimental results and the loss-function decomposition are accessible to a broader audience.

Authors’ abstract

The Koopman operator provides a powerful framework for modeling dynamical systems and has attracted growing interest from the machine learning community. However, its infinite-dimensional nature makes identifying suitable finite-dimensional subspaces challenging, especially for deep architectures. We argue that these difficulties come from suboptimal representation learning, where latent variables fail to balance expressivity and simplicity. This tension is closely related to the information bottleneck (IB) dilemma: constructing compressed representations that are both compact and predictive. Rethinking Koopman learning through this lens, we demonstrate that latent mutual information promotes simplicity, yet an overemphasis on simplicity may cause latent space to collapse onto a few dominant modes. In contrast, expressiveness is sustained by the von Neumann entropy, which prevents such collapse and encourages mode diversity. This insight leads us to propose an information-theoretic Lagrangian formulation that explicitly balances this tradeoff. Furthermore, we propose a new algorithm based on the Lagrangian formulation that encourages both simplicity and expressiveness, leading to a stable and interpretable Koopman representation. Beyond quantitative evaluations, we further visualize the learned manifolds under our representations, observing empirical results consistent with our theoretical predictions. Finally, we validate our approach across a diverse range of dynamical systems, demonstrating improved performance over existing Koopman learning methods. The implementation is publicly available at https://github.com/Wenxuan52/InformationKoopman.