Follow the Mean: Reference-Guided Flow Matching

Abstract: 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.

• The paper introduces a novel method for controllable generation by directly adjusting the endpoint mean using a reference set.
• It leverages both closed-form and semi-parametric guidance to achieve precise attribute, geometric, and compositional control without retraining.
• It demonstrates significant improvements in prompt alignment and diversity on benchmarks such as MNIST and AFHQv2.

Reference-Guided Flow Matching: A Mean-Shift Control Framework

Introduction

"Follow the Mean: Reference-Guided Flow Matching" (2605.10302) introduces a principled and practical mechanism for controllable generation in flow-matching (FM) generative models, leveraging explicit modification of the endpoint mean via a reference set. Contrasting with standard approaches that rely on parameter updates, auxiliary reward models, or expensive test-time search, this work formulates control as a direct function of the data distribution specified by user-selected references, without any need for model fine-tuning, additional network components, or repeated sampling.

Theoretical Framework

The central technical insight is that deterministic interpolants in FM are governed entirely by the conditional endpoint mean. The velocity field at each point in the FM bridge is directed towards this mean, and thus any shift in the mean induces a corresponding shift in the generative trajectory. This property underpins both closed-form test-time control ("Reference-Mean Guidance," RMG) and a semi-parametric variant with a learnable residual ("Semi-Parametric Guidance," SPG).

Closed-Form Reference-Mean Correction

Given an empirical reference distribution $\mathcal{R}$, the endpoint mean can be computed as a softmax-weighted sum over the reference set in latent space. The model’s velocity field is corrected by the difference between the endpoint mean of the reference samples and the original training data, scaled appropriately with respect to the interpolation timestep. This operation is efficient, requiring only a single pass through a frozen model and no auxiliary classifier or gradient computation.

Semi-Parametric Amortization

SPG introduces an attention-based mechanism to estimate the reference conditional mean, combined with a parameterized residual refiner. Notably, this anchor is formed by cross-attention between the current latent state and the reference set, and the refiner is explicitly trained to predict deviations from this mean, enabling robust suppression of reference-set artifacts while retaining fine-grained control.

Experimental Results

Mechanistic Validation and Steering

Controlled experiments on synthetic (two-moons) and canonical (MNIST) data demonstrate that varying the reference set alone—while keeping model, sampling noise, and prompts fixed—systematically changes the generated outputs. The mechanism is robust: even with small reference sets, generation is steered effectively, and the endpoint mean acts as a reliable control signal.

Image Generation with Pretrained FM Models

Applying RMG to a frozen FLUX.2-klein (4B) latent rectified-flow model yields precise control of attributes such as color, style, and object identity. By swapping reference banks corresponding to, e.g., "pink elephants" versus "blue elephants," outputs shift accordingly—demonstrating attribute control without modifying seeds, prompts, or model parameters. RMG further enables non-trivial geometric and pose-level control that is challenging for reward gradient or classifier-guided methods, including structural priors for complex compositions and articulated shapes.

Quantitative and Compositional Control

On the GenEval compositional text-to-image benchmark, RMG shows substantial improvements in prompt alignment relative to both parameteric and gradient/search-based baselines under equivalent computational budgets. The most pronounced gains (+28.75%) are observed in spatial compositional tasks, indicating that reference banks provide more efficient control signals than scalar rewards or classifiers in structurally complex regimes.

SPG and Distributional Control

SPG is evaluated on AFHQv2, matching the unconditional sample quality of a DiT-B/4 baseline (FID 23.26 vs. 23.11). Generated attribute proportions tightly track the composition of the reference set, confirming that reference-anchored control operates predictably at inference. Critically, increasing the size or diversity of the reference set increases perceptual diversity of outputs, avoiding mode collapse and retrieval-like degeneracy. Unlike RMG, SPG robustly suppresses nuisance correlations such as shared backgrounds, isolating only salient attributes.

Relation to Existing Work

The proposed framework generalizes and extends several conceptual strands:

• Fine-tuning and Parameter Transfer: Unlike methods such as LoRA or DreamBooth, no model parameters are updated; adaptation is achieved purely by modifying the data distribution at inference.
• Auxiliary-Model and Reward-Based Guidance: Existing approaches require classifier gradients or repeated reward evaluation [e.g., classifier-free guidance, Tilt Matching]; RMG achieves similar or better control through closed-form endpoint mean shifts, with no external models.
• Search and Prompt Optimization: Search-based and prompt-opt methods incur significant inference cost; RMG operates within a single deterministic sampling trajectory.
• Retrieval-Augmented Generation: RMG and SPG reinterpret retrieved samples not as additional context, but as template-defining modulations of the generative path itself. In SPG, the connection between cross-attention and kernel posterior averaging yields theoretically grounded amortization.
• Prior Work on Posterior Means: While prior analyses observed kernel-averaging structures in generative models, this is the first explicit exploitation of reference means as a test-time control interface.

Numerical Strengths and Empirical Claims

• Test-time Training-Free Steering: RMG enables strong attribute and compositional control in a frozen 4B parameter FM model, with prompt, seed, and weights fixed.
• Quality Preservation: SPG achieves FID nearly identical to a strong DiT-B/4 baseline while allowing reference-set swaps at inference.
• Structural and Semantic Control: Large gains in compositional alignment (+28.75% in position tasks), outperforming search and optimization interfaces.
• Diversity: Both RMG and SPG exhibit increasing output diversity with reference-set size, contradicting claims that stronger controllability implies reduced generative diversity.

Implications and Future Prospects

This data-driven control paradigm reframes adaptation and controllability in generative modeling. Practical implications include:

• Personalization and Concept Transfer: Users can achieve granular, compositional, or rare concept adaptation using only curated reference examples, avoiding data-hungry fine-tuning.
• Efficient Attribute Swap and Editing: The approach provides a new interface for interactive and low-latency editing of generative outputs by dynamically modifying reference banks.
• Modality Generalization: While the focus is vision and latent FM, the core principle applies to any domain where endpoint means admit closed-form or attention-like approximations.
• Connection to Nonparametric Bayesian Methods: The centrality of posterior mean control links FM to nonparametric estimation and could inspire further cross-pollination.
• Augmented Prompting: The composability of prompt and reference-set guidance hints at hybrid text/example-driven interfaces.

Challenges remain in scaling to extremely large reference sets (posterior computation), curating high-quality references for complex semantics, and adapting the framework to modalities with challenging metric structures. Mitigation of potential misuse—where reference set curation enables harmful generation—necessitates responsible frameworks for example selection.

Conclusion

Reference-guided flow matching establishes a theoretically grounded, efficient, and highly flexible mechanism for controlling FM-based generative models through direct manipulation of the endpoint mean. The approach unifies data-driven test-time adaptation and amortized control architectures, suggesting a future for generative modeling in which data selection—not parameter change—is the main axis of adaptation (2605.10302).