Minimum-Excess-Work Guidance

TL;DR

This work introduces minimum-excess-work (MEW) guidance, a physics-inspired regularization for pre-trained probability-flow generative models (e.g., diffusion models) to operate under sparse observational constraints by minimizing excess work. It derives two practical strategies—Observable Guidance, which aligns generated distributions with experimental observables while preserving entropy, and Path Guidance, which concentrates sampling on user-defined transition regions—and connects these to upper bounds in optimal transport via $W_2^2(p_0,p'_0)$ and $D_{KL}(p'_0\|p_0)$. The authors validate MEW on toy systems and a coarse-grained Boltzmann emulator (cgBE), achieving substantial reductions in KL divergence and improved folding free energy accuracy, while maintaining multi-modality and structural validity; path guidance further demonstrates robust sampling of high-energy transition states with MEW regularization, outperforming loss-guidance baselines. Overall, MEW provides a stable, principled alternative to fine-tuning in data-scarce scientific applications, bridging thermodynamic principles and modern generative architectures for efficient, bias-reducing sampling in molecular simulations and beyond.

Abstract

We propose a regularization framework inspired by thermodynamic work for guiding pre-trained probability flow generative models (e.g., continuous normalizing flows or diffusion models) by minimizing excess work, a concept rooted in statistical mechanics and with strong conceptual connections to optimal transport. Our approach enables efficient guidance in sparse-data regimes common to scientific applications, where only limited target samples or partial density constraints are available. We introduce two strategies: Path Guidance for sampling rare transition states by concentrating probability mass on user-defined subsets, and Observable Guidance for aligning generated distributions with experimental observables while preserving entropy. We demonstrate the framework's versatility on a coarse-grained protein model, guiding it to sample transition configurations between folded/unfolded states and correct systematic biases using experimental data. The method bridges thermodynamic principles with modern generative architectures, offering a principled, efficient, and physics-inspired alternative to standard fine-tuning in data-scarce domains. Empirical results highlight improved sample efficiency and bias reduction, underscoring its applicability to molecular simulations and beyond.

Figures (19)

• Figure 1: Schematic comparison of observable and path guidance. Both panels show the evolution of probability density over time $t$ (blue heat map) with marginal distributions $p(\mathbf{x}_1)$ and $p(\mathbf{x}_0)$ on the sides (blue: reference model, red dashed: guided model). (A) Observable guidance perturbs the score function (red arrows) to match experimental observables (yellow) with unknown data distribution $p(\mathbf{x}_0)$ using minimal excess work. (B) Path guidance steers sampling trajectories (black solid) toward specific regions defined by guidance samples $\mathcal{X}^g = \{\mathbf{x}_0^i\}_{i=1}^M$ (dotted grey).
• Figure 1: Metrics for $O(\mathbf{x})$ and KL divergence across models.
• Figure 2: Comparison of Probability Distributions Before and After Observable Guidance for a 1D Energy Potential. The top plot shows the probability distributions for three models: the biased reference model (blue), the ground truth model (yellow), and the guided model (red). Guidance helps to align the reference model with that of the ground truth model using only the expectation of an observable function (bottom) while minimizing excess work.
• Figure 3: Observable Guidance of a Protein Folding Model. (A) Folding free energy comparison between reference model (blue), experimental data (yellow), and guided model (red). (B) Free energy profiles as a function of N- to C-terminal C$^\alpha$ distance. (C) Ensemble of 50 generated protein structures colored by their energy.
• Figure 4: Path Guidance vs. Loss-Guidance for sampling Transition States. (A) Sample quality and diversity, measured by the Wasserstein Distance (WD) and Vendi score (VS), show that path guidance preserves diversity and quality even at high guiding strengths, whereas loss guidance deteriorates. (B) Without MEW regularization ($\gamma = 0$), the sampled transition states collapse and exhibit little to no diversity (VS). Regularization also improves sample quality (WD). (C) Guiding success rate, measured as the percentage of transition states sampled, for different regularization strengths.

Theorems & Definitions (14)

• Proposition 3.0
• Proposition 3.0
• Lemma A.1
• proof
• Proposition A.2
• proof
• Proposition A.3
• proof
• Proposition A.3
• proof