On the Generalisation of Koopman Representations for Chaotic System Control

TL;DR

This paper explores the transferability of Koopman-based representations for chaotic systems, demonstrating their effectiveness in prediction and control tasks through a three-stage learning process using the Lorenz system.

Contribution

It introduces a novel three-stage methodology for learning Koopman embeddings that are reusable across different tasks, outperforming traditional baselines.

Findings

Koopman embeddings outperform PCA baselines in accuracy and data efficiency.

Pre-trained transformer weights can be fixed during fine-tuning without performance loss.

Koopman embeddings enable multi-task learning in physics-informed machine learning.

Abstract

This paper investigates the generalisability of Koopman-based representations for chaotic dynamical systems, focusing on their transferability across prediction and control tasks. Using the Lorenz system as a testbed, we propose a three-stage methodology: learning Koopman embeddings through autoencoding, pre-training a transformer on next-state prediction, and fine-tuning for safety-critical control. Our results show that Koopman embeddings outperform both standard and physics-informed PCA baselines, achieving accurate and data-efficient performance. Notably, fixing the pre-trained transformer weights during fine-tuning leads to no performance degradation, indicating that the learned representations capture reusable dynamical structure rather than task-specific patterns. These findings support the use of Koopman embeddings as a foundation for multi-task learning in physics-informed machine learning. A project page is available at https://kikisprdx.github.io/.