MEND: RL For Flow Models via Proximal Velocity Matching

MEND. RL For Flow Models via Proximal Velocity Matching. By Shreshth Saini and four others, from the University of Texas at Austin, Google, and one other institution. Published on the fifth of October. Its code repository is worked out: about one hundred and sixty source files, with a readme and a licence.

Reward post-training of flow models tends to come in two forms, and the authors find fault with each. One reweights the model's own samples under a KL penalty or against a frozen reference, often for thousands of updates. The other backpropagates the reward directly and moves every sample, without checking whether a given move is worth its size. The first is slow. The second spends displacement indiscriminately, including on samples that did not need to change.

MEND is a reinforcement learning method built on what the authors call proximal velocity matching. Its starting point is that a move should have to earn its place. Rewards are capped within each prompt group, so a sample that already scores well receives no move at all. Only below the cap does MEND consider changing anything.

For those samples, it proposes moves along the reward gradient, and then puts each proposal to a test. A move is accepted only when its capped reward gain exceeds a quadratic price on displacement. Small moves are cheap. Large ones must buy a correspondingly large gain. The accepted moves define velocity targets, and the model simply regresses onto them. There is no KL term, no frozen reference model, and no advantage weights. What is left of the usual apparatus is a regression target and a rule for which moves make it in.

The comparisons are framed around update budgets. In one hundred updates, MEND outperforms Flow-GRPO, which uses about four thousand, on five of six evaluators, measured at the same distance from the base model's images. Under an equal-budget protocol, it surpasses ReFL and DiffusionNFT at every evaluated update, across four training rewards.

On PickScore it reaches twenty-four point zero three. ReFL reaches twenty-three point nine two, and DiffusionNFT reaches twenty-three point four three. A further run of three hundred updates, trained on three rewards, surpasses the DiffusionNFT model trained on five rewards, on all three of the rewards it trains on.

MEND is presented as general and easy to adopt: it applies to any flow backbone that has a differentiable reward. Its proposition is that much of the machinery of reward post-training can be replaced by a single question asked of each sample. Is this move worth what it costs? Samples that answer no are left alone, and the results are reported in hundreds of updates rather than thousands.