PaperlayerView exact revision on arXiv

Source-linked AI summary

Follow the Mean: Reference-Guided Flow Matching

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

arXiv:2605.10302v3cs.LG

TL;DR

Controllable generation needs alternatives to fine-tuning, auxiliary networks, and test-time search. The paper shifts flow-matching endpoint means using reference examples, instantiating this as training-free Reference-Mean Guidance and semi-parametric guidance. These methods provide reference-controlled generation, including structural gains in frozen FLUX.2-klein and quality-preserving inference-time reference swapping in SPG.

  • Problem

    Controllable generation requires target-specific control, but existing methods rely on fine-tuning, auxiliary guidance, or test-time search.

  • Method

    Reference-guided flows steer pretrained flow-matching models by shifting endpoint means computed from reference sets, using either closed-form correction or an anchored residual refiner.

  • Results

    RMG improves prompt alignment in one sampling trajectory, with gains of +28.75 for position and +8.08 for two-object generation under matched model and sampling settings.

  • Takeaways & Limitations

    The results support controlling generation through data references rather than parameter updates, while retaining frozen-model operation for RMG and inference-time reference-set control for SPG.

  • Takeaways & Limitations

    Reference-Mean Guidance depends on reference-set quality, can incur costly posterior-mean computation for large sets, and may require domain-specific choices in other modalities.

Abstract

from arXiv · show

Existing approaches to controllable generation typically rely on fine-tuning, auxiliary networks, or test-time search. We show that flow matching admits a different control interface: adaptation through examples. For deterministic interpolants, the velocity field is solely governed by a conditional endpoint mean; shifting this mean shifts the flow itself. This yields a simple principle for controllable generation: steer a pretrained model by changing the reference set it follows. We instantiate this idea in two forms. Reference-Mean Guidance is training-free: it computes a closed-form endpoint-mean correction from a reference bank and applies it to a frozen FLUX.2-klein (4B) model, enabling control of color, identity, style, and structure while keeping the prompt, seed, and weights fixed. Semi-Parametric Guidance amortizes the same idea through an explicit mean anchor and learned residual refiner, matching unconditional DiT-B/4 quality on AFHQv2 while allowing the reference set to be swapped at inference time. These results point to a broader direction: generative models that adapt through data, not parameter updates.

1 Introduction

The paper introduces reference-guided flows, controlling pretrained generators by shifting endpoint means through reference examples rather than retraining, auxiliary guidance, or search. It develops training-free and semi-parametric forms for controllable generation.

  • Existing control methods rely on parameter updates, auxiliary signals, or test-time search, and none simultaneously avoids all three.Fine-tuning and adapters modify weights; guidance uses classifiers or rewards; search repeatedly samples, filters, or optimizes.
  • Reference-guided flows control generation by shifting the endpoint mean associated with a noisy interpolant.Because flow-matching velocity points toward this mean, changing it shifts the induced generated distribution.
  • Reference examples provide implicit guidance without requiring an explicitly defined reward or classifier.The references need only shift the mean in the desired direction rather than perfectly represent the target distribution.
  • Reference-guided control: A frozen text-to-image model changes an elephant’s generated color when conditioned on a small set of pink-elephant images.The prompt remains “an elephant in a jungle,” while the reference examples induce the color change.
  • Reference-Mean Guidance: Reference-Mean Guidance computes a closed-form endpoint-mean correction for training-free test-time control without extra training, auxiliary networks, gradients, or model evaluations.The method operates on a pretrained model while leaving those additional components out of the control procedure.
  • Semi-Parametric Guidance: Semi-Parametric Guidance combines an explicit posterior-mean anchor with a learned residual refiner to preserve useful reference information while suppressing nuisance correlations.It also preserves unconditional generation quality and permits inference-time reference-set control.

2 Background

Flow matching learns a continuous-time transport between source and target distributions through a time-dependent probability path. For linear interpolants, its velocity field can be represented through an endpoint mean, which has a closed-form empirical approximation as a weighted training-set sum.

  • Flow Matching: Flow matching learns a continuous-time transport model mapping a source distribution p0 to a target distribution p1.It defines a time-dependent distribution pt called the probability path.
  • Flow Matching: The velocity field transports samples along the probability path when it satisfies the continuity equation.The associated ordinary differential equation is ẋ = u_t(x).
  • Linear interpolants: For a linear interpolant, the velocity field moves samples toward the endpoint mean at a rate inversely proportional to the remaining time 1 − t.This gives the endpoint mean a direct role in determining the flow direction.
  • Equivalent parameterizations: The endpoint mean and velocity field are invertibly related, so pretrained flow models equivalently specify both parameterizations.The identity is μ_t(x) = x + (1 − t)u_t(x).
  • Equivalent parameterizations: Flow matching objectives can be written using either predicted velocity or predicted endpoint mean, with either standard or variational flow-matching loss.The training samples include x0 from p0, x1 from p1, and t uniformly from [0, 1].
  • Empirical endpoint mean: Training uses an empirical target distribution over a finite dataset, yielding a closed-form endpoint mean.For a standard normal source, the empirical endpoint mean is a weighted sum over training examples.
  • Empirical endpoint mean: The paper leverages the weighted empirical endpoint mean to design reference-guidance methods, although the observation itself was previously known.The stated novelty is applying this structure to guidance design.

3 Reference-Guided Flows

Reference-guided flows steer a pretrained flow by shifting its endpoint mean toward one induced by reference samples. The paper develops a closed-form reference-mean method and a semi-parametric anchor-plus-residual variant.

  • 3.1 Steering a Flow by Shifting the Endpoint Mean: The velocity fields of flows sharing a source and bridge differ through their endpoint means, so approximating the mean shift steers generation toward a target distribution.The pretrained model’s endpoint mean can be obtained directly or by inverting the flow equation.
  • 3.1 Steering a Flow by Shifting the Endpoint Mean: A reference set defines a target endpoint mean, with the goal of shifting generation while preserving the pretrained model’s diversity and quality.The reference set is sampled from a distribution and supplies the endpoint mean used for guidance.
  • 3.1 Steering a Flow by Shifting the Endpoint Mean: Geometric and arithmetic mixtures provide constructions for combining training and reference endpoint distributions under the same affine bridge.The geometric construction uses a scalar guidance schedule, while the arithmetic construction has an exact posterior mean by Bayes’ rule.
  • 3.2 Reference-Mean Guidance (RMG): Reference-Mean Guidance computes the reference-induced mean in closed form and shifts the velocity without auxiliary models or gradient computations.The method uses an empirical reference mean and can change the guided velocity by changing the reference bank.
  • 3.3 Semi-Parametric Guidance (SPG): Semi-Parametric Guidance combines a learned cross-attention anchor with a parametric residual refiner through time-dependent gates.The model is trained with interpolated samples conditioned on leave-one-out reference sets; stopping gradients through the anchor focuses learning on the positive residual.
  • 3.3 Semi-Parametric Guidance (SPG): The leave-one-out training structure prevents self-attention to the held-out sample, while experiments show that the reference set measurably controls generation at test time.The observed behavior suggests an implicit exchangeability structure in which samples are conditionally independent given an unobserved latent reference measure.

4 Results

The experiments validate reference-set control from mechanistic two-moons tests through frozen FLUX.2-klein and semi-parametric generation. Changing references steers attributes, structure, and output distributions while preserving fixed-model settings and, for SPG, unconditional quality.

  • 4.1.1 Mechanistic Validation: Reference-set composition shifts the two-moons flow field and final attractor, isolating the posterior mean as a steering mechanism.The test varies only reference-set composition while using N = 500 samples; as few as M = 5 references approach the hard-filter upper bound.
  • 4.1.2 Training-Free Control in FLUX.2-klein (4B): On frozen FLUX.2-klein, fixed prompts, seeds, and weights still produce systematic shifts in color, object identity, and style when reference sets change.Each reference set contains 20 images encoding a target attribute.
  • 4.1.2 Training-Free Control in FLUX.2-klein (4B): RMG transfers coarse structural priors in keyhole, hand, and gymnastics examples without copying reference content or requiring gradients, retraining, or extra model evaluations.The keyhole correction changes the global silhouette while preserving the interior scene; broader quantitative evaluation remains open.
  • 4.1.2 Training-Free Control in FLUX.2-klein (4B): +28.75 on position and +8.08 on two-object generation are the largest reported RMG gains in GenEval under matched backbone, sampler, prompts, and seeds.RMG uses a single sampling trajectory without classifier, reward model, LLM, gradients, or candidate selection.
  • 4.2 Semi-Parametric Guidance: SPG matches DiT-B/4 on AFHQv2 while reference-set swaps control generated class proportions at inference time.Generated proportions track the reference composition across reference sizes, and outputs remain semantically aligned with references without being copies.

5 Related Work

The paper positions reference-guided flows against parameter updates, auxiliary-model guidance, search, and retrieval-based control. Its distinction is a closed-form drift correction derived from endpoint means, without reward signals, auxiliary models, or additional evaluations.

  • Existing control interfaces: Prior control methods use fine-tuning, auxiliary models or rewards, or search, and do not jointly avoid training, auxiliary networks, and test-time search.Endpoint-posterior methods likewise use scalar rewards and require training or repeated evaluation.
  • Semi-parametric guidance: SPG demonstrates the same reference-set interface while preserving generation quality and allowing inference-time distribution shifts through reference-set substitution.Its figures report AFHQv2 quality matching and systematic shifts under swapped references.
  • Reference-guided flows: Reference-guided flows use reference-set endpoint means to produce a closed-form drift correction without reward signals, auxiliary models, or additional evaluations.Under a Gaussian bridge, the posterior mean is a softmax-weighted aggregation over reference points.

6 Limitations

Reference-mean guidance is constrained by the quality and size of its reference set, and extending it across modalities may require domain-specific choices.

  • Scope and practical constraints: Noisy or poorly curated references can introduce unwanted artifacts, while posterior-mean computation over large sets can be costly.The paper notes subsampling and approximate retrieval as practical remedies.
  • Scope and practical constraints: Applying the framework to other modalities may require domain-specific design choices.The paper also emphasizes responsible reference-set curation to reduce misuse for harmful or misleading generation.

7 Conclusion

The paper reframes controllable generation as shifting endpoint means and steering a flow by changing its reference set. Reference-Mean Guidance enables training-free control, while Semi-Parametric Guidance amortizes the same mechanism in a learnable architecture.

  • Shifting the endpoint mean shifts the induced flow, making reference-set choice a direct control mechanism.The velocity field is determined by the conditional endpoint mean under deterministic interpolants.
  • Reference-Mean Guidance computes a closed-form reference-mean correction for training-free control without auxiliary networks, gradients, or extra model evaluations.
  • The linear-bridge velocity specializes to ut(x) = (µt(x) − x)/(1 − t), but becomes singular as t approaches 1, motivating termination before the endpoint.
  • Under a Gaussian bridge, the endpoint mean becomes a softmax-weighted aggregation over reference points, algebraically matching cross-attention.

A.3 Proof of Proposition 3.3 (Reference-Mean Guided Dynamics)

The proof derives guided flow dynamics by expressing the velocity through conditional endpoint means and then interpolating training and reference posterior means. Empirical reference sets yield a closed-form softmax-weighted guidance rule, with explicit scope conditions.

  • The geometric mixture combines training and reference scores, and score-to-mean linearity produces a convex combination of their posterior means.
  • Substituting the mixed posterior mean into the affine-bridge parameterization yields the guided velocity used by reference-mean guidance.
  • For an empirical reference set, the posterior mean is computed from softmax weights over the selected references, giving a closed-form guided velocity.
  • The score-motivated marginal mixture is not generally an exact probability-flow construction; the endpoint-level mixture is valid by construction and exact under Gaussian posterior conditions.
  • Under a Gaussian bridge, the endpoint posterior mean is a Nadaraya–Watson estimator and is algebraically equivalent to cross-attention over data points.

C.1 SPG Architecture and Training

SPG uses cross-attention to form a reference posterior-mean anchor and a learned residual refiner to predict endpoints. Its design supports swapped reference sets while retaining unconditional generation quality.

  • Architecture: SPG augments flow matching with a reference-set attention module that approximates the posterior-mean anchor, followed by a learned residual refiner.
  • Architecture: The endpoint prediction combines the noisy state, attention anchor, and time-gated residual correction.
  • Training: The refiner is trained on the positive residual between the ground-truth endpoint and anchor, with gradients stopped through the anchor.
  • Training: During training, references are sampled from the data distribution, while inference-time distribution shifts in the reference set steer generation.
  • RMG implementation: FLUX.2 experiments keep the prompt, noise seed, and model weights fixed while changing only the reference set.

D Mechanistic Validation

Controlled experiments isolate the posterior mean as the mechanism by which reference-set changes steer flows. Additional ablations show how reference size and guidance schedules affect steering, stability, and diversity.

  • Two-moons: Changing only the reference composition reverses the flow direction and changes the final attractor, directly evidencing posterior-mean control.
  • Mechanistic validation: M = 5 labeled references already produce strong soft-posterior steering, while M = 50 soft-labeled MNIST references provide reliable class control.
  • Schedule ablation: The quadratic schedule front-loads guidance, decays to zero near t = 1, and is used in the main experiments to cancel late-time divergence.
  • Guidance strength: Above β0 ≈1, the constant schedule develops oversaturated colors and structural distortion near t = 1, whereas bell-shaped guidance remains stable at larger strengths.
  • Reference-set size: The diversity experiment varies only reference count while fixing the prompt, model, schedule, and seeds, and measures variation using average pairwise LPIPS.
  • Reference-set size: Diversity increases consistently with reference-set size, with larger sets producing broader outputs instead of collapsing to a single mode.

E.5 Number of Function Evaluations (NFE)

The experiments vary solver evaluations and guidance strength while holding the prompt, reference set, and random seeds fixed. They also test whether prompt and reference-set controls act independently.

  • NFE and guidance-strength setup: The ring-leap task sweeps NFE and guidance strength under a fixed prompt, reference set, random seeds, and quadratic-decay guidance schedule.The tested guidance strengths are β0 ∈ {0.1, 0.2, 0.4, 0.5, 1.0}.
  • NFE and guidance-strength setup: Runtime scales approximately linearly with NFE, exposing a practical trade-off between control quality and computational cost.
  • Prompt–reference interaction: Prompt and reference sets provide independent, composable control axes: neutral references suppress prompted pinkness, while pink references introduce or amplify it.
  • Prompt–reference interaction: The pinkness score measures similarity to “a pink elephant in a jungle” minus similarity to “a gray elephant in a jungle,” with higher values indicating stronger pinkness.

F.2 Reference Composition

Reference composition provides continuous control: mixing attribute-specific banks while fixing the prompt changes generated attribute prevalence and semantic alignment. Larger reference sets also broaden SPG diversity, while SPG can suppress shared-background artifacts.

  • Reference-composition setup: Reference sets mix two attribute-specific banks while the prompt and sampling procedure remain fixed, with target proportions varied from 0% to 100%.The standard proportions are {0, 25, 50, 75, 100}%, and each bank contains 20 images unless otherwise stated.
  • Evaluation metrics: Attribute frequency is estimated with Qwen2-VL-7B, while CLIP provides discrete target assignment and continuous similarity to target and distractor prompts.
  • Reference-composition results: Generated target-attribute prevalence increases correspondingly as the target attribute becomes more prevalent in the reference distribution.The qualitative grids illustrate the same progression in generated samples.
  • SPG diversity: LPIPS increases with SPG reference-set size, indicating that larger reference sets support more diverse generations rather than retrieval-like collapse.
  • Nuisance-artifact suppression: With shared white-background references, RMG transfers the nuisance background, whereas SPG preserves object-level guidance without copying it.
Loading 2605.10302v3…