STEP: Post-Training Few-Step Flow Models by Propagating Reward Gradients Through Every Generation Step

We study reward fine-tuning of few-step Gaussian-mixture (GM) flow policies. Each network evaluation defines a velocity field over a sampling interval, allowing analytic estimates of the clean latent at intermediate states. We introduce STEP, which uses these estimates as probes for constructing surrogate reward gradients. One probe is sampled within each deployment interval, and a straight-through chain retains dependencies between intervals. Weighted fusion connects all probes to a shared terminal reward while preserving the sampled image in the forward pass. The probes require no additional transformer evaluations. The full training objective combines this terminal loss with auxiliary rewards and frozen-teacher velocity regularization. After $500$ optimization steps on FLUX.1-dev, STEP achieves $0.530$ GenEval, $0.580$ OCR, and $15.80$ held-out HPSv3 with four neural function evaluations. The official distilled student scores $0.396$, $0.577$, and $13.47$, respectively. Within the full objective, removing the chain or replacing interior probes with zero-depth estimates reduces OCR and GenEval. These results support using the interval structure of GM policies for reward fine-tuning at a fixed four-step deployment budget.