Follow the Mean: Reference-Guided Flow Matching

Pedro M. P. Curvo; Maksim Zhdanov; Floor Eijkelboom; Jan-Willem van de Meent

NeurIPS 2026 · Main conference

Follow the Mean: Reference-Guided Flow Matching

Pedro M. P. Curvo
Maksim Zhdanov
Floor Eijkelboom
Jan-Willem van de Meent

University of Amsterdam

Paper PDF Code AlphaXiv

The same initial noise leads to a regular elephant or a pink elephant when guided by pink reference examples. — The reference set supplies the attribute to transfer. The prompt, initial noise, and model weights stay fixed while guidance changes the result.

Show it examples. Change what it generates.

Abstract

Flow matching offers a direct interface for controlling generation: the conditional endpoint mean determines the velocity field. Changing the examples used to form this mean changes the flow, making a reference bank a practical control signal.

Reference-Mean Guidance (RMG) applies a closed-form correction to a frozen model, controlling color, identity, style, and structure without retraining. Semi-Parametric Guidance (SPG) combines a mean anchor with a learned residual refiner, retaining generation quality while allowing reference sets to change at inference time.

Key contributions

Control through examples
The endpoint-mean view connects a reference bank directly to the flow model’s velocity.
Training-free guidance
RMG steers a frozen FLUX model through a closed-form correction, with no gradient search or auxiliary reward model.
Swappable reference sets
SPG learns a residual around an explicit mean anchor, allowing the reference bank to be replaced during inference.

Reference-guided generation

The velocity of a rectified-flow model is determined by its conditional endpoint mean. Shifting this mean changes the generation trajectory. Under a Gaussian bridge, a finite reference set provides an exact, closed-form endpoint mean: a posterior-weighted sum of the reference images.

A velocity correction toward this reference mean transfers attributes such as color, style, or structure. The references need only encode the desired attribute; they can differ from the text prompt and are supplied at inference time without retraining the model.

Method

The mechanism has three steps, each requiring only what the frozen model already computes.

01

Recover the endpoint mean

For any pretrained rectified-flow model, the endpoint mean is:μ_t(x) = x + (1 − t)u_t(x)No extra computation is needed: it falls out of the velocity prediction.

02

Compute a reference mean

Given M reference images, compute their posterior mean in closed form as a softmax-weighted sum. This is exact under a Gaussian bridge: the same math that underlies cross-attention.

03

Apply the correction

Add a scheduled velocity correction proportional to the difference between reference mean and model mean. No gradients, no backward pass, no auxiliary model.

Derivation

The three steps above have a clean derivation. This is where the endpoint-mean correction comes from, and why the reference mean takes the form of a softmax over examples.

Linear bridge

x_t = (1 - t)x₀ + tx₁

Flow matching begins with a straight path from source noise to a data endpoint.

We want a target distribution that interpolates between the model's prior and the reference density. A geometric mixture does this: β = 0 recovers the base model, β = 1 follows the reference entirely. This is an approximation — exact computation under the bridge is intractable — but under the Gaussian posterior assumption that holds in FLUX.2's VAE latent space, it gives the closed-form result that follows.

≡Cross-attention over R

Reference-Mean Guidance

Training-free. Closed-form. Applied to a frozen 4B-parameter FLUX.2-klein model. The next three sections show what it does.

Reference-set control

Take the prompt "elephant in a jungle." A frozen FLUX.2-klein generates a reasonable image — grey elephant, green canopy. Now hand it twenty images of pink toy elephants. None of them contain jungles. None of them match the prompt. The prompt hasn't changed. The seed hasn't changed. The 4-billion model weights haven't changed. The output is now pink.

Swap the reference bank to blue elephants. The output turns blue.

That's the whole interface. A folder of images. Swappable at inference time, with no retraining. The nearest-neighbor column shows the closest image in the reference bank, so you can see that the model is being guided rather than simply copying.

The reference banks used below

Examples don't need to be perfect

Style change

Any reference set R that shifts the endpoint mean in the right direction works.

Style change: nearest neighbor — nearest neighbor

Structural control

Structural correctness is hard to define as a reward: VLMs struggle to judge whether a hand is correctly oriented or a silhouette matches a target shape. Reference-mean guidance sidesteps this entirely: just show it examples of the target structure. No pose estimator, no spatial loss, no gradient.

Keyhole shape

a miniature forest...all inside a keyhole on a black background

Keyhole shape: nearest neighbor — nearest neighbor

GenEval benchmark

This comparison is between control interfaces, not a controlled ablation. RMG uses a fixed 20-image visual reference bank per category; baselines must express the same constraint through text prompts, classifier scores, or reward gradients. All methods share the same backbone, resolution, sampler, prompts, and seeds. The largest gains appear on compositional categories: position (+28.75) and two-object generation (+8.08), where text alone is insufficient to ground the target structure.

Method	Time down	NFE down	Aux. down	Mean up	Two up	Position up	Attribution up
FLUX.2-klein (4B) baseline	1.00x	1x	-	80.10	91.41	65.25	58.75
Search-based
+ Prompt Opt.	7.87x	8x	8C+2L	84.18	95.45	69.75	64.00
+ Best-of-4	4.07x	4x	4C	83.35	95.96	67.75	65.25
+ SMC	6.17x	4x	81C	80.28	95.71	61.75	57.25
Gradient-based
+ ReNO	19.44x	4x	4C	83.46	93.18	65.50	64.75
Ours (RMG)	1.02x	1x	-	91.17	99.49	94.00	75.25

The largest gains are on the categories where text prompts consistently fail: position (+28.75 points) and two-object composition (+8.08). Showing the model examples of correct spatial arrangements works where describing them in text does not.

Semi-Parametric Guidance

Reference-Mean Guidance has a limitation. The closed-form correction uses all of the reference bank's information — including things you don't want. A bank of pink toy elephants, all photographed against a white studio backdrop, will steer toward pink backgrounds as well as pink subjects. RMG can't distinguish between the attribute you care about and the ones that happen to correlate with it.

Semi-Parametric Guidance fixes this by amortizing the same idea into a learned architecture. Instead of applying the closed-form correction directly, SPG uses a cross-attention anchor to summarize the reference set, and a learned residual refiner to decide what to keep and what to suppress. The reference set is still swappable at inference time — nothing about that changes. But now the model can learn which aspects of the reference distribution are relevant to the generation task.

Trained on AFHQv2, SPG matches DiT-B/4 baseline quality (FID 23.26 vs 23.11), meaning the reference-set anchor doesn't degrade generation. And by swapping the reference set at inference time, you get continuous control over what the model generates: show it only cats, it generates cats; show it only dogs, it generates dogs; show it a mixture, it tracks the mixture proportions.

Unconditional generation quality on AFHQv2

Method	FID down	KID down	IS up
DiT-B/4 baseline	23.111	0.012	6.554
SPG (ours)	23.256	0.013	6.227

SPG matches DiT-B/4 baseline (FID 23.26 vs 23.11), confirming the reference-set anchor does not degrade generation quality.

Inference-time control

Generated class proportions closely track the reference-set composition across a wide range of reference sizes M, demonstrating continuous control without modifying model parameters.

Discussion

If control is represented by parameters, adaptation requires optimization and can create interference between concepts. If control is represented by a reference set, adaptation becomes a data operation: adding, removing, or reweighting examples changes the posterior mean and therefore the flow. This makes reference-guided flows naturally suited to low-data regimes, personalization, and continual adaptation settings where the target distribution changes faster than one would like to retrain a model.

Fine-tuning, reward networks, test-time search — these are the standard answers to controllable generation. Reference-guided flow matching suggests a different answer: a reference set. Swappable, interpretable, and requiring no gradient.

BibTeX

@misc{curvo2026followmeanreferenceguidedflow,
  title={Follow the Mean: Reference-Guided Flow Matching},
  author={Pedro M. P. Curvo and Maksim Zhdanov and Floor Eijkelboom and Jan-Willem van de Meent},
  year={2026},
  eprint={2605.10302},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2605.10302}
}