ReNFT: Repairing Mode Collapse in Reward Post-Training via Internal Probability-Mass Recalibration

Let $y$ denote a text prompt, $x_{0}$ a clean latent, and $\varepsilon\sim\mathcal{N}(0,I)$ the initial noise. Under rectified-flow interpolation,

the target transport velocity is $v=\varepsilon-x_{0}$, and the model predicts $v_{\theta}(x_{t},t,y)$ trained with the standard flow-matching regression objective. We denote three routes through this sampler: $b$ (frozen base), $\theta$ (conditional policy), and $u$ (unconditional route of the current policy); an EMA-smoothed copy of the policy provides a stable reference, with its velocity written $v_{\mathrm{old}}$.

Reward-based post-training reshapes the distribution of a pretrained generator using a reward model $R(x_{0},y)$. DiffusionNFT (Zheng et al. 2026a) avoids likelihood-ratio policy gradients by applying reward supervision through forward-process regression. Each generated endpoint $x_{0}$ (from a rollout of $N_{i}$ ODE steps) is re-noised at $N_{t}$ training timesteps to $x_{t}=(1-\sigma_{t})x_{0}+\sigma_{t}\xi$, with fresh $\xi\sim\mathcal{N}(0,I)$ and noise level $\sigma_{t}{=}t$, and the model regresses the target velocity $v=\xi-x_{0}$. High-reward endpoints are reconstructed through the ordinary prediction

while low-reward endpoints are pushed away through a mirrored prediction around the EMA reference. With the official default $\beta=1$ in $(1+\beta)v_{\mathrm{old}}-\beta v_{\theta}$, this reduces to

The combined objective over a rollout group is

This MSE formulation trains quickly without likelihood ratios or SDE rollouts, and the EMA reference in $v_{\theta}^{-}$ anchors the update and stabilizes post-training (ablated in the appendix). Equivalently, the loss can be written in endpoint space as $\ell(\hat{x}_{0}(v),x_{0})$ with $\hat{x}_{0}(v)=x_{t}-\sigma_{t}v$; our implementation uses this form with a stop-gradient factor for loss-order adjustment. In our repair setting, the pull and push branches will operate on different endpoints rather than the same $x_{0}$; only at the pure-noise step $\sigma_{t}=1$ do the two trajectories share an identical state.

Online post-training supervises the model with its own sampled endpoints; when reward concentrates in a narrow mode, this acts as self-distillation that sharpens it. ReNFT retains the NFT parameterization but replaces this one-sided supervision with a matched local comparison (Section 4).

We study a repair setting in which an adapter $\theta_{H}$ already achieves high reward but maps different initial noises under the same prompt to a narrow set of outputs. The frozen base generator and $\theta_{H}$ share the same pretrained support; reward adaptation changes which parts of that support are easy to reach, but suppressed modes may remain accessible through routes closer to the base model. Collapse is therefore probability-mass compression rather than deletion, and repair should be possible from within.

A coarse mode-level view makes the feedback loop explicit. If $m$ indexes a visual mode, reward-only adaptation can be heuristically written as

where $\bar{R}_{k}(m,y)$ is the average reward of sampled outputs in mode $m$ and $\eta>0$ scales the reward-induced reweighting. Equation (5) is an interpretation of this feedback, not the exact update law. Figure 1 illustrates the view: per-prompt modes contract into narrow spikes as the loop reinforces the favored mode, then reopen during repair. Conditional modes that had lost probability regain output mass, and the unconditional distribution range expands accordingly.

This feedback also explains why re-optimizing reward cannot repair collapse: among semantically equivalent samples, reward differences encode preference rather than quality, so the reweighting only sharpens that preference. Repair instead needs reward gaps that reflect genuine quality differences. ReNFT achieves this by constructing matched counterfactuals from different internal routes (Sections 4.2–4.4), making the reward gap between paired endpoints informative rather than a restatement of the model bias. We validate this reversibility in Section 5.

The unconditional route is the empty-prompt velocity $u_{\theta}(x,t)=v_{\theta}(x,t,\varnothing)$. Although it shares $\theta$’s updated parameters, it has no separate reward target; its changes therefore expose prompt-independent side effects of conditional post-training: recurring structures, styles, or textures the adapter injects even when not requested, rather than directed optimization. Unconditional samples from a fixed noise grid contract from a diverse base into a narrow hub after post-training and progressively reopen during repair (see the appendix for unconditional samples across training steps). We use these tendencies as an internal bias probe, whose diagnostic validity is evaluated rather than assumed.

We first use this readout to prioritize prompts on which the bias is easiest to expose, as shown in Figure 3a. Offline, the prompt bank is encoded once by the SigLIP2 text encoder (Tschannen et al. 2025) and cached. Online, $M$ empty-prompt probes $\{u_{j}\}_{j=1}^{M}$ are generated from the current policy and encoded by the SigLIP2 image encoder. Each candidate prompt $y$ is scored by

and the bottom-$K$ prompts (lowest mean similarity) are selected as anti-hub prompts, those farthest from the current unconditional tendency, where the exposed bias is most likely to yield distinguishable proposals. The frozen encoder serves only prompt prioritization; it never enters the reward, advantage, or repair loss. The same route also serves constructively in the mixed sampling of Section 4.3, where it exposes the bias inside conditional trajectories for reward ranking.

The frozen base route and the unconditional route play complementary roles. A base step introduces a base-like vector-field direction and can propose alternatives that are sampled less often by the post-trained policy; an unconditional step exposes the prompt-independent bias identified in Section 4.2. The conditional route $\theta$ occupies most steps by design, limiting deviation from the high-reward policy. Crucially, neither route is hard-coded as a pull or push branch: mixed routing generates nearby alternatives, and the reward decides their roles (Figure 3b).

For a routing pattern $P$, let $T_{P,y}$ map initial noise to a terminal sample. The final configuration repeats two three-step blocks across the rollout,

Both routes diverge from the same initial noise $\varepsilon$, keeping two thirds of each route conditional and avoiding long $bb$ or $uu$ blocks that stray too far from the policy (cf. stage-wise analyses (Jin, Shi, and Gu 2026)).

For a selected prompt $y$ and shared initial noise $\varepsilon$, the matched endpoints are $x_{0}^{A}=T_{P_{A},y}(\varepsilon)$ and $x_{0}^{B}=T_{P_{B},y}(\varepsilon)$. Sharing $(y,\varepsilon)$ makes routing the main changed factor: two pure-$\theta$ trajectories would coincide, while a direct base-versus-policy comparison is too one-sided; the two mixed routes instead perturb the policy in different directions, and either may win for a particular $(y,\varepsilon)$. Trajectory-level branching has likewise been exploited for alignment and credit assignment in TMPO and DenseGRPO (Li et al. 2026a; Deng et al. 2026); we instead branch two route mixtures from identical noise to construct a reward-testable counterfactual pair.

We rank the endpoints with the original reward (Figure 3b, right) and obtain $(x_{0}^{H},x_{0}^{L})$; when rewards differ, the higher-reward endpoint is always the pull target and the lower the push target, regardless of route. The matched reward gap

provides a local preference without claiming global superiority of either route. For continuous rewards (PickScore), ties are rare; an adaptive flipping guard (Figure 3c) requires at least half of the pull targets to come from route $A$, flipping the lowest-margin route-$B$ win when the ratio falls below $0.5$. This guard keeps the coverage-restoring $b\theta\theta$ branch dominant and prevents the hacking-prone $\theta\theta u$ branch from re-hacking the pull distribution. For rule-based discrete rewards (GenEval), ties are frequent; we break them with a near 50/50 split that slightly favors route $A$, without activating the flipping guard.

We realize the matched preference with the NFT mirror mechanism from Section 3.2. The loss in v-space is

The mirror reference $v_{\theta}^{-}$ uses an EMA-smoothed copy of the policy with constant decay, since repair starts from a late-stage checkpoint.

For each paired timestep, we sample a fresh perturbation $\xi$ and construct

with targets $v^{H}=\xi-x_{0}^{H}$ and $v^{L}=\xi-x_{0}^{L}$. The pull and push share $\xi$, treating corruption as a matched nuisance; sampling $\xi$ independently of $\varepsilon$ avoids a per-(prompt, noise) fitting shortcut (Wen et al. 2024).

At $\sigma_{t}=1$, a fresh $\xi$ would break the coupling, so we use the native rollout noise $\varepsilon$ for a joint update (Figure 3c), with $x_{t}^{H}=x_{t}^{L}=\varepsilon$ and targets $v^{H}=\varepsilon-x_{0}^{H}$, $v^{L}=\varepsilon-x_{0}^{L}$. This anchor matters because early denoising decisions strongly influence global structure (Jin, Shi, and Gu 2026). Each repair step applies one joint update with $\varepsilon$ at $\sigma_{t}=1$, followed by an even number of paired updates with fresh $\xi$ at $\sigma_{t}<1$ (see the appendix for ablations on the joint-noise choice and rollout length). Mixed routes only construct training supervision; inference uses the standard conditional forward pass with no routing overhead.

We use SD3.5-M (Esser et al. 2024) with LoRA adapters and train separate PickScore (Kirstain et al. 2023) and GenEval (Ghosh, Hajishirzi, and Schmidt 2023) checkpoints. ReNFT reloads the DiffusionNFT LoRA at step 700 (PickScore) or 250 (GenEval) and repairs for 50 steps with the default configuration in the appendix. Model-based rewards use all 1,024 PickScore test prompts; the rule-based reward uses all 553 GenEval prompts; diversity uses 100 fixed prompts with 12 fixed-noise samples each. NFT, E²PO, and ReNFT use CFG=1; base and GRPO-style checkpoints use CFG=4.5, so cross-CFG comparisons are end-to-end method comparisons, not controlled CFG ablations.

We compare the untrained SD3.5-M reference, Flow-GRPO (Liu et al. 2025), DiffusionNFT (NFT) (Zheng et al. 2026a), DiverseGRPO (Liu et al. 2026a), and E²PO (Hu et al. 2026). E²PO is not open-sourced; we match its official training budget and reuse its paper-reported reward values (marked ^†), omitting it from the diversity ranking. Quality metrics are PickScore-v1, Aesthetic (Radford et al. 2021; Schuhmann et al. 2022), ImageReward-v1.0 (Xu et al. 2023), and HPSv2.1 (Wu et al. 2023). Diversity metrics are LPIPS-Div (Zhang et al. 2018), DreamSim-Div (Fu et al. 2023), and DINOv3-Div (Siméoni et al. 2025), each computed as the mean pairwise distance among 12 samples per prompt. Anti-hub uses 24 probes and 1,000 candidates; further details are in the appendix.

Figures 4 and 5 compare Base, NFT, DiverseGRPO, and ReNFT on GenEval+OCR and PickScore prompts, respectively, with a shared noise set.

The ablations examine two design decisions in the main text: anti-hub prompt prioritization and internal route construction. All variants branch from the same checkpoint $H$ under the same 50-step budget and checkpoint-evaluation protocol.

Appendix ablations support the default rollout length, joint-noise choice, and EMA decay; sensitivity to the anti-hub pool size and flipping threshold remains future work.