Discrete diffusion and flow matching models capture complex, non-additive and non-autoregressive structure in high-dimensional objective landscapes through parallel, iterative refinement. However, their implicit generative nature precludes direct integration with principled variational frameworks for online black-box optimisation, such as variational search distributions (VSD) and conditioning by adaptive sampling (CbAS). We introduce Active Flow Matching (AFM), which reformulates variational objectives to operate on conditional endpoint distributions along the flow, enabling gradient-based steering of flow models toward high-fitness regions while preserving the rigour of VSD and CbAS. We derive forward and reverse Kullback-Leibler (KL) variants using self-normalised importance sampling. Across a suite of online protein and small molecule design tasks, forward-KL AFM consistently performs competitively compared to state-of-the-art baselines, demonstrating effective exploration-exploitation under tight experimental budgets.
Figures & tables
Figure 1: Budget-constrained optimisation on Ehrlich synthetic landscapes and the AAV capsid design task. We report the simple regret
Figure 2: Structure-based protein design with FoldX oracles. We plot the best score found vs. oracle calls;
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Value / Schedule
Generative Flow & Classifier Architecture
Architecture
Bi-LSTM
Hidden Dimension
256
Embedding Dimension
128
Layers
3
Dropout
0.1
Appendix
Table 1: Hyperparameters and dynamic scheduling for Ehrlich experiments.
Parameter
Value / Schedule
Mixing Coefficients (Round r )
Replay ( αrb )
0.3 (fixed)
Flow ( αflow )
(1−αrb)×0.15 ( r<5 )
(1−αrb)×0.35 ( r≥5 )
Prior ( α0 )
1−αrb−αflow
Buffer Temp ( γ )
0.2
Appendix
Table 2: Hyperparameters for FoldX experiments. Unlisted parameters match Table 1 .
Parameter
Value / Schedule
Generative Flow Model
Architecture
Transformer
Layers
5
Heads
8
Hidden Dim
256
Dropout
0.1
Appendix
Table 3: Hyperparameters for AAV Capsid Optimisation.
Parameter
Value / Schedule
Generative Flow Model
Architecture
Bi-LSTM
Hidden Dimension
512
Embedding Dimension
256
Layers
4
Dropout
0.1
Appendix
Table 4: Hyperparameters for Dockstring (F2/Thrombin) experiments.
Flow Matching (FM) is a simulation-free method for learning a continuous, invertible flow that interpolates between two distributions, and in particular generates data from noise. Inspired by the variational nature of the diffusion process as a gradient flow, we introduce a stepwise FM model, Local Flow Matching (LFM), which sequentially learns a sequence of FM submodels, each matching a diffusion process up to the time-step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-flow model are closer than those in the full-flow matching model, which interpolates data to noise distributions, enabling smaller models with more efficient training. This variational perspective also allows us to prove a theoretical generation guarantee for the proposed flow model in terms of the χ2-divergence between the generated and true data distributions, leveraging the contraction property of the diffusion process. In practice, the stepwise structure of LFM is naturally amenable to model distillation, and various distillation techniques can be applied to accelerate generation. We empirically demonstrate that LFM achieves competitive generative performance compared to FM on unconditional generation of tabular and image datasets, and on conditional generation of robotic manipulation policies.
Chen Xu, Xiuyuan Cheng, Yao Xie
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University
Flow matching (FM) trains a time-dependent vector field that transports samples from a simple prior to a complex data distribution. However, for high-dimensional images, each training sample supervises only a single trajectory and intermediate point, yielding an extremely sparse and high-variance training signal. This under-constrained supervision can cause flow collapse, where the learned dynamics memorize specific source-target pairings, mapping diverse inputs to overly similar outputs, failing to generalize. We introduce Posterior-Augmented Flow Matching (PAFM), a theoretically grounded generalization of FM that replaces single-target supervision with an expectation over an approximate posterior of valid target completions for a given intermediate state and condition. PAFM factorizes this intractable posterior into (i) the likelihood of the intermediate under a hypothesized endpoint and (ii) the prior probability of that endpoint under the condition, and uses an importance sampling scheme to construct a mixture over multiple candidate targets. We prove that PAFM yields an unbiased estimator of the original FM objective while substantially reducing gradient variance during training by aggregating information from many plausible continuation trajectories per intermediate. Finally, we show that PAFM improves over FM by up to 3.4 FID50K across different model scales (SiT-B/2 and SiT-XL/2), different architectures (SiT and MMDiT), and in both class and text conditioned benchmarks (ImageNet and CC12M), with a negligible increase in the compute overhead. Code: https://github.com/gstoica27/PAFM.git.
George Stoica, Sayak Paul, Matthew Wallingford +6
Georgia Tech · University of Washington · Hugging Face +2
Standard flow matching scales well but typically relies on an unstructured source distribution, limiting its ability to learn interpretable latent structure. Latent-variable models, by contrast, capture structure but often sacrifice generative quality. We bridge this gap by proposing Structured Coupling for Flow Matching (SCFM), a cooperative framework that augments flow matching with structured latent representation learning. By introducing structured latent variables and exogenous noise into the source, SCFM jointly learns a structured prior (via latent variable modeling) and a continuous transport map (via flow matching). It uses a shared time-dependent recognition network for both latent variable model variational inference and intermediate-time flow velocity estimation. This yields a structurally informed yet unconditional, simulation-free flow model, where the latent variable model can also assist flow sampling. Empirically, SCFM facilitates unsupervised latent representation learning for clustering, disentanglement and downstream tasks, while remaining competitive with flow matching in sample quality, showing that meaningful structure can be learned without sacrificing generative fidelity.