Geometric regularization of autoencoders via observed stochastic dynamics
| Authors | Sean Hill & Felix X. -F. Ye |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2604.16282 |
| Download | |
| Categories | cs.LG |
Abstract
Stochastic dynamical systems with slow or metastable behavior evolve, on long time scales, on an unknown low-dimensional manifold in high-dimensional ambient space. Building a reduced simulator from short-burst ambient ensembles is a long-standing problem: local-chart methods like ATLAS suffer from exponential landmark scaling and per-step reprojection, while autoencoder alternatives leave tangent-bundle geometry poorly constrained, and the errors propagate into the learned drift and diffusion. We observe that the ambient covariance~ already encodes coordinate-invariant tangent-space information, its range spanning the tangent bundle. Using this, we construct a tangent-bundle penalty and an inverse-consistency penalty for a three-stage pipeline (chart learning, latent drift, latent diffusion) that learns a single nonlinear chart and the latent SDE. The penalties induce a function-space metric, the -metric, strictly weaker than the Sobolev norm yet achieving the same chart-quality generalization rate up to logarithmic factors. For the drift, we derive an encoder-pullback target via Itô's formula on the learned encoder and prove a bias decomposition showing the standard decoder-side formula carries systematic error for any imperfect chart. Under a W^{2,\infty} chart-convergence assumption, chart-level error propagates controllably to weak convergence of the ambient dynamics and to convergence of radial mean first-passage times. Experiments on four surfaces embedded in up to ambient dimensions reduce radial MFPT error by -- under rotation dynamics and achieve the lowest inter-well MFPT error on most surface--transition pairs under metastable Müller--Brown Langevin dynamics, while reducing end-to-end ambient coefficient errors by up to an order of magnitude relative to an unregularized autoencoder.
Engineering Breakdown
Plain English
This paper addresses a critical problem in reduced-order modeling of high-dimensional stochastic systems: learning accurate low-dimensional representations from short bursts of observational data. The authors propose a three-stage pipeline that uses geometric constraints derived from ambient covariance to properly regularize autoencoders, ensuring the learned latent space respects the true tangent-bundle structure of the underlying manifold. Unlike prior methods like ATLAS (which suffer exponential scaling in landmarks) or standard autoencoders (which leave tangent geometry poorly constrained), their approach introduces tangent-bundle penalties and inverse-consistency penalties that prevent error propagation into learned drift and diffusion terms. The key insight is that the ambient covariance matrix Λ already encodes coordinate-invariant information about the tangent space, allowing them to construct explicit geometric regularizers without expensive per-step reprojection.
Core Technical Contribution
The paper's core novelty is a geometric regularization framework that enforces consistency between the autoencoder's learned latent space and the true tangent-bundle structure of the data manifold. Rather than treating the manifold chart learning as a black-box dimensionality reduction problem, the authors exploit the fact that the range of ambient covariance spans the tangent bundle, deriving coordinate-invariant penalties that constrain the encoder-decoder pair. This enables a clean decomposition: first learn a faithful chart via geometric constraints, then learn latent drift and diffusion without error accumulation from geometric mismatch. The three-stage pipeline (chart learning → latent drift → latent diffusion) with dual penalties (tangent-bundle + inverse-consistency) is the technical contribution that makes this tractable for long-timescale metastable systems.
How It Works
The pipeline operates in three distinct stages. First, an autoencoder with geometric regularization is trained to learn a chart (encoding/decoding pair) that respects the tangent-bundle structure: the tangent-bundle penalty forces the encoder's Jacobian to align with directions spanned by the ambient covariance, while the inverse-consistency penalty ensures φ(ψ(z)) ≈ z to maintain invertibility. Second, with the chart fixed, the latent drift (mean vector field) is estimated directly from the latent coordinates using standard SDE inference, benefiting from the geometric consistency of the chart. Third, the latent diffusion (volatility/noise structure) is learned similarly in latent space, with errors now minimal because the chart faithfully represents the manifold's geometry. The entire approach avoids expensive operations like ATLAS's per-step manifold reprojection or landmark exponential scaling, instead leveraging the ambient covariance as a built-in geometric signal.
Production Impact
For engineers building digital twins or surrogate models of complex physical systems—climate, molecular dynamics, fluid flows—this approach offers a concrete win: faster training pipelines that don't require manual landmark selection or repeated projection steps, coupled with mathematically grounded guarantees that learned dynamics won't drift off the true manifold. In practice, this means deploying reduced-order models trained on short observational bursts (hours of simulation or experiment) that remain accurate over long prediction horizons (days or months), which is impossible with naive autoencoders or local-chart methods. The three-stage design is modular: you can swap drift/diffusion learners without retraining the chart, reducing iteration time. Compute cost is lower than ATLAS by avoiding exponential landmark scaling, though slightly higher than vanilla autoencoders due to geometric penalty overhead (typically 10-20% training time increase). Integration with existing frameworks is straightforward—standard autoencoder + custom penalty layers + SDE learner.
Limitations and When Not to Use This
The method assumes the ambient covariance Λ reliably spans the tangent bundle, which may fail for systems with heterogeneous noise or highly nonlinear manifolds where low-variance directions are not spurious. The paper does not address how to choose the latent dimension d or validate whether the learned chart is truly injective, both critical for practical deployment; selecting d too small will force manifold projection, while too large wastes capacity. The approach targets metastable or slow-fast systems where the manifold structure is stable across timescales; for genuinely time-varying or bifurcating systems, retraining or online adaptation may be necessary. Finally, the paper is silent on how to handle partial or noisy observations, a common real-world constraint—the three-stage pipeline assumes access to clean, full-state ambient data in short bursts.
Research Context
This work sits at the intersection of manifold learning (building on diffusion maps and ATLAS) and scientific machine learning (surrogate modeling of SDEs, a problem central to climate and molecular dynamics). It addresses a gap in the literature: prior autoencoder methods for reduced-order modeling either ignore geometry entirely (leading to drift errors) or use expensive local-chart methods (ATLAS, LLE) that scale poorly. The paper builds on classical insights from differential geometry (tangent bundle constraints) and applies them to modern deep learning, much in the spirit of recent work on geometric deep learning and neural SDEs. This opens a research direction toward geometry-preserving autoencoders for other structured domains—parametric PDEs, Hamiltonian systems, stochastic processes with unknown generators—where the ambient structure provides free regularization signals.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::