# Learning the Koopman Operator using Attention Free Transformers

arXiv: https://arxiv.org/abs/2606.23957 · HTML: https://mlsanity.com/paper/2606.23957/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 50 references that are in the map.
- First posted on arXiv 22 Jun 2026.

## Builds on (most cited first, 19 of 19)

- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Deep learning for universal linear embeddings of nonlinear dynamics](https://mlsanity.com/paper/1712.09707/) (2017)
- [Linearly-Recurrent Autoencoder Networks for Learning Dynamics](https://mlsanity.com/paper/1712.01378/) (2017)
- [An Attention Free Transformer](https://mlsanity.com/paper/2105.14103/) (2021)
- [On Convergence of Extended Dynamic Mode Decomposition to the Koopman Operator](https://mlsanity.com/paper/1703.04680/) (2017)
- [Generalizing Koopman Theory to allow for inputs and control](https://mlsanity.com/paper/1602.07647/) (2016)
- [Physics-Informed Probabilistic Learning of Linear Embeddings of Non-linear Dynamics With Guaranteed Stability](https://mlsanity.com/paper/1906.03663/) (2019)
- [Rigorous data-driven computation of spectral properties of Koopman operators for dynamical systems](https://mlsanity.com/paper/2111.14889/) (2021)
- [Koopman Neural Forecaster for Time Series with Temporal Distribution Shifts](https://mlsanity.com/paper/2210.03675/) (2022)
- [Deep Learning Enhanced Dynamic Mode Decomposition](https://mlsanity.com/paper/2108.04433/) (2021)
- [Temporally Consistent Koopman Autoencoders for Forecasting Dynamical Systems](https://mlsanity.com/paper/2403.12335/) (2024)
- [Koopman Operators in Robot Learning](https://mlsanity.com/paper/2408.04200/) (2024)
- [Dynamic mode decomposition for multiscale nonlinear physics](https://mlsanity.com/paper/1903.12480/) (2019)
- [Augmented Invertible Koopman Autoencoder for long-term time series forecasting](https://mlsanity.com/paper/2503.12930/) (2025)
- [An Optimal Control Formulation of Pulse-Based Control Using Koopman Operator](https://mlsanity.com/paper/1707.08462/) (2017)
- [Data-driven approximations of dynamical systems operators for control](https://mlsanity.com/paper/1902.10239/) (2019)
- [Course Correcting Koopman Representations](https://mlsanity.com/paper/2310.15386/) (2023)
- [Data-Driven Observability Decomposition with Koopman Operators for Optimization of Output Functions of Nonlinear Systems](https://mlsanity.com/paper/2210.09343/) (2022)
- [Loss Terms and Operator Forms of Koopman Autoencoders](https://mlsanity.com/paper/2412.04578/) (2024)

## Abstract (arXiv)

Learning Koopman operators with autoencoders enables linear prediction in a latent space, but long-horizon rollouts often drift off the learned manifold, leading to phase and amplitude errors on systems with switching, continuous spectra, or strong transients. We introduce two complementary components that make Koopman predictors more robust. First, we add an attention-free latent memory (AFT) block that aggregates a short window of past latents to produce a corrected latent before each Koopman update. Unlike multi-head attention, AFT operates in linear time and adds only ≈30k parameters (3d2 + T2, fewer than matched multi-head attention), yet captures the local temporal context needed to suppress error divergence. Second, we propose dynamic re-encoding: lightweight, online change-point triggers (EWMA, CUSUM, and sequential two-sample tests) that detect latent drift and project predictions back onto the autoencoder manifold. Across three benchmark systems -- Duffing oscillator, Repressilator, IRMA -- our model consistently reduces error accumulation compared to a Koopman autoencoder and matched-capacity multi-head attention. We also compare against GRU and Transformer autoencoders, evaluated both from initial conditions and with a 50-step context, and find that Koopman+AFT (with optional re-encoding) attains markedly lower long-horizon error while maintaining lower inference latency. We report improvements over horizons up to 1000 steps, together with ablations over trigger policies. The result is a fast, compact predictor that stays on the learned manifold over long horizons.