Three Forms of Stochastic Injection for Improved Distribution-to-Distribution Generative Modeling

TL;DR

The paper tackles distribution-to-distribution learning when both source and target distributions are learned from unpaired data, identifying data sparsity as a core bottleneck for standard flow matching. It introduces three stochastic injections—two-stage transfer learning, perturbing the source with Gaussian noise, and perturbing the interpolant with stochastic noise—along with stochastic interpolants and a VAE-based latent model to densify supervision. Across five high-dimensional imaging datasets, the method delivers substantial gains, averaging $13$ FID points over deterministic flow matching and $9$ points over baselines, while also reducing transport costs and improving source-target alignment. This approach makes flow matching a practical and scalable tool for simulating scientifically meaningful distribution transformations in biology, medicine, astronomy, and beyond.

Abstract

Modeling transformations between arbitrary data distributions is a fundamental scientific challenge, arising in applications like drug discovery and evolutionary simulation. While flow matching offers a natural framework for this task, its use has thus far primarily focused on the noise-to-data setting, while its application in the general distribution-to-distribution setting is underexplored. We find that in the latter case, where the source is also a data distribution to be learned from limited samples, standard flow matching fails due to sparse supervision. To address this, we propose a simple and computationally efficient method that injects stochasticity into the training process by perturbing source samples and flow interpolants. On five diverse imaging tasks spanning biology, radiology, and astronomy, our method significantly improves generation quality, outperforming existing baselines by an average of 9 FID points. Our approach also reduces the transport cost between input and generated samples to better highlight the true effect of the transformation, making flow matching a more practical tool for simulating the diverse distribution transformations that arise in science.

Figures (6)

• Figure 1: Our objective is to learn a flow from source onto target distributions, given unpaired training samples. Left: In the illustrated example, we learn to simulate cellular response to a chemical intervention. Right: We introduce stochastic injections that alleviate the sparsity challenges of distribution-to-distribution learning from finite target and source training examples, by 1) transfer learning from the noise-to-target task, 2) perturbing source samples, and 3) perturbing the training interpolant.
• Figure 2: Our stochastic injections enable denser supervision, as visualized for $d=2$ConcentricShells.
• Figure 3: Sparsity stress tests on ConcentricSpheres demonstrate that flow matching struggles in high dimensions and with few training examples. Our stochastic injections help 'de-sparsify' the supervision for the flow model, learning a more robust, generalizable velocity field.
• Figure 4: Qualitative examples. The top row is from the source side of the test set. The middle row is the corresponding model generation with the top image as the source, using a flow model trained with all three forms of stochastic injection. The bottom row is a random sample from the target side of the test set.
• Figure 5: Ablations suggest that the sine-squared function (Figure \ref{['fig:gamma_func']}) at scale $a=1.0$ (Figure \ref{['fig:gamma_scale']}) is the best choice of interpolant noise schedule $\gamma$. For inference, we find that ODE sampling with source noising performs better than SDE sampling. All metrics are FID are reported on the test set of BBBC.

Theorems & Definitions (6)

• Lemma 1
• Theorem 1: Informal
• Theorem 2
• proof
• Lemma 1
• proof