Sparse Koopman Autoencoders Identify Local Dynamical Regimes in Multibasin Systems

Aidan Li; Uday Kiran Reddy Tadipatri; Mahan Fathi; Sarath Chandar; Ross Goroshin

cs.LGAug 29, 2026

Sparse Koopman Autoencoders Identify Local Dynamical Regimes in Multibasin Systems

Authors: Aidan Li, Uday Kiran Reddy Tadipatri, Mahan Fathi, Sarath Chandar, Ross Goroshin

Organizations: Chandar Research Lab · Mila – Quebec AI Institute · Université de Montréal · University of Pennsylvania ·

Abstract

Koopman autoencoders (KAEs) seek a higher-dimensional latent representation in which nonlinear dynamics evolve linearly. However, many interesting systems have multiple basins of attraction, and both theoretical and empirical work has shown these multibasin systems cannot generally admit a single finite-dimensional global Koopman embedding under standard assumptions. We posit that encoders with a sparsity-inducing objective encouraging few active latent coefficients will provide latent supports as an inspectable basin-modeling principle for Koopman autoencoders. We use these encoders producing sparse latents in training Sparse Koopman Autoencoders (SKAEs) without basin labels or other regime annotations, and treat the learned latent supports as model-produced regime variables after training. Across a range of procedurally generated multibasin systems and chaotic flows, we show that SKAEs have superior forecasting performance compared to dense-latent KAEs. We also perform a mechanistic study that shows latent supports produced by SKAEs are both essential for the quality of the representation and useful for identifying basins on basin-interior states, whereas dense-latent KAEs collapse to an uninformative single support family. These results identify sparse latents and their corresponding supports as label-free, interpretable regime variables for Koopman learning in nonlinear systems with multiple local dynamical laws.

Figures & tables

Figure 1: Top: procedurally generated multibasin systems where gray streamlines are the true vector field, colored dots mark support-family/basin matches, and black dots mark support-family/basin mismatches. The learned LISTA (seed 0) support families are mapped post hoc to each family’s dominant validation basin, which colors each dot. Bottom: Dysts phase portraits of $20$ test starts forecast for $5{,}000$ observation steps, with held-out truth in gray and seed-0 LISTA forecasts in blue.

	Forecasting MSE $\downarrow$						Support diagnostics
	Multibasin, 15 systems			Dysts, 10 systems			Multibasin, 15 systems
Model	H100	H500	H1000	H100	H2000	H4000	$H(B\!\mid\!F_{\rm abs})\downarrow$	$\overline{\|F_{\rm abs}\|}$
LISTA	${0.0292}^{\ast\ast}$	${0.0895}^{\ast\ast}$	${0.106}^{\ast\ast}$	${\mathbf{7.59{\times}10^{-5}}}\,{}^{\ast\ast}$	${\mathbf{0.339}}^{\ast\ast}$	${\mathbf{1.27}}^{\ast\ast}$	${0.0348}^{\ast\ast}$	$4.3$
LISTA-BD	${0.0359}^{\ast\ast}$	${0.109}^{\ast\ast}$	${0.125}^{\ast\ast}$	${8.64{\times}10^{-5}}\,{}^{\ast\ast}$	${0.341}^{\ast\ast}$	${1.29}^{\ast\ast}$	${0.0560}^{\ast\ast}$	$4.2$
LISTA-SB	${0.0348}^{\ast\ast}$	${0.102}^{\ast\ast}$	${0.119}^{\ast\ast}$	${8.47{\times}10^{-5}}\,{}^{\ast\ast}$	${0.359}^{\ast\ast}$	${1.30}^{\ast\ast}$	${\mathbf{0.0326}}^{\ast\ast}$	$4.3$
Sparse MLP-BD	${0.0283}^{\ast\ast}$	${0.0881}^{\ast\ast}$	${0.104}^{\ast\ast}$	${2.15{\times}10^{-4}}\,{}^{\ast\ast}$	${0.467}^{\ast\ast}$	${1.53}^{\ast\ast}$	${0.102}^{\ast\ast}$	$3.9$

Table 1: Forecasting performance and basin-support diagnostics.

H(B\mid F_{\rm abs})

is the entropy of the basin label given the support family and

\overline{|F_{\rm abs}|}

is the number of support families. Values are means over systems of per-system seed IQMs. Direct baselines lift the state once whereas re-lifted ones re-lift it every

m

steps (all neural models use periodic re-encoding).

\ast\ast

: lower than Dense MLP at

p<0.05

, Holm-corrected, in every system.

\ast

denotes the same significance in all systems but one. Bold : best in each

\downarrow

column.

Panel of Figure 2 — Figure 2: Forecasting MSE at the horizons of Table 1 , one panel per horizon. Sparse models forecast better than dense models. Markers show arithmetic means across per-system seed-IQM MSEs on a log scale. Shaded bars show fixed-system seed-bootstrap 95% intervals after system-wise log-relative normalization, anchored to the displayed mean, so they summarize seed uncertainty within an average system.

Figure 3: Forecast error over time on the ten Dysts systems, for the re-encoded forecasts of Table 1 . At every step, we take the median squared error over the $100$ test starts. The lines show the median over the $15$ seeds and bands the interquartile range. Dotted lines mark the Table 1 horizons $H=2000$ and $H=4000$ . Each system’s dominant period is about $1{,}000$ steps. In every system, LISTA’s and Sparse MLP’s errors are below Dense MLP’s from step $304$ to step $2{,}000$ .

Panel of Figure 4 — Figure 4: Support-coordinate interventions on a multibasin system. Interventions severely increase forecasting errors compared to the standard rollout. (a) The standard rollout is compared with cumulative dropping of the largest active coordinates of the initial sparse code. Shaded bands are bootstrap 95% intervals over $100$ initial conditions. (b) Active coefficient values are preserved but reassigned to randomly chosen inactive coordinates. The band is the interquartile range over shuffles and initial conditions ( $20$ shuffles each).

Rollout	Mean $\pm$ SD	Median [IQR]
Standard	$\mathbf{0.0112{\pm}0.0140}$	$\mathbf{0.00735\,[0.00488,0.0122]}$
Drop top-1	$0.670{\pm}0.639$	$0.542\,[0.189,0.700]$
Drop top-2	$1.02{\pm}0.625$	$0.922\,[0.605,1.19]$
Drop top-3	$2.04{\pm}1.49$	$1.67\,[1.30,2.02]$
Drop top-5	$4.07{\pm}3.82$	$2.33\,[1.77,4.48]$
Drop top-10	$10.5{\pm}5.67$	$8.41\,[6.22,13.3]$

Table 2: Accumulated MSE at

H=21

for the support interventions (lower is better).

Panel of Figure 5 — Figure 5: Support shuffling and support–basin alignment on the multibasin benchmark. (a) Rollouts after randomly shuffling the active coordinates (orange) move away from the standard rollouts (black). (b) Conditional entropy $H(B\mid F_{\rm abs})$ of the basin label given the support family: the support families of the sparse models identify basins, whereas those of Dense MLP do not. Dots show system–seed pairs, gray I-bars the first-to-third quartile range of the per-system seed IQMs, and black bars the values of Table 1 .

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Encoder	Code sparsity	Latent transition
LISTA	LISTA, ReLU-MLP pre-code	shrinkage, $\ell_{1}$	dense
LISTA-BD	LISTA, ReLU-MLP pre-code	shrinkage, $\ell_{1}$	block-diagonal
LISTA-SB	LISTA, ReLU-MLP pre-code	shrinkage, $\ell_{1}$	dense, off-block penalty
Sparse MLP	MLP, ReLU hidden and output	ReLU, $\ell_{1}$	dense
Sparse MLP-BD	MLP, ReLU hidden and output	ReLU, $\ell_{1}$	block-diagonal
Dense MLP	MLP, tanh hidden, linear output	none	dense

Appendix

Table 3: Neural models. The last two rows are the encoder ablations of Section D.2 .

Architecture
	Multibasin	Dysts
LISTA code	$128$ coordinates, split by sign into $256$ ; decoder initialized as $[D_{0},-D_{0}]$	$256$ coordinates, nonnegative shrinkage $\max(a-\tau,0)$ ; decoder initialized as $D_{0}$
LISTA threshold $\tau$	fixed, $0.15/\rho$	learned per coordinate, initialized at $0.15/\rho$ , shared by the refinements
BD and SB blocks	one per basin, contiguous, near-equal size	two of $128$ coordinates
Training
Seeds	$0,\ldots,14$	$0,\ldots,14$

Appendix

Table 4: Settings of the neural models. The encoder and decoder learning rate also applies to the learned LISTA thresholds on Dysts.

	Multibasin	Dysts
Training data	the $2{,}048$ windows of the neural run with the same seed	the full training split
Fits	one per system and seed	DMD and polynomial EDMD: one per system, so their seeds replicate only the test starts; RBF EDMD: one per system and seed
Selection starts	validation starts of the same seed	validation starts of a reserved seed ( $9000$ ) that enters no reported result, per system
Selection criterion	lowest validation MSE at $H=500$ of the row’s own rollout, non-finite last; ties to fewer features, then the larger ridge, then the shorter period	lowest finite validation MSE at $H=500$ of the row’s own rollout
Ridge $\lambda$	$0$ , $10^{-10}$ , $10^{-8}$ , $10^{-6}$ , $10^{-4}$ , $10^{-2}$ , $10^{-1}$ , $1$ , $10$ ; added to $\Phi^{\top}\Phi/n$	$10^{-12}$ , $10^{-10}$ , …, $10^{-2}$ ; applied as $\lambda\,\mathrm{tr}(G)/p$ for the Gram matrix $G$ of the $p$ features
Period $m$ (re-lifted)	$10$ , $25$ , $50$ , $100$	$10$ , $25$ , $50$ , $100$ , $150$ , $200$

Appendix

Table 5: Training data and selection of the classical baselines. The ridge and, for re-lifted rollouts, the period

m

are selected jointly, as is

N

for RBF EDMD. The Dysts grids were fixed before the selection on the ten systems; the last row counts the selections at an edge of the grid.

Forecasting MSE $\downarrow$
	Dense MLP	Dense MLP,	Sparse MLP,
	(tanh hidden)	ReLU hidden	$\lambda_{\rm sp}=0$	Sparse MLP §[1/15]	LISTA
$H=100$	$0.654$	${0.0541}^{\ast\ast}$	${0.0347}^{\ast\ast}$	${0.0269}^{\ast\ast}$	${0.0292}^{\ast\ast}$
$H=500$	$2.54$	${0.196}^{\ast\ast}$	${0.0997}^{\ast\ast}$	${0.0784}^{\ast\ast}$	${0.0895}^{\ast\ast}$
$H=1000$	$2.91$	${0.229}^{\ast\ast}$	${0.116}^{\ast\ast}$	${0.0897}^{\ast\ast}$	${0.106}^{\ast\ast}$
Systems with significantly lower / higher MSE than the ReLU-hidden Dense MLP ( $H=100/500/1000$ )

Appendix

Table 6: Sparsity mechanisms of the encoder on the multibasin benchmark (

15

systems,

15

paired seeds; tests of Appendix C ). MSE: means over systems of per-system seed IQMs;

\ast\ast

: significantly lower than Dense MLP in all

15

systems;

\lx@sectionsign[1/15]

: significantly lower than Sparse MLP with

\lambda_{\rm sp}=0

1

system at every horizon. Sparsity rows, as means over systems on

64

held-out trajectories per run: exactly zero encoded coordinates; coordinates with

|z_{i}|\leq 10^{-3}

in the direct

K^{h}

rollout over

8

steps; and exactly zero activations in each hidden layer of the encoder MLP (for LISTA, its pre-code MLP).

Model	(b) $<$ (a)	(c) $<$ (b)	(c) $<$ (a)	(c)/(b)	(b) $<$ (a)	(c) $<$ (b)
	Cells with lower test MSE			Ratio	(b) at its own period
LISTA	204/225	224/225	225/225	$0.022$	208/225	224/225
Sparse MLP	205/225	225/225	225/225	$0.015$	208/225	225/225
Dense MLP	213/225	192/225	225/225	$0.72$	223/225	175/225

Appendix

Table 7: Re-encoding ablation on the multibasin benchmark at

H=100

: (a) no re-encoding, (b) support-constrained, and (c) unrestricted re-encoding. Cells: system–seed cells in which the first condition has strictly lower test MSE than the second (there are no ties). Ratio: geometric mean over the

225

cells of the test-MSE ratio of (c) to (b). Own period: (b) at the period that minimizes its own validation error.

		Direct rollout			Re-encoded (reported)
Model	Window	H100	H500	H1000	H100	H500	H1000
LISTA	$9$ states	$14.3$	$64.9$	$40.1$	$0.0273$	$0.0843$	$0.098$
	$64$ states	$0.027$	$2.16$	$25$	$0.0232$	$0.0704$	$0.0825$
Sparse MLP	$9$ states	$16.3$	$46.8$	$25.3$	$0.0305$	$0.0804$	$0.0916$
	$64$ states	$0.0234$	$1.23$	$14.4$	$0.0217$	$0.0653$	$0.0775$
Dense MLP	$9$ states	$21.6$	$99.7$	$75.1$	$0.665$	$2.53$	$2.82$

Appendix

Table 8: Training-window ablation on the multibasin benchmark (

15

systems, seeds

0,\ldots,4

, paired with the main runs of the same seeds, so the

9

-state rows differ from the

15

-seed values of Table 1 ). Test MSE as means over systems of per-system seed IQMs, for the direct rollout and the reported re-encoded forecast; lower is better.

System	Kept	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Local-linear gates	$0.224$	15	15	15	15	15
Transfer-gated local-linear	$0.202$	15	13	15	15	11
Arrested spiral	$0.349$	2	1	0	0	0
Asymmetric three-well	$0.323$	15	15	15	15	15
High-cross three-well	$0.255$	15	15	15	15	15
Hexagonal six-well	$0.255$	14	10	11	7	10

Appendix

Table 9: Exact one-family-per-basin matches on the high-margin subset: seeds of each system, out of

15

(last row: system–seed cells, out of

225

), whose test support families correspond one-to-one to the basins. Kept: mean fraction of the test states in the high-margin subset, the same for every model.

	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP	Dense MLP
$H(B\mid F_{\rm abs})$ (nats)	$0.0348$	$0.0560$	$0.0326$	$0.102$	$0.0679$	$1.31$
$H(B\mid F_{\rm abs})/H(B)$	$0.0221$	$0.0375$	$0.0206$	$0.0740$	$0.0470$	$1.00$

Appendix

Table 10: Conditional entropy of the basin label given the support family on the high-margin subset, unnormalized (as in Table 1 ) and divided by the basin-label entropy

H(B)

: mean over the

15

systems of per-system seed IQMs; lower is better.

		Jaccard threshold
Model	Absolute threshold	$0.3$	$0.4$	$0.5$	$0.6$	$0.7$
LISTA	$10^{-4}$	$0.258$ ( $99$ )	$0.107$ ( $149$ )	$0.035$ ( $169$ )	$0.005$ ( $194$ )	$0.001$ ( $177$ )
	$10^{-3}$	$0.200$ ( $114$ )	$0.092$ ( $151$ )	$0.032$ ( $175$ )	$0.004$ ( $195$ )	$0.000$ ( $181$ )
	$10^{-2}$	$0.127$ ( $134$ )	$0.056$ ( $156$ )	$0.018$ ( $178$ )	$0.006$ ( $193$ )	$0.001$ ( $182$ )
Sparse MLP	$10^{-4}$	$0.327$ ( $70$ )	$0.130$ ( $135$ )	$0.073$ ( $153$ )	$0.020$ ( $190$ )	$0.005$ ( $192$ )
	$10^{-3}$	$0.272$ ( $91$ )	$0.109$ ( $142$ )	$0.055$ ( $158$ )	$0.014$ ( $193$ )	$0.002$ ( $192$ )

Appendix

Table 11: Threshold sensitivity of the support–basin alignment on the high-margin subset. Each cell gives the mean of

H(B\mid F_{\rm abs})/H(B)

over the

225

system–seed cells (lower is better) and, in parentheses, the number of exact matches. Bold: the paper’s setting; its entropy is a mean over cells, not the mean of per-system seed IQMs of Table 10 .

System	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Local-linear gates	${9.78{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${8.28{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${9.41{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${4.45{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${5.38{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$
Transfer-gated local-linear	${0.235}^{\ast\dagger\ddagger}$	${0.271}^{\ast\dagger\ddagger}$	${0.236}^{\ast\dagger\ddagger}$	${0.218}^{\ast\dagger\ddagger}$	${0.223}^{\ast\dagger\ddagger}$
Arrested spiral	${0.0122}^{\ast\dagger\ddagger}$	${0.00927}^{\ast\dagger\ddagger}$	${0.0147}^{\ast\dagger\ddagger}$	${0.0125}^{\ast\dagger\ddagger}$	${0.0124}^{\ast\dagger\ddagger}$
Asymmetric three-well	${0.00814}^{\ast\dagger\ddagger}$	${0.0266}^{\ast\dagger\ddagger}$	${0.0157}^{\ast\dagger\ddagger}$	${0.00751}^{\ast\dagger\ddagger}$	${0.0112}^{\ast\dagger\ddagger}$
High-cross three-well	${0.0835}^{\ast\dagger\ddagger}$	${0.0865}^{\ast\dagger\ddagger}$	${0.105}^{\ast\dagger\ddagger}$	${0.0709}^{\ast\dagger\ddagger}$	${0.0771}^{\ast\dagger\ddagger}$
Hexagonal six-well	${0.0210}^{\ast\dagger\ddagger}$	${0.0245}^{\ast\dagger\ddagger}$	${0.0293}^{\ast\dagger\ddagger}$	${0.0235}^{\ast\dagger\ddagger}$	${0.0143}^{\ast\dagger\ddagger}$

Appendix

Table 12: Multibasin test MSE per system at

H=100

: IQM over

15

seeds; lower is better. Top: sparse models; bottom: baselines.

\ast

\dagger

\ddagger

: significantly lower MSE than Dense MLP, DMD, or both polynomial EDMD rollouts in that system (Holm-corrected one-sided paired Wilcoxon test,

p<0.05

). Tests against RBF EDMD are not marked.

System	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Local-linear gates	${0.00154}^{\ast\dagger\ddagger}$	${0.00146}^{\ast\dagger\ddagger}$	${0.00109}^{\ast\dagger\ddagger}$	${5.81{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${8.28{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$
Transfer-gated local-linear	${0.745}^{\ast\dagger\ddagger}$	${0.862}^{\ast\dagger\ddagger}$	${0.751}^{\ast\dagger\ddagger}$	${0.664}^{\ast\dagger\ddagger}$	${0.662}^{\ast\dagger\ddagger}$
Arrested spiral	${0.0200}^{\ast\dagger\ddagger}$	${0.0168}^{\ast\dagger\ddagger}$	${0.0230}^{\ast\dagger\ddagger}$	${0.0200}^{\ast\dagger\ddagger}$	${0.0165}^{\ast\dagger\ddagger}$
Asymmetric three-well	${0.0304}^{\ast\dagger\ddagger}$	${0.0728}^{\ast\dagger\ddagger}$	${0.0511}^{\ast\dagger\ddagger}$	${0.0401}^{\ast\dagger\ddagger}$	${0.0403}^{\ast\dagger\ddagger}$
High-cross three-well	${0.156}^{\ast\dagger\ddagger}$	${0.157}^{\ast\dagger\ddagger}$	${0.212}^{\ast\dagger\ddagger}$	${0.146}^{\ast\dagger\ddagger}$	${0.155}^{\ast\dagger\ddagger}$
Hexagonal six-well	${0.0353}^{\ast\dagger\ddagger}$	${0.0395}^{\ast\dagger\ddagger}$	${0.0497}^{\ast\dagger\ddagger}$	${0.0392}^{\ast\dagger\ddagger}$	${0.0247}^{\ast\dagger\ddagger}$

Appendix

Table 13: Multibasin test MSE per system at

H=500

, as in Table 12 .

System	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Local-linear gates	${0.00167}^{\ast\dagger\ddagger}$	${0.00161}^{\ast\dagger\ddagger}$	${0.00110}^{\ast\dagger\ddagger}$	${6.34{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${8.73{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$
Transfer-gated local-linear	${0.811}^{\ast\dagger\ddagger}$	${0.913}^{\ast\dagger\ddagger}$	${0.824}^{\ast\dagger\ddagger}$	${0.724}^{\ast\dagger\ddagger}$	${0.711}^{\ast\dagger\ddagger}$
Arrested spiral	${0.0210}^{\ast\dagger\ddagger}$	${0.0175}^{\ast\dagger\ddagger}$	${0.0241}^{\ast\dagger\ddagger}$	${0.0209}^{\ast\dagger\ddagger}$	${0.0170}^{\ast\dagger\ddagger}$
Asymmetric three-well	${0.0332}^{\ast\dagger\ddagger}$	${0.0785}^{\ast\dagger\ddagger}$	${0.0555}^{\ast\dagger\ddagger}$	${0.0440}^{\ast\dagger\ddagger}$	${0.0440}^{\ast\dagger\ddagger}$
High-cross three-well	${0.166}^{\ast\dagger\ddagger}$	${0.166}^{\ast\dagger\ddagger}$	${0.225}^{\ast\dagger\ddagger}$	${0.155}^{\ast\dagger\ddagger}$	${0.165}^{\ast\dagger\ddagger}$
Hexagonal six-well	${0.0371}^{\ast\dagger\ddagger}$	${0.0413}^{\ast\dagger\ddagger}$	${0.0522}^{\ast\dagger\ddagger}$	${0.0411}^{\ast\dagger\ddagger}$	${0.0260}^{\ast\dagger\ddagger}$

Appendix

Table 14: Multibasin test MSE per system at

H=1000

, as in Table 12 .

System	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Chua	${1.64{\times}10^{-5}}\,{}^{\ast\dagger\ddagger}$	${1.71{\times}10^{-5}}\,{}^{\ast\dagger\ddagger}$	${1.74{\times}10^{-5}}\,{}^{\ast\dagger\ddagger}$	${3.19{\times}10^{-5}}\,{}^{\ast\dagger\ddagger}$	${3.86{\times}10^{-5}}\,{}^{\ast\dagger\ddagger}$
Dadras	${1.70{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${1.78{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${1.74{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${4.24{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$	${4.40{\times}10^{-4}}\,{}^{\ast\dagger\ddagger}$
Dequan Li	${5.24{\times}10^{-5}}\,{}^{\ast\dagger}$	${5.15{\times}10^{-5}}\,{}^{\ast\dagger}$	${5.65{\times}10^{-5}}\,{}^{\ast\dagger}$	${1.10{\times}10^{-4}}\,{}^{\ast\dagger}$	${1.50{\times}10^{-4}}\,{}^{\dagger}$
Hadley	${4.68{\times}10^{-5}}\,{}^{\ast\dagger}$	${4.79{\times}10^{-5}}\,{}^{\ast\dagger}$	${4.50{\times}10^{-5}}\,{}^{\ast\dagger}$	${1.06{\times}10^{-4}}\,{}^{\ast\dagger}$	${1.34{\times}10^{-4}}\,{}^{\ast\dagger}$
Lu–Chen–Cheng	${7.83{\times}10^{-6}}\,{}^{\ast\dagger}$	${8.72{\times}10^{-6}}\,{}^{\ast\dagger}$	${8.27{\times}10^{-6}}\,{}^{\ast\dagger}$	${2.59{\times}10^{-5}}\,{}^{\ast\dagger}$	${2.99{\times}10^{-5}}\,{}^{\ast\dagger}$
Qi–Chen	${2.33{\times}10^{-5}}\,{}^{\ast\dagger}$	${2.73{\times}10^{-5}}\,{}^{\ast\dagger}$	${2.64{\times}10^{-5}}\,{}^{\ast\dagger}$	${7.43{\times}10^{-5}}\,{}^{\ast\dagger}$	${8.97{\times}10^{-5}}\,{}^{\ast\dagger}$

Appendix

Table 15: Dysts test MSE per system at

H=100

observation steps, as in Table 12 .

System	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Chua	${0.00625}^{\ast\dagger\ddagger}$	${0.00738}^{\ast\dagger\ddagger}$	${0.0116}^{\ast\dagger\ddagger}$	${0.0181}^{\ast\dagger\ddagger}$	${0.0248}^{\ast\dagger\ddagger}$
Dadras	${0.777}^{\ast\dagger\ddagger}$	${0.803}^{\ast\dagger\ddagger}$	${0.927}^{\ast\dagger\ddagger}$	${1.16}^{\ast\dagger\ddagger}$	${1.26}^{\ast\dagger\ddagger}$
Dequan Li	${0.0505}^{\ast\dagger}$	${0.0446}^{\ast\dagger}$	${0.0559}^{\ast\dagger}$	${0.0781}^{\ast\dagger}$	${0.0826}^{\ast\dagger}$
Hadley	${0.0117}^{\ast\dagger}$	${0.0111}^{\ast\dagger}$	${0.0124}^{\ast\dagger}$	${0.0234}^{\ast\dagger}$	${0.0294}^{\ast\dagger}$
Lu–Chen–Cheng	${0.00140}^{\ast\dagger}$	${0.00176}^{\ast\dagger}$	${0.00170}^{\ast\dagger}$	${0.00490}^{\ast\dagger}$	${0.00629}^{\ast\dagger}$
Qi–Chen	${0.395}^{\ast\dagger}$	${0.329}^{\ast\dagger}$	${0.325}^{\ast\dagger}$	${0.456}^{\ast\dagger}$	${0.581}^{\ast\dagger}$

Appendix

Table 16: Dysts test MSE per system at

H=2000

observation steps, as in Table 12 .

System	LISTA	LISTA-BD	LISTA-SB	Sparse MLP-BD	Sparse MLP
Chua	${0.117}^{\ast\dagger\ddagger}$	${0.105}^{\ast\dagger\ddagger}$	${0.129}^{\ast\dagger\ddagger}$	${0.209}^{\ast\dagger\ddagger}$	${0.274}^{\ast\dagger\ddagger}$
Dadras	${2.50}^{\ast\dagger\ddagger}$	${2.54}^{\ast\dagger\ddagger}$	${2.58}^{\ast}$	${2.92}^{\ast}$	${3.03}^{\ast}$
Dequan Li	${0.127}^{\ast\dagger}$	${0.126}^{\ast\dagger}$	${0.148}^{\ast\dagger}$	${0.189}^{\ast\dagger}$	${0.233}^{\ast\dagger}$
Hadley	${0.108}^{\ast\dagger}$	${0.102}^{\ast\dagger}$	${0.121}^{\ast\dagger}$	${0.179}^{\ast\dagger}$	${0.216}^{\ast\dagger}$
Lu–Chen–Cheng	${0.00822}^{\ast\dagger}$	${0.00999}^{\ast\dagger}$	${0.00988}^{\ast\dagger}$	${0.0321}^{\ast\dagger}$	${0.0355}^{\ast\dagger}$
Qi–Chen	${2.03}^{\ast\dagger}$	${1.92}^{\ast\dagger}$	${1.95}^{\ast\dagger}$	${2.31}^{\ast\dagger}$	${2.46}^{\ast\dagger}$

Appendix

Table 17: Dysts test MSE per system at

H=4000

observation steps, as in Table 12 .

Display name and system key	Generator family	Parameters	$d$	$\Delta t$	Initial-condition box	Basins
Local-linear gates gated_local_linear	piecewise local-linear gates	$c_{b}=1.75(\cos\alpha_{b},\sin\alpha_{b})$ , core radius $1.05$ ; $\widetilde{A}_{b}$ , $G$ in text	2	0.04	$[-2.6,2.6]^{2}$	3
Transfer-gated local-linear gated_transfer_linear	piecewise transfer gates	$c_{b}=1.85(\cos\alpha_{b},\sin\alpha_{b})$ ; radii, widths, speeds, and templates in text	2	0.04	$[-2.8,2.8]^{2}$	3
Arrested spiral arrested_spiral	spiral capture with Gaussian traps	damping $0.3$ , rotation $2.0$ ; traps at $(1.5,0)$ , $(0,1.8)$ , $(-1,-1)$ , $(0.8,-1.5)$ , and the origin	2	0.02	$[-3,3]^{2}$	5
Asymmetric three-well cal_asymmetric_3	Gaussian wells with independent rotation	Table 19	2	0.03	$[-3,3]\times[-2.5,3]$	3
High-cross three-well cal_high_cross_3	Gaussian wells with stronger rotation	Table 19	2	0.03	$[-3.3,3.3]^{2}$	3
Hexagonal six-well cal_hexagon_6	Gaussian wells with independent rotation	Table 19	2	0.03	$[-3.3,3.3]^{2}$	6

Appendix

Table 18: The

15

multibasin benchmark systems. Every system has state dimension

d=2

. Each stored observation is one fourth-order Runge–Kutta step of size

\Delta t

\dot{x}=f_{s}(x)

from the previous stored state, and initial conditions are drawn uniformly from the listed box. Parameters of the Gaussian-well systems are in Table 19 ; the other systems are defined in the text. “Basins” is the number of designed attractors and of trajectory-fate labels ( Section B.4 ). On Transfer-gated local-linear, about

40\%

of uniform starts instead stall at region boundaries, and the labeling groups these stall points with the designed attractors.

Figure 6: Ground-truth vector fields for the retained

15

-system multibasin benchmark. The models are trained only on stored state trajectories; these continuous-time vector fields and design-center annotations are shown here only to document the benchmark geometry and are not provided to the learner.

System	$\{c_{i}\}$	$(a_{i},\sigma_{i})$	$(\omega,\gamma)$
High-cross three-well	$\{p_{i}^{(3)}(1.8,\pi/2)\}$	$(3.0,0.5)$	$(2.0,0.03)$
Hexagonal six-well	$\{p_{i}^{(6)}(1.7,0)\}$	$(2.0,0.5)$	$(1.0,0.03)$
Octagonal eight-well	$\{p_{i}^{(8)}(2.2,0)\}$	$(3.0,0.5)$	$(0.90,0.02)$
Pentagonal five-well	$\{p_{i}^{(5)}(1.8,\pi/2)\}$	$(2.0,0.5)$	$(1.1,0.03)$
Square four-well	$\{p_{i}^{(4)}(1.8,\pi/4)\}$	$(3.0,0.5)$	$(1.0,0.03)$
Diamond four-well	$\{(0,2.2),(2.2,0),(0,-2.2),(-2.2,0)\}$	$(2.5,0.5)$	$(1.0,0.02)$

Appendix

Table 19: Parameters for retained Gaussian-well systems. Repeated pairs in the

(a_{i},\sigma_{i})

column apply to every listed center.

Display name	System key	Dimension	Native $dt$	Stored step $5dt$
Chua	Chua	3	$2.8474745791{\times}10^{-4}$	$1.4237372895{\times}10^{-3}$
Dadras	Dadras	3	$6.5782963827{\times}10^{-4}$	$3.2891481914{\times}10^{-3}$
Dequan Li	DequanLi	3	$1.6763993999{\times}10^{-5}$	$8.3819969995{\times}10^{-5}$
Hadley	Hadley	3	$2.9086847808{\times}10^{-4}$	$1.4543423904{\times}10^{-3}$
Lu–Chen–Cheng	LuChenCheng	3	$1.8469678280{\times}10^{-4}$	$9.2348391399{\times}10^{-4}$
Qi–Chen	QiChen	3	$7.8371061844{\times}10^{-5}$	$3.9185530922{\times}10^{-4}$

Appendix

Table 20: The ten Dysts systems used in the forecasting benchmark. The “stored step” column is

5

times the native Dysts integration interval, which gives

999.8

stored steps per dominant period for every system.

Koopman Operator Learning

Koopman operator learning builds finite-dimensional linear models of nonlinear dynamics by finding useful observable representations from data.

See all 79 papers

Learning the Koopman Operator using Attention Free Transformers
Jun 22, 2026Mohammed Nagdi, Evangelos-Marios Nikolados, Alexey Yermakov +3Long-Term Time Series Forecasting Koopman Operator Learning
Institute for Bioengineering, School of Engineering and Centre for Engineering Biology, University of Edinburgh, Edinburgh, United Kingdom · Electrical and Computer Engineering and Applied Mathematics, University of Washington, Seattle, WA, USA · Electrical and Computer Engineering and Computer Science & Engineering, University of Washington, Seattle, WA, USA
Deep Embedded Multiplicative DMD for Algebra-Preserving Koopman Learning
Jun 3, 2026Kelan Gray, Finlay Brown, Nicolas Boullé +1Koopman Operator Learning Latent Dynamics Modeling
Department of Mathematics, Imperial College London, London, SW7 2AZ, UK. · Department of Applied Mathematics and Theoretical Physics, University of Cambridge, Cambridge, CB3 0WA, UK.
Variational Augmented Invertible Koopman Autoencoder for probabilistic time series forecasting
Sep 29, 2026Anthony Frion, Lucas Drumetz, Guillaume Tochon +3Probabilistic Time Series Forecasting Long-Term Time Series Forecasting
Institute of coastal modelling Helmholtz-Zentrum Hereon · IMT Atlantique Lab-STICC, UMR CNRS 6285, Brest, France · LRE EPITA, Le Kremlin-Bicêtre, France +1

Jun 22, 2026J/K move · Enter open · S save

Learning the Koopman Operator using Attention Free Transformers

Mohammed Nagdi, Evangelos-Marios Nikolados, Alexey Yermakov, Mars Gao+2

Institute for Bioengineering, School of Engineering and Centre for Engineering Biology, University of Edinburgh, Edinburgh, United Kingdom · Electrical and Computer Engineering and Applied Mathematics, University of Washington, Seattle, WA, USA · Electrical and Computer Engineering and Computer Science & Engineering, University of Washington, Seattle, WA, USA

Long-Term Time Series Forecasting Koopman Operator Learning Latent Dynamics Modeling Transformer

Open paper AskPDF 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

\approx

30k parameters (

3d^2 + T^2

, 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.

Abstract

Koopman Operator Learning

Explore similar work

Learning the Koopman Operator using Attention Free Transformers

Deep Embedded Multiplicative DMD for Algebra-Preserving Koopman Learning

Variational Augmented Invertible Koopman Autoencoder for probabilistic time series forecasting