[Paper] Geometric regularization of autoencoders via observed stochastic dynamics

Source: arXiv - 2604.16282v1

Overview

The paper tackles a classic challenge in modeling complex stochastic systems: how to learn a low‑dimensional representation (a “chart”) of high‑dimensional dynamics from short simulation bursts, and then use that representation to build accurate reduced‑order simulators. By exploiting geometric information already present in the observed data’s covariance, the authors devise a regularized autoencoder that learns both the manifold chart and the latent stochastic differential equation (SDE) governing the dynamics, achieving far better fidelity than prior autoencoder‑based approaches.

Key Contributions

• Geometric regularization via ambient covariance – Shows that the covariance matrix of short‑burst samples encodes the tangent space of the underlying manifold, and leverages this to construct a novel tangent‑bundle penalty.
• Three‑stage learning pipeline – Separates chart learning, latent drift estimation, and latent diffusion estimation, each equipped with consistency penalties that enforce a well‑behaved encoder‑decoder pair.
• ρ‑metric analysis – Introduces a function‑space metric weaker than the Sobolev (H^1) norm but provably sufficient for chart‑quality generalization, matching Sobolev rates up to logarithmic factors.
• Bias‑corrected drift estimator – Derives an encoder‑pullback target for drift using Itô’s formula, and proves that the usual decoder‑side drift formula incurs systematic error when the chart is imperfect.
• Theoretical error propagation – Under a (W^{2,\infty}) chart‑convergence assumption, demonstrates how chart errors translate into weak convergence guarantees for the reconstructed ambient dynamics and accurate estimates of radial mean first‑passage times (MFPTs).
• Empirical validation – Experiments on four synthetic manifolds (up to 201‑dimensional ambient space) show 50‑70 % reduction in MFPT error for rotation dynamics and state‑of‑the‑art performance on metastable Müller‑Brown Langevin dynamics, with up to an order‑of‑magnitude reduction in coefficient errors versus unregularized autoencoders.

Methodology

1. Data collection – Generate many short “burst” trajectories from the high‑dimensional stochastic system. Each burst provides a local cloud of points and an empirical covariance matrix ( \Lambda ).
2. Tangent‑bundle penalty – Use the range of ( \Lambda ) as an estimate of the local tangent space. Penalize the encoder’s Jacobian so that its column space aligns with this estimated tangent space, encouraging the learned chart to respect the underlying geometry.
3. Inverse‑consistency penalty – Add a loss term that forces the decoder to be a near‑inverse of the encoder, preventing drift in the latent representation that would otherwise accumulate over time.
4. Three‑stage training
   • Chart learning: Train the autoencoder with the geometric penalties to obtain a single smooth nonlinear chart ( \phi ).
   • Latent drift: Pull back the observed drift to latent space using Itô’s formula applied to ( \phi ), and train a neural drift model on this corrected target.
   • Latent diffusion: Similarly estimate the diffusion matrix in latent coordinates and train a diffusion model.
5. Simulation – Integrate the learned latent SDE, map its trajectory back to ambient space via the decoder, and evaluate dynamical quantities (e.g., MFPTs).

The authors back the pipeline with rigorous functional‑analysis arguments, showing that the introduced penalties induce the ρ‑metric and lead to provable generalization bounds.

Results & Findings

Key takeaways:

• The tangent‑bundle regularization dramatically improves the quality of the learned manifold, leading to more accurate drift/diffusion estimates.
• The encoder‑pullback drift estimator eliminates a systematic bias that plagues standard decoder‑side drift learning.
• Even with a single global chart (instead of many local patches), the method attains high fidelity, simplifying downstream simulation pipelines.

Practical Implications

• Reduced‑order modeling for high‑dimensional stochastic systems – Engineers can now train a single autoencoder to capture the essential dynamics of, e.g., molecular simulations, climate sub‑models, or robotics control systems, without resorting to expensive local‑chart constructions.
• Faster surrogate simulators – Once the latent SDE is learned, simulating long‑time behavior (e.g., rare event statistics) becomes orders of magnitude cheaper than running full‑scale Monte Carlo bursts.
• Improved reliability of data‑driven dynamical models – The geometric penalties provide a principled way to enforce physical consistency (tangent‑space alignment), reducing the risk of spurious drift that can destabilize downstream control or optimization pipelines.
• Plug‑and‑play with existing ML stacks – The three‑stage pipeline can be built on top of standard deep‑learning libraries (PyTorch, TensorFlow) and integrates naturally with existing autoencoder architectures, making adoption straightforward for ML‑savvy developers.

Limitations & Future Work

• Assumption of smooth, well‑sampled manifolds – The theoretical guarantees rely on a (W^{2,\infty}) convergence of the chart, which may be violated for highly noisy or discontinuous data.
• Scalability of covariance estimation – Computing and storing the full ambient covariance for very high dimensions (e.g., > 10⁴) can become memory‑intensive; approximate or low‑rank techniques would be needed.
• Single global chart – While effective on the tested manifolds, more complex topologies (e.g., manifolds with holes) may require multiple overlapping charts; extending the regularization to a multi‑chart setting is an open direction.
• Extension to non‑Gaussian burst statistics – The current method assumes the burst covariance captures the tangent space; bursts with strong anisotropic noise or non‑Gaussian distributions could degrade the tangent‑bundle estimate.

Future research could explore hierarchical multi‑chart regularization, stochastic covariance approximations for massive dimensions, and integration with physics‑informed neural networks to further tighten the link between learned dynamics and underlying governing equations.

Authors

• Sean Hill
• Felix X. -F. Ye

Paper Information

• arXiv ID: 2604.16282v1
• Categories: cs.LG, math.DS, math.PR
• Published: April 17, 2026
• PDF: Download PDF