Active Flow Matching (AFM) Framework

• Active Flow Matching (AFM) is a framework for online black-box optimization that reformulates variational objectives over implicit marginal densities into tractable conditional endpoint distributions.
• It employs discrete flow matching techniques combined with importance-weighted training and round-based updates to steer generative models toward high-fitness regions.
• AFM has demonstrated robust empirical performance in protein and molecular design by integrating gradient-based optimization with adaptive sampling of discrete combinatorial spaces.

Active Flow Matching (AFM) is a framework for online black-box optimisation with discrete flow matching and discrete diffusion models that replaces variational objectives over intractable marginal endpoint densities with objectives over tractable conditional endpoint distributions along the flow. In the formulation introduced for protein and small-molecule design, AFM enables gradient-based steering of implicit discrete flow models toward high-fitness regions while preserving the variational structure associated with variational search distributions (VSD) and conditioning by adaptive sampling (CbAS) (Grewal et al., 1 Mar 2026).

1. Problem setting and motivation

AFM is formulated for online active generation over a discrete combinatorial space $\mathcal{X}$, with an expensive oracle

$$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$

and a superlevel-set objective

$$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$

The optimisation protocol is round-based: at round $r$, one has data $\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$, maintains a search distribution $q^\phi_{\mathcal{D}_r}(x)$, samples a batch of candidate designs, evaluates them with the oracle, augments $\mathcal{D}_r$, and updates the search distribution (Grewal et al., 1 Mar 2026).

The motivation for AFM is specific to implicit discrete generative models. Autoregressive sequence models factorize

$p(x)=\prod_{i=1}^L p(x^i\mid x^{<i}),$

which commits to each token before future context is known. By contrast, discrete diffusion and discrete flow matching models generate through parallel iterative refinement and condition on the full sequence context at each step. This makes them attractive on highly epistatic landscapes such as protein fitness, but it also creates a variational obstacle: their marginal endpoint density is implicit and generally intractable, so the density ratios and score-function terms required by standard VSD and CbAS are unavailable (Grewal et al., 1 Mar 2026).

AFM addresses exactly this mismatch. Rather than requiring tractable $\log q_\phi(x)$ or $\nabla_\phi \log q_\phi(x)$, it reformulates the variational objective in terms of conditional endpoint distributions that are already parameterized by the discrete flow model. This preserves the use of principled variational optimisation while retaining the non-autoregressive, parallel-refinement structure of discrete flow models (Grewal et al., 1 Mar 2026).

2. Discrete flow matching substrate

AFM is built on the standard discrete flow-matching setup. The state space is

$$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$0

with source distribution $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$1, often chosen as an all-mask distribution or a uniform distribution, and target distribution $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$2. A coupling $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$3 is defined between $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$4 and $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$5; the simplest case is the independent coupling $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$6 (Grewal et al., 1 Mar 2026).

The conditional probability path is

$$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$7

with scheduler $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$8, $$y = f(x) + \epsilon,\qquad \epsilon \sim p(\epsilon),\ \mathbb{E}[\epsilon]=0,$$9, and $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$0. The associated marginal path is

$$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$1

The dynamics are represented as a continuous-time Markov chain with probability velocity

$$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$2

where $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$3 is the true endpoint posterior (Grewal et al., 1 Mar 2026).

A neural network $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$4 approximates this posterior. Standard discrete flow matching trains it by cross-entropy,

$$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$5

AFM does not replace this substrate. Instead, it changes the effective target distribution by reweighting endpoint conditionals toward high-fitness regions, while leaving the discrete flow machinery itself intact (Grewal et al., 1 Mar 2026).

3. Variational reformulation along the flow

The central idea of AFM is to replace KL divergences over the intractable marginal $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$6 with KL divergences over the conditional endpoint distribution $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$7. Given a classifier $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$8 and a base flow $$\mathcal{S}_\tau = \{x\in\mathcal{X}: y \ge \tau\}.$$9, AFM defines the target conditional

$r$0

This conditional expresses the high-fitness reweighting directly at the endpoint level (Grewal et al., 1 Mar 2026).

The forward-KL objective is

$r$1

Because direct sampling from $r$2 is unavailable, AFM uses self-normalised importance sampling (SNIS). With proposal $r$3, endpoint samples $r$4, weights

$r$5

and intermediate samples $r$6, the forward objective is approximated by

$r$7

Under the stated discrete-flow assumptions—masked source, convex interpolation, and strictly positive scheduler—the global minimizer of forward-KL AFM yields terminal distribution

$r$8

so the forward objective is consistent with the desired variational target (Grewal et al., 1 Mar 2026).

AFM also defines a reverse-KL objective,

$r$9

and a symmetric objective

$\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$0

The reverse-KL variant has no analogous end-to-end consistency theorem in the paper and behaves more mode-seekingly in experiments; the forward-KL variant is the theoretically preferred form (Grewal et al., 1 Mar 2026).

4. Proposal distributions and online optimisation loop

AFM uses importance-weighted training, so proposal design is operationally important. The proposal $\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$1 is a mixture of a prior component, a base-flow component, and a replay-buffer component:

$\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$2

The replay-buffer weights are sampled with

$\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$3

In practice, AFM samples from one proposal component per batch and uses the corresponding simplified importance weight formula (Grewal et al., 1 Mar 2026).

The full optimisation loop is round-based. At each round, the active flow proposes candidates by sampling from the source and simulating the discrete flow to the terminal endpoint. The oracle evaluates the batch, the dataset is updated, the classifier $\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$4 is retrained, and AFM is optimized for several gradient steps with the SNIS objective. Thresholds $\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$5 may be fixed or increased over rounds, so the search progressively concentrates on higher-fitness regions (Grewal et al., 1 Mar 2026).

This yields a variational search mechanism for models whose marginals are implicit. In contrast to guidance-based methods that modify the sampling dynamics at inference time, AFM absorbs oracle feedback into the training objective itself. The test-time generator remains the original discrete flow; what changes is the effective training distribution over endpoint conditionals (Grewal et al., 1 Mar 2026).

5. Empirical behavior and theoretical interpretation

AFM was evaluated on Ehrlich synthetic landscapes, AAV capsid design, FoldX structure-based protein design, and Dockstring molecular docking. The main empirical pattern is that forward-KL AFM is the most robust variant. On the Ehrlich landscapes it converges fastest to near-optimal regret; on AAV capsid design it achieves the lowest regret alongside symmetric-KL AFM; on FoldX stability it performs best, while on the SASA objective it underperforms VSD and CbAS; on Dockstring thrombin docking it substantially outperforms VSD. Reverse-KL AFM consistently lags and appears prone to collapse (Grewal et al., 1 Mar 2026).

The broader flow-matching literature gives a useful interpretation of these outcomes. Generator Matching places diffusion and flow matching inside a common Markov-generator framework and argues that flow-like generators inherit robustness from first-order transport dynamics rather than from the inversion of a smoothing process. This suggests that AFM’s active reweighting is especially compatible with flow matching because the underlying dynamics are less sensitive to approximation and discretization errors than reverse diffusions (Patel et al., 2024).

A separate theoretical result on flow matching shows that controlling the pathwise $\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$6 velocity error controls the terminal KL divergence and, through Pinsker’s inequality, the total variation distance. The paper also notes that the relevant quantity is the expected $\mathcal{D}_r=\{(x_n,y_n)\}_{n=1}^{N_r}$7 error with respect to the true path measure, rather than uniform time sampling specifically. This suggests that AFM-style importance weighting and non-uniform sampling can remain theoretically meaningful so long as they reduce the effective path-measure error rather than merely oversampling easy regions (Su et al., 7 Nov 2025).

6. Terminology and acronym ambiguity

In current arXiv usage, the acronym “AFM” is overloaded. In the present sense it denotes Active Flow Matching, but several unrelated flow-matching methods use the same initials.

Expansion	Role	Representative paper
Active Flow Matching	Online black-box optimisation with discrete flows	(Grewal et al., 1 Mar 2026)
Autoregressive Flow Matching	Probabilistic sequence forecasting via one-step conditionals	(Rogalla et al., 13 Apr 2026)
Adaptive Flow Matching	Conditional climate downscaling backbone in CorrDiff++-style models	(Debeire et al., 3 Apr 2026)
Asynchronous Flow Matching	Token-wise non-uniform time scheduling in VLA action generation	(Jiang et al., 18 Nov 2025)
Action Flow Matching	Online action refinement for continual robot learning	(Murillo-Gonzalez et al., 25 Apr 2025)

This ambiguity matters because the technical object called “AFM” changes substantially across domains. In “Autoregressive Flow Matching,” the acronym refers to temporal factorization of future states rather than active experimental design (Rogalla et al., 13 Apr 2026). In “Physics-Constrained Adaptive Flow Matching,” it refers to a conditional generative downscaling model with an encoder-defined base distribution and conservation losses (Debeire et al., 3 Apr 2026). In “Asynchronous Flow Matching,” it denotes token-wise asynchronous denoising schedules for robotic action chunks (Jiang et al., 18 Nov 2025). In “Action Flow Matching,” it denotes a conditional flow in action space that refines planned controls under model misalignment (Murillo-Gonzalez et al., 25 Apr 2025). Within this terminological field, Active Flow Matching is the specific variational framework that lifts VSD/CbAS-style optimisation to conditional endpoint distributions along a discrete flow (Grewal et al., 1 Mar 2026).