PaperlayerView exact revision on arXiv

Source-linked AI summary

DisKO: Deep Koopman Learning in Distribution Space from Unpaired Snapshots

He Ma, Xiaochen Liu, Wanfeng Lu, Ying Wang, Wei Lin, Qunxi Zhu

arXiv:2609.34629v1cs.LGmath.DS

TL;DR

Unpaired distribution snapshots make trajectory-based learning difficult because individual observations are not linked across time, while distribution space is infinite-dimensional. DisKO learns predictive observables, finite-dimensional Koopman dynamics, and a decoder from snapshots; across seven benchmarks it achieves state-of-the-art extrapolation and recovers interpretable spectral structure. The paper also identifies open questions about identifiability and learning guarantees.

  • Problem

    Unpaired snapshots provide no individual correspondences across time, and infinite-dimensional distribution space complicates learning compact, approximately closed representations for long-horizon prediction.

    Source: §1 · p. 1, §3.3 · p. 4

  • Method

    DisKO jointly learns distributional observables, finite-dimensional Koopman dynamics, and a decoder that maps propagated coordinates back to distributions.

    Source: §1 · p. 2

  • Results

    Across seven benchmarks, DisKO achieves state-of-the-art extrapolation with substantially slower long-horizon error accumulation, and recovers leading eigenvalues and eigenfunctions on systems with analytic spectra.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The learned representation supports both accurate distribution extrapolation and interpretable decay and oscillation recovery from unpaired snapshots.

    Source: §6 · p. 10

    Identifiability of Koopman dynamics from unpaired snapshots and convergence and stability guarantees for the learning procedure remain open theoretical questions.

    Source: §6 · p. 10

Abstract

from arXiv · show

Many complex systems are observed only through temporally unpaired distribution snapshots, making trajectory-based dynamical learning difficult without additional assumptions. We therefore formulate the problem directly in distribution space, treating the distribution itself as the dynamical state. The challenge is that distribution space is infinite-dimensional, making compact and approximately closed representations difficult to learn from finite snapshots. We introduce DisKO, which extends deep Koopman learning to distribution dynamics by jointly learning predictive distributional observables, a finite-dimensional Koopman representation, and a generative map back to the full distribution. Across seven diverse benchmarks, DisKO achieves state-of-the-art extrapolation performance, with substantially slower error accumulation on long-horizon prediction tasks. DisKO further recovers leading Koopman eigenvalues and eigenfunctions on systems with analytic spectra, revealing meaningful dynamical structure in the learned representation.

1 INTRODUCTION

DisKO learns distribution dynamics directly from unpaired snapshots, jointly fitting predictive observables, finite-dimensional Koopman dynamics, and a decoder. Across seven benchmarks, it achieves state-of-the-art extrapolation and recovers interpretable spectral structure on systems with analytic spectra.

1 INTRODUCTION

  • Unpaired distribution snapshots contain samples at limited time points without links between individuals, so the learning target is the temporal evolution of distributions.Such data arise in single-cell sequencing, collective motion and transport, and cosmological simulations.Source: §1 · p. 1
  • Trajectory reconstruction from marginal distributions requires additional assumptions because different couplings or stochastic processes can produce the same snapshots.Modeling dynamics directly in distribution space instead targets population-distribution prediction.Source: §1 · p. 1
  • Distribution space is infinite-dimensional, complicating the learning of compact, approximately closed representations for long-horizon prediction.Koopman methods lift nonlinear dynamics to linear observable evolution, while deep Koopman learning learns observables from data.Source: §1 · p. 2
  • Existing data-driven distributional Koopman methods use prescribed observable dictionaries and provide neither learned predictive observables nor a generative map to full distributions.Table 1 compares methods for learning distribution dynamics.Source: §1 · p. 2, Table 1 · p. 2
  • DisKO jointly learns predictive distributional observables, a finite-dimensional Koopman representation, and a generative map back to full distributions.Its pipeline encodes snapshots, evolves observables, then decodes distributions for reconstruction and extrapolation.Source: §1 · p. 2, Figure 1 · p. 2
  • Across seven benchmarks, DisKO achieves state-of-the-art extrapolation with substantially slower long-horizon error accumulation, and recovers leading eigenvalues and eigenfunctions on systems with analytic spectra.The recovered modes capture interpretable decay and oscillation.Source: §1 · p. 2

2 RELATED WORK

Related work addresses population prediction from unpaired snapshots through assumptions on between-snapshot evolution, while Koopman methods provide linear observable dynamics and set-based generative models support distribution representation. DisKO brings these strands together for prediction and spectral analysis.

2 RELATED WORK

  • Prior snapshot methods constrain between-snapshot evolution through transport regularization, population growth, learned latent coordinates, or energy and mechanical principles; DisKO instead learns distribution observables for prediction and spectral analysis.Examples include TrajectoryNet, DeepRUOT, VGFM, MIOFlow, scNODE, PRESCIENT, JKOnet*, and WLM.Source: §2 · p. 3
  • Koopman methods approximate linear observable evolution, with representations learned through dictionaries or neural embeddings; distributional Koopman theory makes the observables functionals of distributions.DisKO learns those functionals jointly with continuous dynamics and a generative decoder.Source: §2 · p. 3
  • Permutation-invariant set encoders represent finite snapshots, and conditional generative flows trained by flow matching decode evolved representations into distributions.Source: §2 · p. 3

3 THEORETICAL FRAMEWORK

The framework treats distributions as states evolving under deterministic distribution-space dynamics, observed through unpaired samples at known times. It represents distribution observables in finite dimensions, propagates them with Koopman-inspired dynamics, and decodes them back to distributions.

3.1 PROBLEM FORMULATION

  • Distribution snapshots are unordered state samples from deterministic or stochastic dynamics, with observation times and sequence membership known but individual samples unpaired across time.The model assumes a unique nonexplosive solution for the stated stochastic differential equation; setting σ = 0 recovers an ODE.Source: §3.1 · p. 3
  • The learning task is to extrapolate distribution evolution beyond the training interval from a source snapshot, using future snapshots as evaluation targets.Sequences share dynamics but may begin from different initial distributions; samples are unpaired across time.Source: §3.1 · p. 3

3.2 THE DISTRIBUTIONAL KOOPMAN OPERATOR

  • The distributional Koopman operator composes an observable with distribution evolution, mapping a distribution-valued input to a scalar statistic and evolving linearly on observables.Its generator describes instantaneous evolution of distribution observables.Source: §3.2 · p. 4

3.3 LEARNING FINITE-DIMENSIONAL DISTRIBUTIONAL REPRESENTATIONS

  • A finite-dimensional encoder supplies distribution observables, an affine latent flow propagates their coordinates, and a decoder maps those coordinates back to distributions.Exact propagation requires an invariant observable span; the learned representation instead seeks approximate closure, assessed through prediction and spectral diagnostics.Source: §3.3 · p. 4
  • For an OU system, mean and variance observables form an invariant span with exact affine dynamics, and a Gaussian decoder exactly reconstructs the evolving distribution.The two coordinates evolve as ˙z_t = (−am_t, −2av_t + σ²).Source: §3.3 · p. 4

4 DISKO MODEL AND LEARNING ALGORITHM

DisKO combines permutation-invariant snapshot encoding, affine latent dynamics, and conditional generative decoding. Its training couples reconstruction and prediction with latent consistency and distribution-level losses, using staged fitting or a discrete closed-form alternative.

4 DISKO MODEL AND LEARNING ALGORITHM

  • DisKO instantiates distribution encoding, latent propagation, and decoding with a set network, affine dynamics, and conditional flow matching, with a discrete closed-form alternative to continuous generator learning.Source: §4 · p. 5

4.1 DISTRIBUTION ENCODING WITH SET NETWORKS

  • The permutation-invariant encoder combines learned set aggregation with snapshot mean and standard deviation to produce a fixed-dimensional representation independent of sample order.Deep Sets pooling and Set Transformer attention are the described variants.Source: §4.1 · p. 5

4.2 CONTINUOUS-TIME LATENT DYNAMICS

  • A shared affine generator governs all time intervals, supports irregular observation times and arbitrary future-time prediction, and has eigenvalues interpreted as latent decay, growth, and oscillation rates.Source: §4.2 · p. 5

4.3 DISTRIBUTION DECODING VIA CONDITIONAL FLOW MATCHING

  • Conditional flow matching decodes latent conditions into distributions by learning a velocity field along a noise-to-sample path; generation integrates that field at fixed physical time.The propagated latent condition carries physical evolution, and the decoder accommodates the 30-dimensional Pancreas and 50-dimensional WOT iPSC representations.Source: §4.3 · p. 5

4.4 TRAINING OBJECTIVES AND ALGORITHM

  • Training pairs unpaired snapshots from the same sequence and combines prediction, reconstruction, latent consistency, and sample-distribution losses; staged pretraining precedes joint optimization.For uniformly spaced snapshots, an alternative fits discrete affine updates by ridge regression and alternates refitting with encoder-decoder updates.Source: §4.4 · pp. 5–6

5 RESULTS

Across seven distribution-dynamics benchmarks, DisKO leads the W1 comparisons, while analytic-spectrum tests show accurate eigenvalue and eigenfunction recovery. OU analyses also support approximate closure along one test trajectory and show lower future errors for Koopman dynamics than for a Neural ODE ablation.

5.1 EXPERIMENTAL SETUP

  • The evaluation covers seven datasets, including long-horizon tests on OU and Duffing, single-cell prediction on Pancreas and WOT iPSC, and analytic-spectrum tests on Circle and Torus.The protocol separates observation times into training and extrapolation windows and also splits initial distributions for OU, Duffing, and CAMELS.Source: §5.1 · p. 6

5.2 DISTRIBUTION EXTRAPOLATION

  • DisKO achieves the lowest W1 on all seven datasets and lowest SW1 on six; W1 reductions versus the strongest baseline are 58.6% on OU, 53.9% on Duffing, and 36.2% on Ocean.The corresponding W1 reductions on Pancreas and WOT iPSC are 18.0% and 18.7%.Source: Table 2 · p. 7, §5.2 · p. 7
  • On OU, prediction error remains approximately constant over 50 extrapolation steps; on Duffing, it grows slowly over 80 steps and remains below the baselines shown.Figure 2 shows W1 errors on four benchmarks, with the extrapolation window shaded and error axes on a logarithmic scale.Source: §5.2 · p. 7, Figure 2 · p. 8

5.3 KOOPMAN SPECTRAL APPROXIMATION

  • At Δt = 0.1, eigenvalue errors are below 0.006 for Circle’s first three harmonics and 0.003 for Torus fundamental modes.Discrete and continuous parameterizations both show close agreement between analytical and recovered finite-time eigenvalues.Source: Table 3 · p. 8, §5.3 · p. 8
  • On future test snapshots, complex correlations reach 0.999, 0.994, and 0.976 for the first three Circle harmonics, and exceed 0.995 for both Torus fundamental modes.Discrete DisKO better resolves higher Circle harmonics, while continuous DisKO gives more accurate Torus eigenvalues; these are distinct tests.Source: §5.3 · p. 9

5.4 DYNAMICAL CLOSURE AND EXTRAPOLATION STABILITY

  • Beyond the training interval, propagated and directly encoded OU trajectories agree in rotation and contraction in the PCA projection, supporting approximate closure along this trajectory.Figure 4 compares these encodings and forecasts for one OU test sequence.Source: Figure 4 · p. 9, §5.4 · p. 9
  • Against the Neural ODE ablation, Koopman dynamics reduce future W1 from 0.201 ± 0.026 to 0.149 ± 0.007 and SW1 from 0.107 ± 0.018 to 0.065 ± 0.005.The comparison keeps encoder and decoder architectures, losses, training budgets, and pretrained weights matched; Koopman dynamics also show smaller seed variation in all four metrics.Source: Figure 5 · p. 9, §5.4 · p. 9

6 CONCLUSION AND LIMITATIONS

DisKO learns a shared representation for distribution extrapolation and spectral recovery. The paper leaves identifiability and learning guarantees as open theoretical questions.

6 CONCLUSION AND LIMITATIONS

  • DisKO’s shared representation supports accurate distribution extrapolation and recovery of interpretable decay and oscillation.The reported spectral recovery concerns systems with analytic Koopman spectra.Source: §6 · p. 10
  • Identifiability conditions for Koopman dynamics from unpaired distribution snapshots remain an open question.Source: §6 · p. 10
  • Convergence and stability guarantees for learning the operator when it is identifiable also remain open.Source: §6 · p. 10

AI USE STATEMENT

The authors used generative AI tools for implementation support, visualization, and language polishing, while retaining responsibility for reviewing the assisted content.

AI USE STATEMENT

  • Generative AI assisted code implementation and debugging, including baseline implementations.Source: AI USE STATEMENT · p. 10
  • The tools also supported experimental visualization and language polishing.Source: AI USE STATEMENT · p. 10
  • The authors state that AI was not used for research ideas, core methodology, theoretical framing, or experimental results, and that they reviewed all assisted content.The authors take full responsibility for the final content.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

The experiments use synthetic simulations and previously released datasets, without collecting new human-subject data or using personally identifiable information. The authors call for domain-specific validation and safeguards in sensitive or high-stakes applications.

ETHICS STATEMENT

  • The work uses synthetic simulations and previously released benchmarks, including publicly available single-cell datasets, rather than new human-subject data collection.Source: ETHICS STATEMENT · p. 10
  • The experiments do not use participant identifiers or personally identifiable information.Source: ETHICS STATEMENT · p. 10
  • Sensitive or high-stakes applications of the method should require domain-specific validation and safeguards.Source: ETHICS STATEMENT · p. 10

A.1 DATASETS AND PREPROCESSING

The appendix describes seven distribution-prediction benchmarks, their temporal partitions and prediction protocol, and the evaluation metrics and sampling conventions.

A.1 DATASETS AND PREPROCESSING

  • The seven benchmarks include synthetic Circle and Torus systems, Ornstein–Uhlenbeck and Duffing dynamics, and Boids and Ocean datasets.The Circle and Torus systems provide analytical reference spectra; snapshots are sample sets without individual correspondences across times.Source: §A.1 · pp. 16–17
  • The remaining benchmarks include CAMELS cosmology and the Pancreas and WOT iPSC single-cell datasets, each with specified training and extrapolation periods.Source: §A.1 · p. 17

A.2 PREDICTION PROTOCOL

  • The prediction protocol evaluates distinct sequence and time-window splits where available, while single-sequence datasets use temporal partitions.Predictions are conditioned on the earliest snapshot; later test snapshots are used only for evaluation, without updates from intervening targets.Source: §A.2 · pp. 17–18
  • DisKO’s data standardization uses training sequences within the training time window and is fixed during prediction.Evaluation transforms predictions back to input coordinates and averages errors over relevant sequences and target times.Source: §A.2 · p. 18

A.3 EVALUATION METRICS

  • Evaluation uses W1, SW1, W2, and squared MMD, with smaller values indicating closer agreement; SW1 uses 128 directions and MMD uses at most 512 samples.Source: §A.3 · pp. 18–19, Table A.2 · p. 18
  • DeepRUOT’s Pancreas and CAMELS results use weighted evaluation with 1,024 samples and a different MMD bandwidth convention from the standard configuration.Source: §A.3 · p. 19

B DISTRIBUTION EXTRAPOLATION EXPERIMENTS

Across the seven benchmarks, DisKO has the lowest displayed extrapolation errors on OU, Duffing, Ocean, Pancreas, and WOT iPSC, while remaining close to WLM on CAMELS. The spectral experiments compare discrete and continuous DisKO against analytic eigenpairs on Circle and Torus.

B DISTRIBUTION EXTRAPOLATION EXPERIMENTS

  • DisKO has the lowest displayed extrapolation errors on OU, Duffing, Ocean, Pancreas, and WOT iPSC; on CAMELS, it finishes close to WLM and below DeepRUOT and MIOFlow.The errors increase with horizon for all four methods on CAMELS.Source: B DISTRIBUTION EXTRAPOLATION EXPERIMENTS · p. 20, Figure B.1 · p. 21
  • Visual comparisons show DisKO retaining Duffing’s two future lobes and tracking the changing Boids, Ocean, CAMELS, Pancreas, and WOT iPSC distributions, with residual shape or location errors.Boids predictions blur and shift later; WOT iPSC predictions are displaced toward lower values of the second principal component at the latest times.Source: B DISTRIBUTION EXTRAPOLATION EXPERIMENTS · p. 26

C.1 ANALYTICAL EIGENVALUES AND EIGENFUNCTIONS

  • For Circle and Torus, the analytic distributional eigenfunctions are expectations of Fourier observables, with generator eigenvalues λk and discrete-time eigenvalues ρk = e^(∆tλk).The selected Fourier indices define the reference eigenpairs used for comparison.Source: §C.1 · p. 34

C.2 DISCRETE AND CONTINUOUS DISKO

  • The discrete model uses affine one-step latent updates and ridge-refits its operator, whereas the continuous model learns an affine generator and propagates it with a matrix exponential.The spectral comparison uses different initialization procedures and dynamics-update budgets, so it compares complete training configurations.Source: §C.2 · p. 35

C.3 EIGENVALUE APPROXIMATION

  • On Circle, mean eigenvalue errors are 0.050311 for the discrete model and 0.175543 for the continuous model; on Torus, they are 0.008816 and 0.002256, respectively.At Circle k = 4, errors are 0.0179 for the discrete model and 0.2145 for the continuous model.Source: §C.3 · p. 36, Figure C.1 · p. 37, Figure C.2 · p. 38

C.4 EIGENFUNCTION EVALUATION ON DISTRIBUTIONS

DisKO’s learned eigenfunctions are evaluated on future observed distributions and on controlled wrapped-Gaussian probes. A separate OU visualization finds that latent forecasts follow the main projected trajectory beyond the training interval, with residual amplitude and timing errors.

C.4 EIGENFUNCTION EVALUATION ON DISTRIBUTIONS

  • On future test distributions, Circle eigenfunction correlations favor the discrete model at k = 2 and 3, while both Torus models exceed 0.995 for the two fundamental indices.Across the six Torus indices, mean future correlations are 0.881 for discrete and 0.849 for continuous DisKO.Source: §C.4 · p. 39, Figure C.3 · p. 40
  • On fixed-width wrapped-Gaussian probes, Circle k = 3 NRMSE is 0.123 for discrete and 0.697 for continuous DisKO, despite similar eigenvalue errors.These probes characterize the functions only on the selected Gaussian family, not across all probability measures.Source: §C.4 · p. 39, Figure C.4 · p. 41, Figure C.5 · p. 42

D LATENT DYNAMICAL CONSISTENCY

  • In the OU test sequence, the projected latent forecast continues the reference’s contracting rotation beyond t = 2.5, with its clearest discrepancy near t = 4.The predicted minimum in the second coordinate is shallower and slightly later; projection and averaging limit the evidence to the displayed dynamics.Source: D LATENT DYNAMICAL CONSISTENCY · p. 43, Figure D.1 · p. 43

E ABLATION STUDY

On OU, the ablation compares continuous affine Koopman dynamics with a similarly sized autonomous Neural ODE under a shared encoder-decoder and training protocol. The Koopman model has lower future errors and less variation across seeds in all four reported metrics.

E ABLATION STUDY

  • The continuous affine generator and Neural ODE have similar parameter counts, 4,160 and 4,192, and share the encoder-decoder architecture and training protocol.Both use 128 OU training initial distributions, 64 test initial distributions, and five training seeds.Source: E ABLATION STUDY · p. 44
  • The Koopman model has lower future prediction errors and smaller variation across training seeds in all four metrics than the Neural ODE.For 2.5 < t ≤ 5, W1 is 0.149 ± 0.007 for Koopman and 0.201 ± 0.026 for Neural ODE; SW1 is 0.065 ± 0.005 and 0.107 ± 0.018, respectively.Source: E ABLATION STUDY · p. 44, Figure E.1 · p. 44
  • Under this OU training protocol, the affine Koopman parameterization is supported as an effective inductive bias for accurate and reproducible extrapolation.The Neural ODE shows a more pronounced increase in SW1 and MMD2 near the training boundary and into the future interval.Source: E ABLATION STUDY · p. 44

F DATA EFFICIENCY

On the OU benchmark, more training initial distributions reduced prediction errors, while a larger available sample pool brought limited additional accuracy under a fixed input budget. The accompanying theory characterizes distributional observables, approximate Koopman closure, and conditions for identifying evolution from observations.

F DATA EFFICIENCY

  • From R = 16 to R = 128, W1 decreases from 0.169±0.010 to 0.150±0.007, and SW1 decreases from 0.086±0.006 to 0.073 ± 0.003.The reported mean errors decrease in all four metrics as R increases.Source: Figure F.1 · p. 45, F DATA EFFICIENCY · p. 45
  • At N = 512, mean errors are lowest in all four metrics; across the tested sample pools, added samples yield limited accuracy gains with 128 particles per input.Mean W1 ranges from 0.149 to 0.161, with nonmonotonic changes; seed variation is largest at N = 256.Source: Figure F.2 · p. 45, F DATA EFFICIENCY · p. 45

G.1 DISTRIBUTION OBSERVABLES AND KOOPMAN OPERATORS

  • Distribution observables are scalar functionals that evolve linearly under the Koopman operator; reconstruction uses a decoder, whose exact recovery requires an injective encoder on the modeled family.Practical decoder accuracy remains empirical.Source: §G.1 · p. 46, §G.2 · p. 48

G.2 LEARNED OBSERVABLE SPACES AND APPROXIMATE CLOSURE

  • Under compact Markov assumptions, the complete family of polynomial moment observables is dense in C(M; C) and is a core for the generator, but fixed dictionaries or trained encoders need separate assessment.The assumptions include compact X, M = P(X) with the weak topology, and a strongly continuous Markov semigroup on C(X).Source: §G.2 · pp. 46–47
  • Small finite-time closure residuals assess learned propagation, while exact closure yields Koopman eigenfunctions and approximate closure gives candidate eigenfunctions without bounding their distance to exact eigenvalues.Complex combinations of learned observables encode candidate eigenfunctions; their decay and oscillation correspond to the real and imaginary parts of the eigenvalue.Source: §G.2 · p. 47

G.3 OBSERVATION SUFFICIENCY AND IDENTIFIABILITY

  • Observed evolution is autonomous exactly when identical observed laws have identical future observed laws.If this condition fails, a predictor using only the current observed law faces ambiguity.Source: §G.3 · p. 48
  • Generator recovery from sampled snapshots requires informative sampling or frequency restrictions, and initial distributions must excite the observables being learned.Convergence and stability of joint encoder, operator, and decoder learning require further analysis within an identifiable model class.Source: §G.3 · p. 48

Related papers

This paper cites

Loading 2609.34629v1…