MeanFlowAdvantage improves few-step velocity-based generation with stable reward fine-tuning
MeanFlowAdvantage: Stable Reward Fine-Tuning for Few-Step Average-Velocity Generators
Artificial IntelligenceComputational Engineering, Finance, and ScienceComputer Vision and Pattern Recognition
Summary
Generating data step-by-step can be tricky when the way you measure success doesn’t match how the model learns. The authors introduce a method called MeanFlowAdvantage that fixes this mismatch by adjusting how rewards are used during training, keeping everything consistent with how the model makes predictions. This approach helps improve performance on tasks like data generation and DNA sequence design, doing better or as well as methods that take many more steps. Their method works well even when only a few steps are used, saving time and computation.
What this means in practice
- •For machine learning engineers: Improve efficiency and accuracy of sequence-generating models by using reward functions aligned with average-velocity representations during fine-tuning.
- •For bioinformatics developers: Design DNA promoters more effectively using a reward-guided fine-tuning approach that supports both reinforcement learning and teacher-guided distillation.
Authors
Haocheng Tang, Tianchi Xie, Xingqiao Lin
Abstract
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.