Bridging the Simulation-to-Experiment Gap with Generative Models using Adversarial Distribution Alignment

Abstract

A fundamental challenge in science and engineering is the simulation-to-experiment gap. While we often possess prior knowledge of physical laws, these physical laws can be too difficult to solve exactly for complex systems. Such systems are commonly modeled using simulators, which impose computational approximations. Meanwhile, experimental measurements more faithfully represent the real world, but experimental data typically consists of observations that only partially reflect the system's full underlying state. We propose a data-driven distribution alignment framework that bridges this simulation-to-experiment gap by pre-training a generative model on fully observed (but imperfect) simulation data, then aligning it with partial (but real) observations of experimental data. While our method is domain-agnostic, we ground our approach in the physical sciences by introducing Adversarial Distribution Alignment (ADA). This method aligns a generative model of atomic positions -- initially trained on a simulated Boltzmann distribution -- with the distribution of experimental observations. We prove that our method recovers the target observable distribution, even with multiple, potentially correlated observables. We also empirically validate our framework on synthetic, molecular, and experimental protein data, demonstrating that it can align generative models with diverse observables. Our code is available at https://kaityrusnelson.com/ada/.

Figures (8)

• Figure 1: ADA aligns generative models trained on approximate simulation data with the real world by leveraging partial experimental observations. (1) ADA starts from a base generative model $\mu_{\text{base}}(x)$ trained on simulation data, such as molecular dynamics simulations using a classical force field. (2) ADA aligns the base model with multiple, potentially correlated partial experimental observations (e.g., radial distribution functions (RDFs) or nuclear magnetic resonance (NMR) measurements). (3) ADA yields an aligned model $\mu_{\theta}(x)$ that approximates the true underlying real-world distribution $\nu(x)$.
• Figure 1: Adversarial Distribution Alignment from a Pretrained Generative Model
• Figure 2: ADA finds the closest distribution $\mu_{\theta}^*$ to a base distribution $\mu_{\text{base}}$ that satisfies the observable constraints $o^{(i)}_\# \mu_{\theta} = o^{(i)}_\# \nu$ for all observables $i \in \{1,\dots,m\}$. The resulting feasible set is denoted by $\mathcal{M}_o(\mathcal{X})$ (shown in orange). The KL term in \ref{['eqn:full_objective']} acts as an entropic regularizer, reflecting the fact that the observable constraints in \ref{['eqn:obs_constraint']} do not uniquely specify the underlying distribution, since many states can map to the same observable values. The discriminator $f^{(i)}_\phi$ learns the observable Wasserstein distance to align the model with the distributional observable constraint.
• Figure 3: Comparison between expectation alignment (EA) and ADA on synthetic mixture-of-Gaussians benchmark. We align a base distribution, defined as a mixture of Gaussians centered at the corners of a cube, with a target distribution where the variances and mixture weights have been perturbed. For the observables, we use pairwise coordinate projections, yielding correlated, multimodal marginals. We report the $\ell_1$ pdf residual between target and generated distributions before and after alignment. ADA consistently achieves larger reductions in distributional error on multidimensional observables, even when EA uses three moments. This is taken from a kernel density estimate of 2000 samples.
• Figure 4: ADA shifts the simulated distribution to align with the experimentally measured Trp-cage structures by using only noisy, high-dimensional cryo-EM images. The base generative model is trained on MD simulations using a classical force field. ADA uses the cryo-EM images as observables to align with the experimental structures from the PDB. The long MD simulations run by a classical force field sample the unfolded states more often than seen in the 2JOF PDB entry, leading to a longer tail in the radius of gyration observable.

Theorems & Definitions (12)

• Theorem 2.1: Existence and uniqueness of the saddle point
• proof
• Theorem 2.2: Asymptotic convergence in Wasserstein
• proof
• Theorem 2.3: Asymptotic convergence to the constraint set
• proof
• Theorem B.1: Existence and uniqueness of the saddle point
• proof
• Theorem B.2: Asymptotic convergence in Wasserstein
• proof