PaperlayerView exact revision on arXiv

Source-linked AI summary

MEND: RL For Flow Models via Proximal Velocity Matching

Shreshth Saini, Neil Birkbeck, Yilin Wang, Balu Adsumilli, Alan C. Bovik

arXiv:2610.05954v1cs.LGcs.CV

TL;DR

Reward post-training must turn scalar preference scores into selective updates, rather than only reweighting samples or moving every sample by reward ascent. MEND uses capped, per-sample proximal velocity matching to decide which samples move. It outperforms Flow-GRPO in 100 updates on five of six evaluators at the same base distance and surpasses ReFL and DiffusionNFT at equal budget across four training rewards.

  • Problem

    Existing post-training either reweights samples under a KL penalty or reference, or backpropagates reward through sampling without checking whether each sample's gain merits its displacement.

    Source: §1 · p. 2

  • Method

    MEND caps rewards within prompt groups, accepts gradient-based sample moves only when capped reward gain exceeds a quadratic displacement price, and regresses onto the resulting velocity targets.

    Source: §1 · p. 2

  • Results

    In 100 updates, MEND outperforms Flow-GRPO on five of six evaluators at the same base distance and surpasses ReFL and DiffusionNFT on all four training rewards at equal budget, reaching PickScore 24.03 versus 23.92 and 23.43.

    Source: §1 · p. 3, §4.3 · p. 7

  • Takeaways & Limitations

    MEND applies across flow backbones, improving SD3.5-M, SD3-M, and distilled Z-Image-Turbo at the same update budget.

    Source: §1 · p. 3, §6 · p. 10

    MEND requires a differentiable reward; longer training can over-optimize a reward, and some single-reward runs lose held-out score within 100 updates.

    Source: §6 · p. 10

Abstract

from arXiv · show

Reward post-training of flow models either reweights the model's own samples under a KL penalty or a frozen reference, often for thousands of updates, or backpropagates the reward and moves every sample without checking that the move is worth its size. We introduce MEND, a reinforcement learning method built on proximal velocity matching. MEND caps rewards within each prompt group, so samples that already score well receive no move. Below the cap, it proposes moves along the reward gradient and accepts one only when its capped reward gain exceeds a quadratic displacement price. The model then regresses onto the resulting velocity targets, with no KL term, frozen reference model, or advantage weights. In 100 updates, MEND outperforms Flow-GRPO (about 4k updates) on five of six evaluators at the same distance to base-model images. Under an equal-budget protocol, it surpasses ReFL and DiffusionNFT at every evaluated update across four training rewards, reaching PickScore 24.03 versus 23.92 and 23.43, respectively. A 300-update three-reward run also surpasses the five-reward DiffusionNFT model on all three rewards it trains on. MEND is general and easy to adopt: it applies to any flow backbone with a differentiable reward.

1 INTRODUCTION

Reward post-training must turn scalar scores into selective updates: reweighting changes sample weights, while reward ascent moves samples without checking whether the gain merits the displacement. MEND uses a per-sample reward-versus-distance check and outperforms longer-trained baselines in the reported evaluations.

Motivation

  • Reweighting changes how much existing samples count, while reward backpropagation moves every sample without checking whether the reward gain is worth the move.A KL penalty or reference term limits distribution-level drift but makes no per-sample decision.Source: §1 · p. 2

Method

  • MEND accepts a sample move only when capped reward gain exceeds a quadratic displacement price, then trains by regression onto the resulting velocity targets.It uses no KL term, frozen reference model, or advantage weights.Source: §1 · p. 2, Figure 4 · p. 4

Theory

  • MEND's guarantees state that accepted targets improve capped reward by more than their quadratic price, with displacement shrinking toward zero near the cap.At behavior parameters, the regression descent direction backpropagates normalized reward gradients through accepted samples only, scaled by selected steps.Source: §1 · p. 2

Results

  • In 100 updates, MEND outperforms Flow-GRPO on five of six evaluators at the same base distance and surpasses ReFL and DiffusionNFT on all four training rewards at equal budget.The matched visual comparisons share prompts and seeds across methods.Source: §1 · p. 3, Figure 3 · p. 3

Generality

  • MEND is backbone-agnostic and improves SD3.5-M, SD3-M, and distilled Z-Image-Turbo at the same update budget.The reported 100-update comparison is at the same distance to base-model images.Source: §1 · p. 3

2 METHOD

MEND builds per-sample targets through a cap, reward-gradient proposals, and a reward-versus-distance acceptance check, then fits those targets with velocity-matching regression. Its behavior adapter follows training through an exponential moving average rather than serving as a frozen reference.

Round overview

  • A round rolls out the behavior adapter, selects targets, and takes one regression update; kept samples receive zero displacement and moved samples receive their accepted correction.The process is organized as rollout, cap/propose/verify target selection, and regression.Source: §2 · pp. 4–5

Setting and notation

  • A rectified flow interpolates a clean latent and Gaussian noise, while its velocity provides a clean prediction for a prompt.The behavior adapter generates deterministic 10-step trajectories and stores a state nearest t = 0.278 for each sample.Source: §2 · p. 4

Cap

  • Within each prompt group, rewards are capped at the 0.75 quantile subject to a slowly rising global floor, and samples at or above the cap receive no displacement.Rewards above the cap receive no additional credit.Source: §2 · p. 4

Propose

  • Below the cap, samples with nonzero reward gradients receive three decoded proposals along the normalized gradient, making each step length independent of reward scale.Samples with zero gradient are kept.Source: §2 · pp. 4–5

Verify

  • The selection rule maximizes capped reward minus a quadratic displacement penalty over a finite candidate set, including the unchanged sample.A move is accepted only when its capped gain exceeds its distance price; τ sets that price and is adjusted between rounds to target 30%–60% acceptance.Source: §2 · p. 5

Regress

  • For moved samples, regression targets the behavior prediction at the stored state plus accepted displacement, rather than the endpoint itself.The endpoint differs from the one-step prediction by the remainder of the rollout, which would add a reward-independent gap to the update.Source: §2 · p. 5
  • Equation 4 is a velocity-matching loss: kept samples target the behavior velocity, while moved samples target that velocity shifted by the accepted correction.The target is equivalently the behavior clean prediction plus displacement.Source: §2 · p. 5

Behavior model and optimizer

  • The behavior adapter follows the trained adapter by EMA, so the keep term anchors velocities to a lagged model rather than a frozen reference.Each round uses one AdamW step; MEND does not backpropagate through the sampler.Source: §2 · p. 5

3 THEORETICAL ANALYSIS

MEND’s target-selection rule accepts a move only when its objective exceeds that of leaving the sample unchanged. The resulting move is bounded by reward gain and shrinks as a sample approaches the reward cap.

3 THEORETICAL ANALYSIS

  • Accepted targets are bounded by the reward they gain, and the regression update realizes reward backpropagation only on accepted seeds.Kept seeds contribute zero gradient at the old parameters.Source: §3 · pp. 5–6
  • The largest permitted move shrinks to zero as a sample’s headroom below the cap vanishes.The bound tightens as κ − Rκ(x) approaches zero.Source: §3 · p. 5
  • Equation 7 backpropagates normalized endpoint reward gradients through clean predictions, restricted to accepted seeds and scaled by their selected steps.It preserves the reward-gradient direction while selecting which samples move and how far.Source: §3 · p. 6

4 EXPERIMENTS

Across the reported comparisons, MEND improves reward against several baselines while matching or staying close to base-model distance. Results span equal-budget training rewards, backbones, ablations, and matched visual examples.

4.1 SETUP

  • The setup trains SD3.5-M on PickScore for 100 updates, using 48 prompts, 24 images per update, rank-32 LoRA, and 10-step 512-pixel rollouts.Evaluation uses DrawBench with 200 prompts and five seeds, and DreamSim distance to matched base images for fidelity.Source: §4.1 · p. 6

4.2 MAIN COMPARISON

  • Five of six evaluators favor MEND over Flow-GRPO at the same base distance of 0.313, despite roughly 40-fold fewer updates.Aesthetic score is the exception: Flow-GRPO scores 5.90 versus MEND’s 5.88.Source: §4.2 · p. 6

4.3 UPDATE EFFICIENCY AND FIDELITY

  • At equal budget, MEND leads ReFL and DiffusionNFT at every evaluated update; at update 100, PickScore is 24.03 versus 23.92 and 23.43.On 64 unseen Pick-a-Pic prompts, MEND is above same-protocol DiffusionNFT at every evaluated update.Source: §4.3 · p. 7, Figure 8 · p. 8

4.4 OTHER REWARDS AND BACKBONES

  • Across four training rewards, each MEND run exceeds ReFL and DiffusionNFT on its trained reward at every evaluated update and finishes above the base on held-out evaluators.Three runs show held-out declines before update 100.Source: §4.4 · pp. 7–8, Table 3 · p. 8
  • MEND improves each tested backbone at the same update budget, and applies to any flow model with a differentiable reward.On Z-Image-Turbo, PickScore training raises PickScore from 22.86 to 24.10; HPSv2.1 training reaches 0.357 with held-out PickScore 23.27.Source: §4.4 · p. 8

4.5 ABLATIONS

  • The full cap-and-price rule preserves 0.037 more diversity than moving every below-cap seed, for 0.35 PickScore less.Moving every below-cap seed yields the highest full-set PickScore, 23.45, and lowest diversity, 0.260.Source: §4.5 · p. 8

4.6 QUALITATIVE RESULTS

  • Matched-noise examples show MEND correcting base-model prompt errors while keeping composition; examples include counts, attributes, and spatial relations.Many examples have higher saturation, reflecting a reward preference examined in Appendix E.3.Source: §4.2 · p. 6, §4.6 · p. 9

5 RELATED WORK

Prior flow-model post-training methods include policy gradients, weighted regression, reward backpropagation, and sample selection. MEND instead constructs local targets and checks each proposed move against its capped reward gain and displacement cost.

5 RELATED WORK

  • Flow-GRPO and DanceGRPO extend group-relative policy gradients to flow models using stochastic sampling and a KL penalty.DDPO and DPOK also optimize rewards with policy gradients.Source: §5 · p. 9
  • DiffusionNFT, AWM, and FlowAWR use reward- or advantage-weighted regression on model-generated samples, adjusting contributions of existing samples.Diffusion-DPO fits offline preferences, while RAFT and related methods select among existing samples.Source: §5 · p. 9
  • MEND constructs new local targets and accepts a proposal only when capped reward gain exceeds a quadratic displacement penalty.This per-target check differs from distribution-level Wasserstein regularization.Source: §5 · p. 9

6 CONCLUSION

MEND uses proximal velocity matching to train flow models on per-sample targets, and the conclusion reports competitive results across evaluators, rewards, and backbones. The section also describes implementation bookkeeping, evaluation settings, and baseline comparisons.

6 CONCLUSION

  • In 100 updates, MEND outperforms Flow-GRPO on every evaluator except aesthetic score at the same distance to base-model images, and surpasses ReFL and DiffusionNFT at equal budget on four rewards.The conclusion also reports improvements on SD3.5-M, SD3-M, and distilled Z-Image-Turbo at the same update budget.Source: §6 · p. 10

A ALGORITHM AS IMPLEMENTED

  • Algorithm 2 adds implementation bookkeeping to Algorithm 1: a global cap floor, tie rules, a price controller, and two EMAs.Rewards, cap statistics, and loss normalizers are pooled across devices.Source: A ALGORITHM AS IMPLEMENTED · p. 13

B EVALUATION SETTINGS

  • Evaluation settings vary by figure and table, including guidance and update budget; Table 2 includes runs at update 100 and a three-reward row at update 300.The settings also specify different guidance levels for some baselines and MEND comparisons.Source: B EVALUATION SETTINGS · pp. 13–14

B.1 BASELINES

  • Flow-GRPO is a PickScore adapter trained for about 4k updates with a KL penalty; matched comparisons use the † rows because other baseline runs do not match Protocol O.The pipeline's reproduced scores are reported as agreeing with the original papers, including Flow-GRPO at 23.52 against 23.53 and DiffusionNFT at 23.82 against 23.80.Source: §B.1 · p. 14

B.2 EVALUATION PROTOCOL

The evaluation uses shared image-generation settings and reports how guidance affects reward and diversity. Held-out checks compare MEND with a base model and Flow-GRPO, and track reward against distance to the pretrained model and diversity.

B.2 EVALUATION PROTOCOL

  • DrawBench evaluation uses 200 prompts with five images each, 40-step Euler flow, 512-pixel images, EMA weights for MEND, shared initial noise, and 95% prompt-bootstrap intervals.Diversity is mean pairwise DreamSim cosine distance within prompts; distance to the base uses the base image from the same noise and guidance.Source: §B.2 · p. 14

B.3 GUIDANCE

  • As guidance rises, HPSv2.1, ImageReward, and CLIPScore rise while diversity falls for every model; MEND's PickScore is highest at guidance 1 and declines slowly.Figure 12 marks guidance 4.5 for the Val64 comparison.Source: §B.3 · p. 15, Figure 12 · p. 15

B.4 OTHER HELD-OUT CHECKS

  • At guidance 1, MEND scores 0.62 on GenEval versus 0.32 for its base; at guidance 4.5, it matches Flow-GRPO on GenEval (0.74 versus 0.74) but trails on OCR (0.53 versus 0.68).Neither GenEval nor OCR was used to train MEND.Source: §B.4 · p. 15
  • The guidance-sweep image grid compares SD3.5-M, Flow-GRPO, DiffusionNFT, and MEND across guidance 1, 2, 3, and 4.5; DiffusionNFT over-saturates as guidance grows.The images show two Val64 prompts at update 100, selected by eye.Source: Figure 13 · p. 16
  • At every update, MEND is above DiffusionNFT in PickScore and closer to pretrained images; at equal diversity, it is at least 0.35 higher in PickScore.These comparisons use the same Protocol O setting.Source: Figure 14 · p. 16, Figure 15 · p. 16

C THEORY: STATEMENTS AND PROOFS

The theory analyzes MEND’s selected targets and the first-order effects of fitting them, with guarantees bounded by explicit assumptions. It also characterizes how target mass can move across region boundaries.

C THEORY: STATEMENTS AND PROOFS

  • The guarantees cover target selection and first-order prediction changes, not finite optimizer steps or full runs; transfer assumptions are not established for AdamW and gradient clipping.The appendix identifies these as limits on what its theoretical results establish.Source: C THEORY: STATEMENTS AND PROOFS · p. 17, §C.1 · p. 18

C.1 VERIFIED TARGET SELECTION

  • Accepted targets strictly improve capped reward by more than their quadratic displacement price, and displacement is bounded by reward headroom and the longest proposal.These guarantees apply to target selection, not to the full training update.Source: §C.1 · p. 18

C.2 FIRST-ORDER FORM OF THE UPDATE

  • At first order, prediction changes combine accepted displacements through a kernel, but this does not guarantee endpoint-reward changes or held-out evaluator outcomes.The result concerns a linearized prediction projection, not endpoint reward or held-out evaluation.Source: §C.2 · pp. 19–20

C.3 TARGET MASS ACROSS REGION BOUNDARIES

  • Target mass can cross a region boundary only from samples within their displacement bounds; samples at the cap have zero displacement radius.The boundary-crossing statement is about targets, while post-training diversity is evaluated empirically.Source: §C.3 · p. 21

D FULL RESULT TABLES

The appendices provide expanded result tables, a held-out-prompt validation split, training-dynamics plots, and per-reward evaluations across backbones and evaluators.

D FULL RESULT TABLES

  • Table 5 reports full SD3.5-M results across evaluators, diversity, distance to base, and 95% intervals.Its caption identifies DrawBench 200 × 5 and marks the best and second-best trained headline rows.Source: D FULL RESULT TABLES · p. 21, Table 5 · p. 21

D.1 VALIDATION SPLIT

  • Table 6 repeats Table 2 on DrawBench prompts 65 to 200, which were not used for design decisions.The caption notes that guidance-4.5 rows reuse Table 2 images; other rows use 136 × 2 images.Source: §D.1 · p. 21, Table 6 · p. 22

E TRAINING DYNAMICS AND OTHER BACKBONES

E.1 OTHER REWARDS AND SD3-M

  • Tables 7 and 8 report one run per training reward and a separate MEND run on SD3-M.Table 7 is at update 100 under Protocol O; the ‡ row marks the last update before a defined held-out decline.Source: §E.1 · p. 22, Table 7 · p. 22, Table 8 · p. 23

E.2 EVERY EVALUATOR PER RUN

  • Table 9 scores each single-reward run with every evaluator, with each row’s training reward shown by its diagonal entry.The HPSv3 run’s last full evaluation is update 90; the table separates guidance-1 and guidance-4.5 blocks.Source: §E.2 · p. 23, Table 9 · p. 23

E.3 TRAINING PAST THE BUDGET AND OVER-OPTIMIZATION

Past the development budget, rising training reward does not ensure continued held-out gains: PickScore gains flatten, and optimizing the LAION aesthetic predictor collapses other evaluators. On Z-Image-Turbo, both runs improve every held-out evaluator except CLIPScore.

E.3 TRAINING PAST THE BUDGET AND OVER-OPTIMIZATION

  • Training on the LAION aesthetic predictor raises aesthetic score from 5.2 to 9.7 while every other evaluator collapses.The paper identifies this reward as known to be over-optimized.Source: §E.3 · p. 23, Table 10 · p. 23

E.4 Z-IMAGE-TURBO

  • Both Z-Image-Turbo runs improve every held-out evaluator except CLIPScore.Z-Image-Turbo is distilled, sampled in 9 steps at 1024 pixels; the evaluation covers every evaluator on DrawBench.Source: §E.4 · p. 23, Table 11 · p. 24

E.3 TRAINING PAST THE BUDGET AND OVER-OPTIMIZATION

  • At 200 PickScore updates, trained reward keeps rising while held-out gains flatten and diversity declines slowly; HPSv3 peaks at update 60.The figure compares changes at update 200 with those at update 100.Source: §E.3 · p. 23, Table 10 · p. 23, Figure 17 · p. 24

F ABLATIONS AND CONTROLS

The controls show reward–diversity trade-offs in design choices, while the toy model tests mode retention and reveals that verification against a misspecified reward only delays off-data moves.

F ABLATIONS AND CONTROLS

  • A looser cap, a broader price band with more acceptances, and a higher learning rate each trade diversity for reward in development-budget ablations.The one-factor arms use the development budget; the table defines “moved” as the share of seeds receiving a move target.Source: F ABLATIONS AND CONTROLS · p. 24, Table 12 · p. 24

F.1 DESIGN CHOICES

  • Nearly every accepted proposal uses the shortest rung, η = 0.1, while the optimizer—not the verdict—determines how far the model moves.Across held-out seeds, the median projection of endpoint motion onto the certified direction is 0.001.Source: §F.1 · pp. 24–25

F.2 SENSITIVITY AT THE FULL BUDGET

  • At update 50, adding HPSv2.1 to the reward raises HPSv2.1 and HPSv3 without PickScore cost; a 0.05 rung lowers PickScore.A second reference launch differs by 0.01 PickScore, while initial price and EMA barely matter.Source: §F.2 · p. 25, Table 13 · p. 25

F.3 ONE-DIMENSIONAL TOY MODEL

  • After one round, MEND mode masses are 0.218, 0.430, and 0.353, versus 0.217, 0.430, and 0.350 before; ascent moves 0.127 of mass off-data.The toy uses a learned reward with an overestimated tail beyond the data.Source: §F.3 · p. 25

F.4 THE TOY OVER REPEATED ROUNDS

  • For 15 rounds, MEND retains every mode and keeps off-data mass at most 0.0003; gradient ascent drains the hard mode and moves 16% of mass off-data.From round 16, the learned reward certifies off-data moves; by round 18 MEND has passed gradient ascent.Source: Figure 19 · p. 26, §F.4 · p. 26

G MORE QUALITATIVE RESULTS

The qualitative results pair models using shared prompts and initial noise, and include both selected examples and an uncurated gallery. Across prompt progressions, compositions mostly remain similar while prompt fidelity mostly improves.

G MORE QUALITATIVE RESULTS

  • Comparisons use shared prompts and initial noise, but sample prompts and seeds are selected by eye; the gallery sheet is shown without quality selection.The tables report numbers over the full prompt sets.Source: G MORE QUALITATIVE RESULTS · p. 26, Figure 21 · p. 28, Figure 25 · p. 31
  • Across 1,200 logged seeds, 33% of proposals were certified, and every certified move used the shortest step.The figure compares capped gain with displacement price for two example seeds and marks the selected targets.Source: Figure 20 · p. 27
  • Across six seeds of two prompts, compositions stay similar while prompt fidelity mostly improves; fixed-noise prompt progressions also show PickScore tiles.The 1024-pixel gallery tiles are shown at native aspect ratios without upscaling.Source: Figure 23 · p. 29, Figure 22 · p. 29, Figure 24 · p. 30

Related papers

Loading 2610.05954v1…