PaperlayerView exact revision on arXiv

Source-linked AI summary

MeanFlowAdvantage: Stable Reward Fine-Tuning for Few-Step Average-Velocity Generators

Haocheng Tang, Tianchi Xie, Xingqiao Lin

arXiv:2609.37670v1cs.AIcs.CEcs.CV

TL;DR

MeanFlow’s few-step sampler directly uses learned average-velocity maps, while existing advantage-based reward objectives typically target instantaneous velocities or equivalent prediction spaces. MeanFlowAdvantage uses a shared detached correction to align reward supervision and regularization with those maps; it improves all eight reported metrics over four-step MeanFlowNFT and matches or exceeds 40-step DiffusionNFT on six, while also supporting DNA promoter RL and distillation.

  • Problem

    MeanFlow’s interval-average representation enables few-step sampling, but existing advantage-based objectives typically target instantaneous velocities or equivalent prediction spaces rather than the map used at inference.

    Source: §1 · p. 1

  • Method

    MeanFlowAdvantage applies signed advantage-weighted least squares in prediction space with a shared detached correction, while anchoring the average-velocity model to rollout and reference maps.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    On SD3.5-Medium, MFA improves all eight reported metrics over four-step MeanFlowNFT and matches or exceeds 40-step DiffusionNFT on six of eight; on DNA promoters, reward-guided distillation achieves the best Sei profile error among compared configurations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The formulation supports both on-policy reward fine-tuning and reward-guided distillation while keeping regularization aligned with the average-velocity maps used by MeanFlow inference.

    Source: §1 · p. 2, §2 · p. 3

    The analysis assumes an exact correction, a calibrated rollout, and free functions, whereas training uses finite differences, LoRA, and sample-dependent scaling.

    Source: §4 · pp. 5–6

Abstract

from arXiv · show

MeanFlow enables efficient few-step generation by predicting interval-average velocities, but this representation creates a mismatch for reward fine-tuning: existing advantage-based objectives are typically defined on instantaneous velocities or equivalent $x_0$-space predictions, whereas inference directly uses the learned average-velocity map. We introduce MeanFlowAdvantage, a signed advantage-weighted least-squares objective for average-velocity generators. Our key construction uses a shared, detached MeanFlow derivative correction to express the reward objective in prediction space while making rollout and reference regularization exact penalties on the average-velocity network deployed at inference. The resulting formulation preserves MeanFlow's native few-step sampler and provides a direct mechanism for transferring reward improvements to the deployed flow map. On SD3.5-Medium, MeanFlowAdvantage improves all eight reported metrics over the matched four-step MeanFlowNFT baseline and, with only four NFEs, matches or exceeds the 40-step DiffusionNFT baseline on six of eight metrics. The same objective also transfers to DNA promoter design, where it supports both teacher-free on-policy RL for a generator defined on a manifold and teacher-guided reward-graded distillation, with the latter yielding the lowest one-step Sei profile MSE among the compared configurations.

1 INTRODUCTION

MeanFlow’s average-velocity representation enables few-step sampling but leaves a mismatch with reward objectives defined for instantaneous predictions. MeanFlowAdvantage addresses this mismatch, with reported gains on image-generation benchmarks and DNA promoter tasks.

1 INTRODUCTION

  • MeanFlow’s interval-average velocity enables finite-interval sampling, but creates a mismatch because existing advantage objectives typically target instantaneous velocities or equivalent prediction spaces.A useful reward objective should improve sample quality while remaining aligned with the average-velocity map used at inference.Source: §1 · p. 1
  • MeanFlowAdvantage is a signed advantage-weighted least-squares objective that uses prediction-space reward supervision while anchoring the average-velocity model to rollout and reference policies.A shared stop-gradient correction makes the regularization terms direct differences between MeanFlow velocity predictions.Source: §1 · p. 2
  • MFA improves all eight reported metrics over four-step MeanFlowNFT and matches or exceeds 40-step DiffusionNFT on six of eight metrics using four NFEs.The comparisons are on SD3.5-Medium.Source: §1 · p. 2
  • On DNA promoters, the objective supports teacher-free RL and reward-guided distillation, with the latter achieving the best Sei profile error among compared configurations.Teacher-free RL uses a manifold generator; distillation transfers graded teacher improvements into a one-step student.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries contrast flow matching’s instantaneous-velocity representation with MeanFlow’s finite-interval average velocities. They establish that reward fine-tuning must account for the average-velocity maps directly used by the MeanFlow sampler.

2 PRELIMINARIES

  • Flow matching represents transport with a time-dependent instantaneous velocity field, whose probability-flow ODE is integrated from noise to data at inference.Rectified flow uses a linear interpolation between data and noise.Source: §2 · p. 2
  • AdvantageFlow’s signed prediction-space objective is natural for instantaneous-velocity models, but MeanFlow’s finite-interval sampler requires fine-tuning to account for its average-velocity map.The objective is strictly convex in the prediction when A + γ + λ > 0, including when A < 0.Source: §2 · p. 3
  • MeanFlow inference uses the predicted average velocity directly in the finite-interval map, rather than integrating an instantaneous velocity field.For the exact average velocity, the map is xs = xt − (t−s)u(xt, s, t); the learned sampler substitutes uθ.Source: §2 · p. 3

3 MEANFLOWADVANTAGE

MeanFlowAdvantage connects prediction-space reward supervision to the finite-interval average-velocity maps deployed by MeanFlow. Its shared correction makes rollout and reference anchoring direct in that representation, while the training procedure samples positive-length intervals as well as boundary intervals.

3.1 LINKING AVERAGE VELOCITY TO PREDICTION SPACE

  • A shared, detached derivative correction converts the average-velocity prediction into the induced instantaneous-velocity representation used for reward supervision.The correction is estimated from the rollout network along the sampled direction and shared across learner, rollout, and reference models.Source: §3.1 · p. 3
  • Because all three networks share the correction, pairwise prediction-space differences reduce to differences between their average-velocity predictions on the sampled interval.Using separate derivatives would add a residual term dependent on the sampled direction.Source: §3.1 · p. 4

3.2 ADVANTAGE-WEIGHTED LEAST SQUARES ON THE INDUCED PREDICTION

  • The resulting loss is exactly an advantage-weighted MeanFlow regression, with rollout and reference anchors directly penalizing deviations in the deployed average-velocity maps.Positive-length intervals are the new setting for few-step generators; at s = t, the loss reduces to AdvantageFlow on the boundary velocity.Source: §3.2 · p. 4

3.3 SAMPLES, ADVANTAGES, AND INTERVALS

  • Training uses N = 4 rollout steps, signed advantages, and a mixture of boundary, full-jump, and general intervals so it trains both boundary velocities and sampler-used maps.The interval types are sampled with probabilities (0.5, 0.25, 0.25), respectively.Source: §3.3 · pp. 4–5

4 WHAT DOES MEANFLOWADVANTAGE LEARN?

The analysis characterizes the reward-reweighted MeanFlow target, the direction and scale of MFA updates, and when improvements transfer to the sampler. At its fixed point, the learned map can sample the reweighted distribution across sampling budgets, subject to stated regularity conditions.

Theorem 1

  • The loss has a unique free-function minimizer blending the q-based MeanFlow target with the reference map; at λ = 0, it is ordinary MeanFlow training under q.The shared correction lets rollout and reference anchors act directly on the deployed average-velocity maps.Source: §4 · p. 5

Proposition 1

  • MFA updates toward high-advantage MeanFlow targets and away from low-advantage targets, with step size scaling as 1/γ and vanishing covariance under constant advantages.When all images consistent with a noisy input are equally good, the update does not move the map.Source: §4 · p. 6

Proposition 2

  • MFA uses gain 1/(γ + Eπold[A | Z]) and targets one fixed q-distribution at every input, whereas NFT’s target has that form only when the conditional expected advantage is constant.NFT can fall short of or overshoot its positive-policy target depending on the conditional reward; the analysis does not claim MFA is uniformly better.Source: §4 · p. 6

Theorem 2

  • At the fixed point, under Lemma 1’s regularity conditions, u* is the exact average velocity of vq, so every step count and grid samples q exactly.The target distribution q does not depend on the sampling budget.Source: §4 · p. 6

Scope of the analysis

  • The theory assumes exact correction, calibrated rollout, and free functions, while training uses finite differences, LoRA, and sample-dependent scaling; multiple rewards yield a result only for their combined score.The authors test several predictions empirically, including the need for positive-length intervals and consistency across sampling budgets.Source: §4 · pp. 5–6

5 EXPERIMENTS

The experiments evaluate MFA on SD3.5-Medium and DNA promoters, including performance across sampling budgets and matched ablations. MFA improves the matched image-model baselines, while promoter experiments show gains in teacher-free RL and reward-graded distillation.

5 EXPERIMENTS

  • The evaluation applies MFA as on-policy RL on SD3.5-Medium and extends it to promoter generation through teacher-free RL and reward-graded distillation.The SD3.5-Medium setup uses four-step CFG-free rollouts and matches the MeanFlowNFT outer loop and hyperparameters.Source: §5 · p. 7, §5.1 · p. 8

5.1 TEXT-TO-IMAGE GENERATION WITH SD3.5-MEDIUM

  • With four NFEs, MFA improves all eight metrics over MeanFlowNFT and matches or exceeds 40-NFE DiffusionNFT on six of eight metrics.On promoters, distillation reaches 0.0467 Sei MSE versus 0.0570 for MeanFlowNFT, an 18% reduction.Source: §5.1 · p. 8, §5.3 · p. 10
  • From four to 32 sampling steps, image scores remain on a plateau, but the semigroup residual grows from 0.40 at two steps to 0.76 at 32, so the learned map is not exact.Two steps recover most quality; four steps offer a reported quality-cost trade-off.Source: §5.1 · p. 9

5.2 TESTING THE ANALYSIS ON THE IMAGE MODEL

  • In matched ablations, boundary-only training collapses after about 670 steps, while positive-only and constant advantages collapse within the first few dozen steps.Removing adaptive scaling raises the final gradient norm from 0.97 to 11.97 and lowers held-out aggregate to 8.70; removing the correction also lowers training reward and held-out aggregate.Source: §5.2 · p. 9, Figure 5 · p. 9

5.3 REWARD-FILTERED ADVANTAGE-WEIGHTED PROMOTER FINE-TUNING

  • On 1024-bp FANTOM5 promoters, teacher-free MFA lowers pretrained Sei MSE from 0.0714 to 0.0482 over three seeds, while its optimized reward is statistically indistinguishable from MeanFlowNFT.The experiment uses a Riemannian MeanFlow generator and negative Sei profile MSE as reward; distillation uses a frozen 10-step teacher and reward-filtered pairs.Source: §5.3 · p. 10

6 CONCLUSION

The conclusion presents MeanFlowAdvantage as an objective aligned with MeanFlow’s deployed average-velocity representation. The accompanying proofs characterize its reweighted target, update direction, and boundary-case relation to existing prediction-space training.

6 CONCLUSION

  • MFA improves all eight reported SD3.5-Medium metrics over matched MeanFlowNFT and matches or exceeds 40-step DiffusionNFT on six of eight with four NFEs.The objective also supports teacher-free RL and reward-graded distillation for DNA promoter generation.Source: §6 · p. 10

A.1 AVERAGE-VELOCITY FORM OF THE LOSS AND BOUNDARY REDUCTION

  • At the boundary s = t, the correction vanishes and the loss reduces to the AdvantageFlow loss applied to the boundary predictor for the same samples, advantages, and coefficients.The same boundary reduction holds for the scaled loss when the comparison uses the same denominator.Source: §A.1 · p. 13

A.2 PROOF OF THEOREM 1

  • With a shared correction, the rollout anchor becomes a MeanFlow regression; separate learner and rollout derivatives introduce dependence on the sampled direction and break that equivalence.At the population level, the objective fits the MeanFlow target of a reward-reweighted distribution, with positive reference weight blending that target with the reference map.Source: §A.2 · p. 13

A.3 PROOF OF PROPOSITION 1

  • The MFA population update direction is governed by the conditional covariance between advantage and MeanFlow target, bounded using their conditional standard deviations.The positive-definite quadratic has a unique minimizer under the stated assumptions.Source: §A.2 · p. 13, §A.3 · p. 14

A.4 PROOF OF PROPOSITION 2 AND THE RELATION TO PREDICTION-SPACE NFT

This appendix section compares MeanFlowAdvantage with MeanFlowNFT through their induced population updates and prediction-space objectives. It identifies a positive-only case in which the two methods generally have different minimizers.

A.4 PROOF OF PROPOSITION 2 AND THE RELATION TO PREDICTION-SPACE NFT

  • The positive-policy example shows that the MFA minimizer equals the positive-policy target everywhere, whereas the NFT minimizer equals it only where 2ρ(Z) = β.Here ρ is the conditional mean of the positive-policy weight given Z.Source: §A.4 · p. 14
  • With shared correction, MeanFlowNFT’s implicit positive and negative velocity loss expands to a quadratic with advantage A = 2r − 1 and an anchor weight γ(A) = β − A.At β = 1 it matches the MFA variant γ(A) = 1 − A with λ = 0; separate adaptive denominators break this equivalence.Source: §A.4 · p. 14

A.5 FLOW-MAP CONSISTENCY AND PROOF OF THEOREM 2

The appendix relates a self-consistent MeanFlow map to exact sampling across step grids and separately bounds variance for the unscaled rollout-anchor gradient. It also defines a semigroup residual as a measure of distance from the fixed point, not reward improvement.

A.5 FLOW-MAP CONSISTENCY AND PROOF OF THEOREM 2

  • A map satisfying the fixed-point condition composes into the exact flow from t = 1 to t = 0 on any grid, so the sampler draws from the reweighted target q for every N.The consistency lemma identifies the interval map with the trajectory displacement divided by the interval length.Source: §A.5 · p. 14

A.6 VARIANCE REDUCTION OF THE SHARED ANCHOR

  • Under the stated conditional-mean condition, the rollout-anchor regression gradient is a conditional expectation of the target-regression gradient, so its covariance is no larger in positive-semidefinite order.This variance result concerns the unscaled rollout anchor only, not the full signed-advantage update with adaptive scaling.Source: §A.6 · p. 15

A.7 SEMIGROUP RESIDUAL

  • The semigroup residual measures distance from the fixed point by comparing the direct map with uniform N-step compositions; a small residual alone does not establish reward improvement.The residual is zero by definition at N = 1.Source: §A.7 · p. 15

B IMPLEMENTATION NOTES

The image-model update samples and scores prompt-conditioned images, forms clipped signed advantages, and trains on prediction-space targets with a shared correction. Scheduler scaling, boundary-safe finite differences, loss scaling, and matched controls specify important implementation conditions.

B IMPLEMENTATION NOTES

  • Each update samples prompt-conditioned images, computes clipped signed advantages, draws training intervals, forms corrected predictions for learner, rollout, and reference, then updates the learner and refreshes the rollout model.When s = t, the correction is zero; otherwise it is computed once from the rollout model.Source: B IMPLEMENTATION NOTES · p. 15
  • The finite-difference increment of 5 is in raw scheduler units and corresponds to δ = 0.005; near boundaries, shorten the stencil or use a valid one-sided difference.The stencil assumes s ≤ t − δ and t + δ ≤ 1; on s = t, it is skipped.Source: B IMPLEMENTATION NOTES · p. 15
  • The shared float64 denominator preserves the three loss-term weights within a sample but changes weights across samples and does not generally preserve unweighted population regression.For constant γ, A ≥ −γ − λ + ε_A ensures positive curvature; γ = 1.1, λ = 10^-3, and A ∈ [−1, 1] satisfy this bound.Source: B IMPLEMENTATION NOTES · p. 16
  • Matched controls compare direct-output and diffusion-only variants, and prediction-space and velocity-space losses require reporting their effective time weights because they differ by t^2.The completed matched endpoints are reported in Appendix Table S5.Source: B IMPLEMENTATION NOTES · p. 16

C EXPERIMENTAL DETAILS

The image experiments follow the published Stage-3 protocol, while the DNA setup reports shared configurations and distinct endpoint conventions. The evaluation protocol fixes rollout and sampling details, and the DNA RL objectives are statistically indistinguishable on the optimized reward.

C EXPERIMENTAL DETAILS

  • Both image methods use K = 24 samples per prompt, with EMA schedule, evaluation resolution, prompt lists, and checkpoint identifiers fixed to the published Stage-3 protocol.Table S1 collects SD3.5-Medium RL configurations; sampling and reward choices follow the baseline, while objectives and regularization coefficients differ.Source: C EXPERIMENTAL DETAILS · p. 16, Table S1 · p. 16
  • All three paired Sei MSE intervals contain zero, making the two objectives statistically indistinguishable on the optimized reward, as Proposition 2 predicts.Table S3 reports per-seed results on 512 test sequences, with paired NFT−MFA test Sei MSE intervals.Source: C EXPERIMENTAL DETAILS · p. 17, Table S3 · p. 17
  • The DNA adaptation applies both losses to continuous RMF x1 endpoint coordinates, with noise at t = 0 and data at t = 1, rather than using the image-time prediction form.The learner calls the RMF endpoint predictor at (t, s) = (0, 1); the data target is the guided endpoint, optionally extrapolated from the frozen rollout endpoint.Source: C EXPERIMENTAL DETAILS · p. 18
  • The rollout model is used to construct training pairs and anchor predictions, while evaluation samples from the updated learner; ρ = 1 freezes rollout parameters but leaves the student trainable.Figure S3b–c evaluates both methods every 10 steps through step 100.Source: Figure S3b · p. 18

D.1 MATCHED IMAGE-MODEL ABLATIONS

Matched image ablations probe derivative sharing, anchor weighting, and the limits of the curvature condition. DNA experiments test teacher-free RL on a manifold generator and reward-graded distillation, where MFA outperforms NFT on identical pairs and advantages.

D.1 MATCHED IMAGE-MODEL ABLATIONS

  • Sharing the derivative correction reaches essentially the same reward regime as separate derivatives while reducing wall-clock time per step from 89.9 to 73.3 seconds.The slightly higher single-seed held-out aggregate with separate derivatives is too small to establish a quality advantage.Source: §D.1 · p. 18
  • Changing fixed γ from 1.1 to 5 leaves reward and held-out aggregate nearly unchanged; dynamic γ(A) = 1 − A is stable but not established as better.The authors retain γ = 1.1 as the conservative default and treat the dynamic variant as a follow-up.Source: §D.1 · pp. 18–19
  • Positive local curvature prevents a concave per-sample quadratic but does not guarantee neural-network optimization stability or cross-interval flow consistency.The s = t diffusion-only run retains a positive minimum coefficient while its data, reference, and gradient terms diverge late in training.Source: §D.1 · p. 19

E DNA PROMOTERS: ON-POLICY RL AND REWARD-GRADED DISTILLATION

  • The DNA generator is a Riemannian MeanFlow model pretrained on 1024-bp FANTOM5 promoters; evaluation uses one-step generation with hard one-hot decoding.The reward is negative Sei profile MSE, and the pretrained generator reaches 0.0714 Sei MSE and 0.928 6-mer correlation.Source: E DNA PROMOTERS: ON-POLICY RL AND… · p. 20
  • 0.0714 to 0.0482 mean Sei MSE after teacher-free MFA RL on the manifold generator; MFA and NFT remain statistically indistinguishable on training reward.The six-mer correlation trend favors MFA across three seeds but is not a tested claim.Source: E DNA PROMOTERS: ON-POLICY RL AND… · p. 20
  • Reward-graded distillation uses a frozen one-step rollout and ten-step Sei-guided teacher; their reward gap determines pair filtering and advantage strength.Kept pairs use an extrapolated teacher target, with the rollout endpoint serving as both anchors.Source: E DNA PROMOTERS: ON-POLICY RL AND… · p. 20
  • 0.0467 Sei MSE for MFA, 18% lower than NFT on identical distillation pairs and advantages; the gap persists across training randomness.The paired difference is 0.0103 with a 95% interval of [0.0064, 0.0143], and the comparison was repeated with four training seeds.Source: E DNA PROMOTERS: ON-POLICY RL AND… · p. 21

E.1 ADDITIONAL PROMOTER RESULTS

Across four training seeds, reward-graded MFA outperformed NFT on every seed, with a paired Sei MSE difference whose interval excludes zero. Additional tests compare matched A ≡1 configurations and examine unfiltered negative directions and diagnostics.

E.1 ADDITIONAL PROMOTER RESULTS

  • 0.0115 was the paired NFT−MFA Sei MSE difference across four seeds, and MFA won on every seed.The 95% interval was [0.0091, 0.0138]; average reduction was 19.7%, while the smallest seed-specific reduction was 18.1%.Source: Table S7 · p. 22, §E.1 · p. 22
  • Matched-learning-rate and nominal-extrapolation A ≡1 configurations had comparable results.Table S8 reports checkpoints selected on validation data (n=256) and evaluated on test data (n=512).Source: Table S8 · p. 22, §E.1 · p. 22
  • 0.0744–0.0755 was the best valid Sei MSE for unfiltered signed variants, versus 0.0510 for filtered graded-advantage MFA.The authors retain filtering and leave distillation-specific negative directions for future work.Source: §E.1 · p. 22
  • Figure S4 covers A ≡1 pretraining stability, paired validation differences, and evaluation-size sensitivity; Figure S5 plots training-pair Sei MSE.Figure S5 distinguishes 1-NFE student scores from 10-NFE teacher scores; its accept-filter scores differ from Table S6's hard-decoded held-out sequences.Source: §E.1 · p. 23, Figure S4 · p. 24, Figure S5 · p. 24

Related papers

Loading 2609.37670v1…