Sparse Koopman Autoencoders Identify Local Dynamical Regimes in Multibasin Systems

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

Source-linked AI summary

Sparse Koopman Autoencoders Identify Local Dynamical Regimes in Multibasin Systems

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

arXiv:2608.29057v1Aug 29, 2026cs.LGmath.DS

TL;DR

Multibasin systems generally cannot be represented by one continuous finite-dimensional global Koopman embedding, motivating local regime-aware representations. The paper trains sparse Koopman autoencoders without basin labels, using latent supports to encode regimes, and reports improved forecasting and interpretable basin-related structure. These supports are most reliable as basin-interior regime variables rather than perfect global classifiers.

Problem
Multibasin systems generally cannot admit a single continuous finite-dimensional global Koopman embedding, while different basins may require incompatible local linearizations.
Method
The paper combines Koopman autoencoders with LISTA-style sparse coding so active latent supports represent local regimes and nonzero coefficients represent within-regime state.
Results
SKAEs improve long-horizon forecasting across multibasin and chaotic benchmarks, while their supports are necessary for accurate rollouts and interpretable as emergent local dynamical regimes.
Takeaways & Limitations
Sparse supports provide a label-free interface between continuous Koopman coordinates and discrete regime structure for multibasin dynamics.
Takeaways & Limitations
Interpretability is evaluated on known benchmark geometries and basin-interior states; deterministic benchmarks leave noise and partial-observability settings outside the demonstrated scope.

Abstract

from arXiv · show

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 held-out basin interior states, whereas dense-latent KAEs collapse to an uninformative single 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.

1 Introduction

Koopman autoencoders learn linear latent dynamics, but multibasin systems generally resist a single continuous finite-dimensional global embedding. The paper proposes sparse latent supports as label-free, inspectable regime variables for learning compatible local representations.

Multibasin systems generally cannot admit one continuous finite-dimensional global Koopman linearization, although separate basins may support distinct local linearizations.
SKAEs induce sparse latent codes whose active supports can distinguish local dynamical regimes, while nonzero coefficients retain continuous predictive coordinates.
Sparse coding and LISTA provide a piecewise-linear, union-of-subspaces inductive bias suited to basin-specific local linearizations.
The paper introduces SKAEs and reports improved long-horizon forecasting, functionally necessary supports for accurate rollouts, and interpretable emergent local regimes across multibasin and chaotic benchmarks.

2 Problem formulation

The paper formulates multibasin Koopman learning as state reconstruction with linear latent prediction when the appropriate law depends on basin location. It uses stored-step dynamics, trajectory windows, repeated latent transitions, and no basin labels during training.

Forecasting advances the discrete observation map by repeatedly applying F∆t, with horizon h denoting h map applications.
Equal forecast horizons can represent different physical times because the system-specific observation interval ∆t may vary across systems.
A basin of attraction contains initial conditions whose trajectories converge to an attractor, and coexisting basins partition the state space up to boundaries.
The task is to learn a representation supporting linear prediction whose appropriate law depends on a state’s location in multibasin geometry.
Koopman autoencoders learn an encoder, decoder, and latent linear transition matrix, applying the same matrix at every rollout step.
Training uses adjacent observation windows, encoding each state and decoding reconstructed observations alongside latent rollouts and forecasts.
Basin labels are withheld from model selection, encoder and transition training, and support selection, serving only for post-hoc evaluation when available.

3 Method

SKAEs combine sparse coding with Koopman autoencoders so active latent supports identify local dynamical regimes while nonzero coefficients represent within-regime coordinates. The method jointly trains encoders, decoders, and latent transitions, then evaluates support structure and model-only periodic re-encoding.

Sparse supports as local regime variables: Sparse latent supports act as regime variables, while nonzero coefficients provide continuous coordinates for the corresponding local linear representation.This couples basin identification and within-basin prediction in one latent code.
Sparse coding objective: The Lasso objective reconstructs the input while penalizing the latent code's l1 norm, encouraging few nonzero coefficients.The sparse code solves an under-specified reconstruction problem with selective activation.
Sparse encoder: LISTA unrolls ISTA iterations into an encoder whose support is the set of active latent coordinates, inducing data-adaptive regions with shared decoder participation.When dictionary and code are learned jointly, states sharing a support occupy the same reconstruction subspace.
Joint training objective: SKAE training jointly optimizes an encoder, linear decoder, and latent transition using reconstructed states, decoded rollouts, latent-linearity terms, and sparsity.Decoder columns are constrained to unit norm; transition variants are dense, block diagonal, or softly block-regularized.
Support objects: Support evaluation distinguishes exact active-coordinate masks from support families, which group overlapping masks to reduce basin fragmentation.Absolute-threshold supports use |z_i| > 10^-3, and families assign masks by Jaccard overlap with threshold τ = 0.5.
Rollout evaluation: Periodic re-encoding advances predicted states for a selected number of Koopman steps, decodes them, and re-encodes them throughout the rollout.The procedure uses only model predictions, without true future states, basin labels, or transition oracles; the period is tuned on training data.

4 Experiments

The experiments compare sparse and dense Koopman autoencoders across multibasin and chaotic systems, evaluating forecasting, support interventions, and basin-support alignment. Sparse models forecast better, require correct active coordinates for accurate rollouts, and produce support families that better identify basin membership than dense models.

Experimental design: Forecasting uses held-out trajectories at the stored observation cadence, with mean squared error summarized using per-system seed IQMs and arithmetic means across systems.Horizon trends use seed-bootstrap bands describing uncertainty over seeds for an average system.
Forecasting performance: At long horizons, the dense MLP has 17-27x larger MSE at H1000 on multibasin systems and about 1.5x larger MSE at H4000 on Dysts systems than most sparse models.Sparse models significantly reduce forecasting error at all horizons on both benchmark sets, while differences among sparse encoders are marginal.
Support interventions: At H = 21, standard rollout MSE is 0.0158 versus 0.508, 1.37, 2.08, 3.26, and 8.51 after dropping the top 1, 2, 3, 5, and 10 active coordinates.Randomly reassigning active coefficient values to inactive coordinates is also highly destructive, isolating the role of support selection from coefficient values.
Support-basin alignment: Support-based basin identification is evaluated on held-out states deep inside basins, because supports may change under small perturbations near separatrices.The alignment diagnostic therefore does not establish reliable identification at basin boundaries.
Support-basin alignment: LISTA support families achieve the lowest conditional entropy, 0.130, for identifying ground-truth basins, whereas the dense MLP collapses to one support family.LISTA models expose the most support families, closer to the benchmark basin-count mean/median of 4.20/4.

5 Discussion

Sparse supports provide a practical, label-free way to represent multiple local Koopman regimes in one model. Their interpretation is strongest in basin interiors and remains bounded by the deterministic benchmark setting.

Sparse Koopman autoencoders use different latent supports to learn multiple local Koopman regimes without basin labels.The supports provide a post-training diagnostic for inspecting basin structure.
Support-based interpretability is evaluated on systems with known basin geometry and states deep inside basins.Near separatrices, support assignments can become fragmented or unstable.
The benchmarks are deterministic, while real systems may include noise, partial observability, nonstationarity, control inputs, or non-basin-like regime changes.
Sparse supports connect continuous Koopman coordinates with discrete regime structure for multi-regime forecasting, analysis, and potentially control.The paper positions this structure as a practical step toward Koopman models for multibasin dynamics.

A Support definitions and sensitivity

The appendix distinguishes exact active supports from grouped support families and specifies deterministic rules for constructing and interpreting them. It also fixes a Jaccard threshold whose value trades off basin merging against support fragmentation.

Support definitions: Exact supports are Boolean active-coordinate masks derived from latent coefficients exceeding a strict magnitude threshold of 10^-3.The measured support size is state-dependent rather than a fixed sparsity budget.
Support families: Support families greedily group exact masks by Jaccard similarity, processing masks by decreasing frequency and using serialized masks as deterministic tie-breakers.Each family retains an actual Boolean representative.
Sensitivity: For top-eight masks, Jaccard similarity of at least 0.5 means at least six shared active coordinates.Lower thresholds can merge distinct basins, whereas higher thresholds can over-fragment them.
Diagnostic roles: Sabs supports exact-support diagnostics and wrong-support interventions, while Fabs groups exact masks for basin-scale alignment analysis.
Diagnostic interpretation: Low H(B | Sabs) or H(B | Fabs) indicates basin prediction on evaluated states but does not establish one exact support per basin.Additional fragmentation diagnostics include H(Sabs | B), exact-support uniqueness, and support counts.

B Additional experimental details

The experiments use fixed, label-free protocols across controlled multibasin systems and deterministic chaotic flows. Forecasting, support alignment, and interventions are evaluated on held-out data with predefined aggregation and checkpoint procedures.

Benchmarks: The controlled benchmark contains 15 two-dimensional multibasin systems, while the external Dysts stress test contains 10 three-dimensional chaotic systems.Learners observe stored states rather than vector fields, integration substeps, basin labels, or assignments.
Models: Six model rows use latent dimension dz = 256, including sparse LISTA, sparse MLP, block-diagonal variants, a soft-block ablation, and Dense MLP.Dense MLP removes the sparsity mechanism by setting its sparsity coefficient to zero.
Training: All controlled models use 15 seeds, rollout length L = 8, and 200,000 optimization steps; Dysts models use L = 10 and 100,000 steps.
Training: A label-free reset procedure retains boundary-adjacent and transient states using perturbation sensitivity and late-rollout motion.Half of controlled training windows start from the retained pool with jitter.
Reporting: Checkpoints are selected by validation rollout error, while reported tables aggregate completed seeds rather than selecting the best seed.
Evaluation: Forecasting uses held-out trajectories, with no-reencoding and periodic rollouts; periodic re-encoding uses predicted states rather than true future states.
Evaluation: Support alignment uses held-out states from the deepest quartile within each represented basin, then computes exact supports, Fabs families, and H(B | Fabs).This slice is intended to provide a clean basin-interior evaluation population.

C Statistical testing protocol

The statistical protocol separates descriptive summaries from inferential tests and matches each analysis to its experimental independent unit. Effects are paired, directionally tested, and corrected for multiple comparisons.

Aggregation: Point estimates use seed-level interquartile means within systems followed by averages across systems, while family counts use descriptive per-system seed means.
Effect measures: Forecasting MSE comparisons use paired log effects against Dense MLP, with negative effects favoring the candidate.Log ratios address multiplicative comparisons and reduce heavy-tailed MSE influence.
Effect measures: Support diagnostics use raw paired differences, and family counts receive no one-sided test because closeness to basin count is non-monotone.
Multibasin inference: Controlled-system analyses use paired Wilcoxon tests over training seeds within systems, requiring at least four finite paired seeds, then Holm-correct across systems.
Multibasin inference: Every shown non-Dense multibasin forecasting cell passes the prespecified corrected test on all 15/15 systems.
Dysts inference: Dysts inference treats systems as independent units after within-system seed summarization and uses Holm-corrected one-sided sign tests.A displayed superscript corresponds to improvement on all 10/10 Dysts systems.
Interventions: Figure 3 is a mechanism diagnostic, reporting bootstrap intervals for coordinate dropping and distribution summaries for random support shuffling.

D Multibasin benchmark inventory and system definitions

The multibasin benchmark comprises 15 procedurally generated two-dimensional systems with documented attractor geometry and varied local vector-field constructions. Models learn from stored trajectories, while basin information and continuous-time vector fields remain evaluation or documentation resources.

Inventory: The benchmark contains 15 procedurally generated two-dimensional multibasin systems, with every system contributing to aggregate results.Four systems are additionally shown as main benchmark visual panels.
Data and annotations: Basin labels are restricted to benchmark annotation, controlled-transfer construction, and evaluation rather than support extraction, routing, or rollout evaluation.For structured-transition ablations, basin counts additionally set a fixed architecture group count.
Data and annotations: All retained systems are continuous-time vector fields integrated with fourth-order Runge–Kutta to produce stored observations, while training windows contain only stored states.The learner does not receive vector fields, integration substeps, or basin assignments.
System families: Most retained systems use two-dimensional Gaussian potentials with independent rotational drift, instantiated with system-specific parameters.The Gaussian-well family includes repeated center, scale, and potential parameters listed in the benchmark tables.
System families: Additional constructions include transition-route systems, native gated local-linear systems, arrested spirals, and nonlinear potential-based systems.These systems use corridor-dependent rotations, region-specific gates, traps, or confinement terms to create specialized local dynamics.

E Dysts benchmark inventory and system definitions

The chaotic-flow benchmark uses ten three-dimensional autonomous systems from the Dysts repository under a pinned package version. Forecasting is performed on stored observations sampled at intervals 30 times each system’s native integration step, without exposing vector fields or solver substeps to the learner.

Inventory: The benchmark includes ten three-dimensional autonomous flows drawn from the Dysts repository.The systems use continuous-time right-hand sides and metadata from pinned package version dysts 0.96.
Stored observations: The stored observation interval is Δt = 30 dtnative, where dtnative is system-specific.Thus, the benchmark’s stored step is 30 times the native Dysts integration interval.
Stored observations: The learner observes only resulting stored states, not the continuous-time vector field, solver, or integration substeps.The systems nevertheless have explicit closed-form right-hand sides, mostly polynomial, with some absolute-value nonlinearities.
System definitions: The inventory spans named systems including Chua, Dequan Li, Lu–Chen–Cheng, San–Um–Srisuchinwong, Shimizu–Morioka, and Wang–Sun.The supplied definitions specify system-specific parameters and equations for these entries.