Geometric regularization of autoencoders via observed stochastic dynamics

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 $H^1$ 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 $201$ ambient dimensions reduce radial MFPT error by $50$--$70\%$ 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.

• The paper presents a method that uses observed ambient covariance to align latent chart tangent spaces, reducing drift and diffusion estimation errors.
• It employs a three-stage pipeline with tangent-bundle and inverse-consistency penalties to lower reconstruction errors and improve mean first passage time accuracy.
• Empirical and theoretical analyses confirm that the approach achieves optimal generalization rates, enabling robust reduced simulators from sparse high-dimensional data.

Geometric Regularization of Autoencoders via Observed Stochastic Dynamics

Introduction and Motivation

The study introduces a theoretically principled method for geometric regularization of autoencoders tailored for stochastic dynamical systems exhibiting slow or metastable evolution on unknown low-dimensional manifolds within high-dimensional ambient spaces. The core objective is robust model reduction: to learn a reduced simulator, from short bursts of high-dimensional observations, that accurately captures ambient drift, diffusion, and dynamical observables such as mean first passage times (MFPTs).

Existing methods such as ATLAS and manifold-based autoencoders face severe limitations. Linear chart-based models are afflicted by exponential scaling in the number of required landmarks (curse of dimensionality), uncontrolled geometric extrapolation, and computational bottlenecks during simulation. Pure reconstruction-based training for autoencoders, meanwhile, leads to charts unconstrained in their tangent space structure, which manifests as systematic bias and error amplification in learned stochastic differential equation (SDE) parameters.

Geometric Regularization via Ambient Covariance and the $\rho$-Metric

The method leverages a coordinate-invariant property of the observed ambient covariance $\Lambda$: its range at any point spans the tangent space of the unknown manifold, regardless of the parameterization. Thus, local eigenanalysis of $\Lambda$ yields a projection operator $P$ onto the tangent bundle, which is itself invariant to chart reparameterization.

The authors propose a three-stage autoencoding pipeline in which:

1. Stage 1: An autoencoder is trained with a composite loss containing the standard reconstruction penalty, a tangent-bundle penalty aligning the decoder's Jacobian-inferred tangent projector to that derived from $\Lambda$, and an inverse-consistency penalty enforcing that the encoder and decoder are approximate local inverses.
2. Stage 2: The latent drift $\mu$ for the reduced SDE is estimated by regressing against an encoder-pullback target (derived via Itô's formula for the learned encoder), which is shown to be strictly less biased than the decoder-side formula except in the ideal invertible (exact) case.
3. Stage 3: The latent diffusion is learned by regressing against the pullback of the ambient covariance through the encoder.

Central to this framework is the introduction of the $\rho$-metric, a function-space distance combining the $L^2$ difference in reconstructed ambient coordinates and the Frobenius norm difference between predicted and observed tangent projectors. The $\rho$-metric is strictly weaker than the Sobolev $H^1$ norm but, crucially, the authors prove that $\Lambda$0-based empirical risk minimization achieves the same statistical generalization rates for chart quality as full $\Lambda$1 supervision up to logarithmic factors.

Algorithmic Efficiency and Coordinate Invariance

All losses and targets in the pipeline are formulated to be coordinate-invariant—the chart geometry is regularized directly through observable quantities, sidestepping the need for ground-truth tangent or curvature labels.

Despite the need to compute $\Lambda$2 objects such as covariance matrices and projection operators, efficient computational techniques are detailed, including the avoidance of explicit Jacobian or Hessian formation in favor of structured products and Hessian-vector products. This reduces per-sample cost to $\Lambda$3 flops and $\Lambda$4 memory, supporting scalability to high-dimensional settings.

Theoretical Analysis

The theoretical contributions include a precise analysis of geometric consistency, coordinate invariance, and error propagation in the pipeline. The main result establishes that, under a mild lower-singular-value assumption on the decoder's Jacobian, the $\Lambda$5-metric admits strong oracle generalization guarantees. Specifically, $\Lambda$6-regularized ERM achieves an out-of-sample error of order $\Lambda$7 (for $\Lambda$8-smooth target charts), matching optimal Sobolev-type rates.

Bias analysis rigorously decomposes errors arising from the use of decoder-side drift targets, revealing the necessity of encoder-based pullbacks to avoid systematic error unless the chart is exact. The propagation of chart-level approximation errors through to ambient SDE coefficients and, ultimately, to weak convergence of the ambient dynamics under learned SDEs is established, contingent on $\Lambda$9 convergence of the chart sequence.

Empirical Validation

Experiments are comprehensive, encompassing four analytically defined manifolds (paraboloid, hyperbolic paraboloid, quartic dome, sinusoidal) embedded in up to $\Lambda$0 dimensions, and two classes of latent SDEs: rotation drift with state-dependent diffusion and overdamped Langevin dynamics in the metastable M\"uller–Brown potential.

Ablation studies probe the influence of the tangent-bundle penalty ($\Lambda$1), inverse-consistency penalty ($\Lambda$2), and standard baselines:

• Tangent-bundle penalty: Yields orders-of-magnitude reduction in tangent error and drift/covariance error relative to unregularized baselines, especially crucial at high dimension or with sparse training data.
• Tangent+Inverse consistency (T+F): Yields the lowest reconstruction, tangent, and end-to-end drift error under all experimental conditions. Notably, radial MFPT error is reduced by 50–70% relative to baseline, and inter-well MFPT under M\"uller–Brown is minimized on the majority of surface/transition pairs.

Importantly, tangent space alignment and inverse consistency are demonstrated to be jointly critical: F alone fails to control the tangent error in the absence of T, and T alone does not suffice for unbiased drift estimation without F.

When evaluated for extrapolation beyond the support of the training set, T+F again provides the slowest growth in reconstruction error, demonstrating improved geometric fidelity outside the training region.

Numerical Result Highlight:

Figure 1: Reconstruction error as a function of extrapolation distance $\Lambda$3 beyond the $\Lambda$4 training domain ($\Lambda$5), showing that T+F regularization maintains lower error compared to all other methods across four surfaces.

Implications and Future Directions

The framework enables more reliable learning of latent SDEs and manifold charts in scientific and engineering applications where only sparse, short-burst observations of high-dimensional stochastic systems are available. Theoretical results indicate that geometric supervision via observed stochastic dynamics provides regularization with optimally controlled generalization properties and proper error propagation to SDE parameterization and downstream statistics.

Practical implications include improved fidelity in reduced simulators for molecular dynamics, atmospheric modeling, or any application where system evolution is primarily constrained to a stiff, low-dimensional manifold.

Future directions include extensions to multi-chart atlases for globally nontrivial manifolds, second-order geometric regularization directly targeting Hessian alignment, and integration of empirical estimates for the ambient drift and covariance. Quantifying the induced error when replacing oracle coefficients with those estimated from finite samples is a key open statistical question.

Conclusion

The paper advances both the theory and practice of geometric learning for autoencoders applied to stochastic dynamical systems. By regularizing through observable dynamical covariances, geometric properties of the manifold are enforced in a coordinate-invariant manner, yielding provably optimal generalization and robust model reduction. The analytic and empirical results together provide a highly compelling case for geometric regularization via stochastic dynamics in manifold learning and latent SDE discovery (2604.16282).