Technical Summary: Follow the Mean: Reference-Guided Flow Matching

Problem Statement

Controlling the outputs of pretrained generative models (e.g., enforcing specific attributes, styles, or structures) remains a significant challenge. Existing approaches typically fall into three categories, each with distinct limitations:

1. Fine-tuning and Adapters: Require modifying model parameters for every new target, which is computationally expensive and inflexible for dynamic changes.
2. Guidance Methods: Leave the generator frozen but rely on auxiliary classifiers, reward models, or reward signals, introducing additional computational overhead and dependency on external networks.
3. Search-Based Methods: Avoid retraining but incur high inference costs through repeated sampling, filtering, or per-prompt optimization.

None of these approaches simultaneously avoids additional training, auxiliary networks, and test-time search while maintaining high-quality generation. This paper proposes an alternative formulation: Reference-Guided Flows, which controls generation by adapting the model through data examples rather than parameter updates or external reward signals.

Methodology

The core insight of the paper is that in Flow Matching (FM), the velocity field $u_t(x)$ is solely governed by the conditional endpoint mean $\mu_t(x) = \mathbb{E}[x_1 | x_t = x]$. For a linear interpolant, the velocity is given by:
$$u_t(x) = \frac{\mu_t(x) - x}{1 - t}$$
Consequently, shifting the endpoint mean shifts the entire flow trajectory. The authors propose that this mean can be shifted by conditioning on a reference set $R$ of images, effectively steering the model toward the distribution induced by $R$ without altering the model weights.

The paper instantiates this principle in two forms:

1. Reference-Mean Guidance (RMG)

RMG is a training-free method for controlling frozen pretrained models.

• Mechanism: It computes a closed-form correction to the velocity field based on the difference between the endpoint mean of the reference set ($\hat{\mu}^\rho_t$) and the endpoint mean of the pretrained model ($\mu^\theta_t$).
• Formula: The guided velocity is defined as:
  $$u^\pi_t(x) \approx u^\theta_t(x) + \beta_t \frac{\hat{\mu}^\rho_t(x) - \mu^\theta_t(x)}{1 - t}$$
  where $\beta_t$ is a guidance schedule and $\hat{\mu}^\rho_t(x)$ is the empirical weighted average of the reference set (computed via a softmax over the reference images).
• Properties: It requires no gradient computation, no auxiliary networks, and no additional model evaluations beyond the base model. The reference set can be swapped at inference time to change color, identity, style, or structure.

2. Semi-Parametric Guidance (SPG)

SPG is a learnable architecture that amortizes the reference-mean guidance concept.

• Architecture: The model $\mu^\theta_t(x_t, R)$ explicitly conditions on a reference set $R$ during training. It uses a cross-attention mechanism to compute an "anchor" $\bar{x}$ (analogous to the closed-form mean) and a learned residual refiner $f^\theta$ to correct for nuisance correlations (e.g., shared backgrounds) that the anchor might inadvertently copy.
• Training: The model is trained with a leave-one-out objective to prevent the anchor from attending to its own endpoint, combined with a residual loss to teach the refiner to capture details beyond the anchor.
• Properties: SPG matches the unconditional generation quality of standard models (e.g., DiT-B/4) while enabling inference-time control via reference-set substitution.

Key Contributions

1. Mean Shift as Guidance: The authors define guidance as a geometric interpolation between the probability path of the training set and that of a reference set. They prove that the resulting guidance term depends only on the difference between the endpoint means of the two distributions.
2. Reference-Mean Guidance (RMG): A training-free method that computes the velocity shift in closed form using a reference bank. It enables control over frozen models (demonstrated on FLUX.2-klein) without gradients, classifiers, or reward models.
3. Semi-Parametric Guidance (SPG): A controllable architecture that learns to approximate the reference-mean anchor and a residual correction. It preserves unconditional generation quality while allowing the reference set to be swapped at inference time, effectively suppressing nuisance correlations like shared backgrounds.

Results

The paper validates the approach through mechanistic analysis and extensive experiments:

• Mechanistic Validation: On the two-moons distribution and MNIST, the authors demonstrate that changing the reference set composition directly shifts the flow field and the final attractor, isolating the causal role of the posterior mean.
• Training-Free Control (RMG): Applied to a frozen FLUX.2-klein (4B) model, RMG successfully steers generation for color (pink/blue elephants), identity (giraffes/zebras), style (Van Gogh/sketches), and structure (keyhole shapes, hand poses).
  • GenEval Benchmark: RMG outperforms search-based (Best-of-4, SMC) and gradient-based (ReNO) baselines on compositional text-to-image tasks (e.g., position, two-object generation) while using a single sampling trajectory and no auxiliary models. It achieves a mean score of 91.17 compared to the baseline's 80.10, with significantly lower inference time and no external model calls.
• Semi-Parametric Guidance (SPG): Trained on AFHQv2, SPG matches the unconditional quality of a DiT-B/4 baseline (FID 23.26 vs. 23.11).
  • Control: Swapping the reference set (e.g., cat-only vs. dog-only) systematically shifts the output distribution. The generated class proportions closely track the reference set composition.
  • Artifact Suppression: Unlike RMG, which may transfer nuisance correlations (e.g., a white background from all reference images), SPG's residual refiner successfully preserves object-level guidance while suppressing such artifacts.

Significance and Claims

The paper argues that generative models can be adapted through data rather than parameter updates. By framing control as a problem of shifting endpoint means, the authors provide a simple, principled alternative to fine-tuning, auxiliary guidance, or search.

• For Frozen Models: RMG demonstrates that a pretrained model can be steered effectively by changing the reference set it implicitly follows, offering a "guide with examples, not rewards" paradigm.
• For Learnable Models: SPG shows that this mechanism can be amortized into a learnable architecture that maintains high generation quality while enabling flexible, inference-time control.
• Broader Direction: The work suggests a path toward generative models that adapt dynamically through data references, reducing the need for retraining or complex reward engineering for controllable generation.

The authors acknowledge limitations, noting that RMG inherits the quality of its reference set (noisy references introduce artifacts) and that computing means over large sets can be costly, though subsampling offers remedies. They emphasize the need for responsible curation of reference sets to prevent misuse.