Technical Summary: Learning Stochastic Multiscale Models through Normalizing Flows

1. Problem Formulation

The paper addresses the challenge of learning effective dynamics in partially observed multiscale stochastic systems. These systems, common in physics, engineering, and biology, consist of a low-dimensional slow component ($X^{(n)}$) evolving under the influence of high-dimensional, rapidly fluctuating fast components ($Y^{(n)}$).

Key Constraints and Challenges:

• Partial Observability: Only a single trajectory of the slow variable $X^{(n)}$ is observed; the fast variables $Y^{(n)}$ remain unobserved and latent.
• Scale Separation: The fast dynamics evolve on a time scale of $O(1/n)$, making direct reconstruction of the latent trajectory computationally intractable (requiring tracking $O(n)$ events over a fixed horizon).
• Intractable Invariant Measures: The effective dynamics of the slow variable depend on the invariant distribution ($\pi_0$) of the fast process. While $\pi_0$ characterizes the equilibrium effect of the fast variables, it is typically unknown and defined implicitly by a high-dimensional, intractable partial differential equation (PDE) (e.g., the stationary Fokker-Planck equation).
• Limitations of Existing Methods: Traditional PDE-based approaches (Fokker-Planck) require dense space-time data or grid-based solvers, which are unsuitable for single-trajectory data. Standard model reduction techniques like PCA fail to respect the underlying dynamical structure.

The goal is to infer a data-driven effective drift ($\bar{b}_{\pi_0}$) and diffusion function for a reduced-order Stochastic Differential Equation (SDE) that mimics the observed slow dynamics, without explicitly knowing the fast dynamics or $\pi_0$.

2. Methodology

The proposed framework combines stochastic averaging for model reduction with normalizing flows for latent density modeling and variational inference for uncertainty quantification.

2.1 Model Reduction via Stochastic Averaging

The authors first establish a principled model reduction. Under Assumption 3.1 (existence of a Lyapunov function and uniform ellipticity), as the time-scale separation parameter $n \to \infty$, the slow process $X^{(n)}$ converges weakly to a reduced process $X^{\pi_0}$ governed by:
$$dX^{\pi_0}(t) = \bar{b}_{\pi_0}(X^{\pi_0}(t)) dt + \bar{\sigma}_{\pi_0}(X^{\pi_0}(t)) dW(t)$$
where the effective drift is the expectation of the original drift with respect to the invariant measure:
$$\bar{b}_{\pi_0}(x) = \int_{\mathbb{R}^{d'}} b(x, y) \pi_0(dy)$$
The inference problem is thus reformulated as learning the unknown invariant measure $\pi_0$ (or directly the functional $\bar{b}_{\pi_0}$) from the observed trajectory of $X$.

2.2 Normalizing Flow Parameterization

Since $\pi_0$ is an infinite-dimensional object, the authors parameterize it using normalizing flows.

• Latent Density: The invariant measure is modeled as a push-forward of a simple reference distribution $\nu_{ref}$ (e.g., standard Gaussian) through a neural network transformation $f_\theta$:
  $$q_\theta = (f_\theta)_\# \nu_{ref}$$
• Expressivity: Leveraging the Universal Approximation Theorem, the authors show that as the complexity of the neural network increases, the family of flow-induced distributions becomes dense in the space of probability measures under the Wasserstein metric. This allows the approximation of complex, non-Gaussian, and multimodal invariant distributions.
• Estimator Construction: The effective drift is approximated via Monte Carlo integration:
  $$\bar{b}_{\theta}^{MC}(x) = \frac{1}{L} \sum_{l=1}^L b(x, f_\theta(Z^{(l)}))$$
  where $Z^{(l)} \sim \nu_{ref}$.

2.3 Optimization and Learning

The learning objective is a penalized maximum likelihood estimation (MLE).

• Likelihood: The likelihood of the observed trajectory is approximated using the Euler-Maruyama discretization.
• Penalization: To ensure stability and prevent overfitting (specifically, to control the moments of the learned distribution), a penalty term $\lambda m_p(\rho)$ (based on the $p$-th moment of the distribution) is added to the negative log-likelihood.
• Convergence: The paper proves that as the neural network complexity ($n$) increases and the penalty parameter ($\lambda$) decreases, the flow-based estimators converge to the set of global MLEs.

2.4 Bayesian Uncertainty Quantification

To address epistemic uncertainty, the authors develop a variational inference procedure:

• Posterior Approximation: A second normalizing flow, $g_\vartheta$, is used to approximate the posterior distribution over the flow parameters $\theta$.
• ELBO Optimization: The Evidence Lower Bound (ELBO) is maximized using the reparameterization trick, enabling efficient gradient-based optimization. This yields a scalable method to quantify uncertainty in the inferred effective dynamics.

3. Key Contributions

1. Formulation: The problem of inferring multiscale SDEs is formulated as a latent-invariant measure learning problem, where the effective drift is expressed as a conditional expectation over an unknown invariant distribution.
2. Flow-Based Parameterization: Introduction of normalizing flows to parameterize the latent invariant measure, facilitating accurate approximation of the averaged functions governing reduced dynamics.
3. Bayesian Extension: Development of a flow-based variational posterior to quantify uncertainty in the inferred effective dynamics, providing a scalable alternative to MCMC in high-dimensional parameter spaces.
4. Theoretical Guarantees: Proof that penalized minimizers over increasingly expressive neural pushforward classes converge to the global MLE set under suitable conditions.
5. Empirical Validation: Demonstration of the method on synthetic multiscale systems (tagged particle in solvent and double-well dynamics), showing accurate recovery of effective drifts and meaningful uncertainty estimates.

4. Experimental Results

The method was validated on synthetic data generated from multiscale SDEs:

• Drift Recovery: The proposed structured estimator significantly outperformed an unstructured neural network baseline (which ignored the integral representation of the drift). In specific tests, the structured flow-based estimator achieved Mean Squared Errors (MSE) approximately 10 times lower than the unstructured NN.
• Trajectory Matching: The SDEs driven by the learned averaged drifts successfully mimicked the true slow trajectories, both within and beyond the observation time window, indicating genuine learning rather than overfitting.
• Uncertainty: The Bayesian extension provided confidence bands that captured the variability in the estimated drift.
• Finite-Time Laws: Comparisons of empirical laws (using Kolmogorov-Smirnov and Wasserstein distances) showed that the learned model accurately reproduced the finite-time probabilistic structure of the averaged system.

5. Significance and Limitations

Significance:
The paper bridges stochastic model reduction and flow-based probabilistic inference. It offers a scalable, data-driven framework for learning effective dynamics in partially observed multiscale systems where traditional PDE solvers fail due to dimensionality and data scarcity. By treating the invariant measure as a learnable latent variable, the method respects the dynamical structure of the system while leveraging modern deep learning capabilities.

Limitations and Future Work (as stated by the authors):

• Strong Coupling: The current approach assumes the fast dynamics are independent of the slow variables (or weakly coupled in a specific averaging regime). It does not directly apply to strongly coupled systems where fast dynamics depend on slow variables.
• Data Quality: The method assumes high-frequency observations of the slow component. It is not yet validated for settings with sparse and noisy observations.
• Architecture: Future work is suggested to extend the method using more expressive architectures, such as continuous-time flows (Neural ODEs), to handle more complex regimes.