ESS-Flow: Training-free guidance of flow-based models as inference in source space

TL;DR

ESS-Flow tackles conditional generation with pretrained flow-based models without retraining or gradients by performing Bayesian inference in the source space and sampling with Elliptical Slice Sampling. By exploiting the Gaussian source prior, it achieves gradient-free sampling from the target distribution $\pi(z) \propto g(T_\theta(z)) p(z)$, avoiding Jacobians and enabling use with non-differentiable simulations and quantization. The method, validated on materials design and protein structure prediction, outperforms several gradient-based and gradient-free baselines in terms of targeted property alignment and structural realism, while introducing a practical multi-fidelity extension for computational efficiency. This approach broadens the applicability of flow priors to scientific inverse problems, offering a flexible, training-free mechanism for controlled generation with principled Bayesian guarantees.

Abstract

Guiding pretrained flow-based generative models for conditional generation or to produce samples with desired target properties enables solving diverse tasks without retraining on paired data. We present ESS-Flow, a gradient-free method that leverages the typically Gaussian prior of the source distribution in flow-based models to perform Bayesian inference directly in the source space using Elliptical Slice Sampling. ESS-Flow only requires forward passes through the generative model and observation process, no gradient or Jacobian computations, and is applicable even when gradients are unreliable or unavailable, such as with simulation-based observations or quantization in the generation or observation process. We demonstrate its effectiveness on designing materials with desired target properties and predicting protein structures from sparse inter-residue distance measurements.

Figures (4)

• Figure 1: Illustration of ESS-Flow. We sample the target distribution $\pi(z)$ in the source space of flow-based models, which typically has a Gaussian prior $p(z)$. This allows gradient-free sampling with elliptical slice sampling, which defines an ellipse through the current MCMC iterate $z_i$ and a sample from the prior $\nu$, and moves to the next iterate by searching along the ellipse. At stationarity, the transformed ellipse passes through regions of high potential in the data space by construction.
• Figure 2: Conditional generation targeting a specific $y$ indicated by the dotted line. Prior samples are shown in gray with marginals along the border. Some D-Flow samples, which follow the source-space gradient, become trapped in disconnected manifolds (indicated by arrow), while ESS-Flow, being gradient-free, can avoid this problem.
• Figure 3: Sample property distributions. Property values are along the X-axis, prior property distribution is shown in gray, and target values are shown as dotted lines. ESS-Flow samples are near target values across all tasks.
• Figure 4: Ground truth protein structure PDB:7r5b, conditional sample with the lowest $\text{RMSD}_\text{gt}$ from each method, and their ELBO from Chroma (in parenthesis). While samples from ADP-3D and DAPS resemble partially collapsed polymers, ESS-Flow achieves a better trade-off between data fidelity and sample realism.

Theorems & Definitions (1)

• Proposition 1: Geometric convergence of ESS-Flow, Theorem 2.2 by Natarovskii2021