Koopman Learning with Episodic Memory

TL;DR

This work addresses predicting non-autonomous dynamical systems by enriching Koopman learning with episodic memory. It introduces a practical online framework (EDMD with memory) that stores spectral objects from local time windows and recalls past episodes using Wasserstein distances on eigenvalues and Euclidean distances on modes to improve future predictions, selecting the most relevant past window when a match is found. Empirical results on synthetic data, U.S. influenza cases, and Washington, D.C. bike-share data show substantial predictive gains versus standard sliding EDMD, including up to $\\sim650\\%$ improvements in synthetic settings and notable gains in real-world forecasting horizons. The approach is online, computationally efficient, and highlights multiple avenues for expansion, such as uncertainty quantification, additional observables, and integration with broader KMD and control methods, potentially broadening the applicability of Koopman-based learning for non-stationary dynamics.

Abstract

Koopman operator theory has found significant success in learning models of complex, real-world dynamical systems, enabling prediction and control. The greater interpretability and lower computational costs of these models, compared to traditional machine learning methodologies, make Koopman learning an especially appealing approach. Despite this, little work has been performed on endowing Koopman learning with the ability to leverage its own failures. To address this, we equip Koopman methods -- developed for predicting non-autonomous time-series -- with an episodic memory mechanism, enabling global recall of (or attention to) periods in time where similar dynamics previously occurred. We find that a basic implementation of Koopman learning with episodic memory leads to significant improvements in prediction on synthetic and real-world data. Our framework has considerable potential for expansion, allowing for future advances, and opens exciting new directions for Koopman learning.

Figures (7)

• Figure 1: Example of repeated structure in non-autonomous dynamical system. (A) Number of flu cases per week in the U.S., from 1997 to 2022, as reported to CDC by the Public Health Laboratories. (B) Zoomed-in four year time windows of (A). Years are colored by their qualitative similarities in shape and peak amplitude. (C) Predicting the future number of cases at a given time point (denoted by the red circle), using a temporally local window (denoted by an orange box), can lead to poor predictions (compare green and gray lines). (D) However, if a similar dynamical regime has occurred in the past, "remembering" this and using it for prediction can lead to improved performance (compare green and gray lines).
• Figure 2: Illustration of sliding EDMD. Data (black line in A) from a temporally local window (orange box) is lifted (B), via a dictionary of functions (Eq. \ref{['eq: EDMD']}---e.g., Gaussian radial basis functions). This lifting enables the approximation of a temporally local Koopman mode decomposition (C), which can then be used to make predictions $\Delta$ time steps into the future (red line in A).
• Figure 3: Overview of episodic memory for Koopman learning. For the current state of the system under study, the data within the sliding window (inset figure on left) is used to compute the KMD. The resulting Koopman eigenvalues and modes are then compared to a memory bank of saved Koopman eigenvalues and modes, associated with past data. If a match of both Koopman eigenvalues and modes is found (illustrated by the red box), then the data that occurred in the past (dashed red line) is used to predict the future behavior of the system (dashed green line). If no match is found, sliding EDMD is used to generate a prediction.
• Figure 4: Episodic memory improves the ability of Koopman learning to predict synthetic non-autonomous exponential data. (A) The predictions of sliding EDMD $\Delta = 1, ..., 5$ time steps ahead (red dots) often poorly match the true data (black line). (B) Utilizing EDMD memory (blue line) leads to $650\%$ improvement on predictions, as compared to sliding EDMD (red line). Note that although a zoomed-in section is shown here, the improvement was computed over a $1000$ time-step simulation. Here $\Delta = 5$, $\omega = 5$, $n_\text{delays} = 1$, $\varepsilon_\lambda = 0.05$, and $\varepsilon_v = 0.10$. (C) Because there are a finite (and small) number of possible values $\lambda_t$ can take, there is repetition of transitions (examples highlighted by light green/dark green and orange/dark red rectangles). (D) These repetitions can be identified by characteristic Koopman spectral objects [colored the same as in (C)], demonstrating how EDMD with memory is able to successfully recall relevant past data.
• Figure 5: Episodic memory improves the ability of Koopman learning to predict U.S. flu cases. (A)--(B) Predicted weekly number of U.S. flu cases one and five weeks ahead ($\Delta = 1$ and $5$, respectively), using sliding EDMD with (blue line) and without (red line) episodic memory. Black line denotes true data. (C)--(D) Zoomed-in windows of (A) and (B). For all subplots, $\omega = 3$, $n_\text{delays} = 4$, $\varepsilon_\lambda = 0.05$, and $\varepsilon_v = 0.25$.