Source-linked AI summary
ReNFT: Repairing Mode Collapse in Reward Post-Training via Internal Probability-Mass Recalibration
Yuchen Bao, Chao Wen, Haowei Wang, Ruoxin Chen, Donghao Luo, Jiahui Zhan, Wenjian Huang, Shen Chen, Yiting Wang, Taiping Yao, Chengjie Wang, Shouhong Ding, Jianguo Zhang
TL;DR
Reward post-training can compress probability onto reward-favored modes, erasing within-prompt diversity, and existing interventions do not repair an already-collapsed adapter while preserving reward. ReNFT recalibrates probability mass internally using unconditional probes and matched base/unconditional routes, retaining nearly all NFT reward while substantially improving diversity. The paper’s supported scope includes repair from inherited generator support, with sensitivity to anti-hub pool size and flipping threshold left for future work.
Problem
Reward post-training concentrates probability on a few favored modes, collapsing within-prompt diversity, while existing approaches do not repair an already-collapsed adapter without external objectives or text-encoder modification.
Method
ReNFT uses anti-hub unconditional probes, matched counterfactual routes from the same prompt and noise, reward-ranked pull and push roles, and joint-and-paired NFT updates.
Results
ReNFT retains 98.9% of NFT’s reward on PickScore and 99.0% on GenEval while improving DreamSim-Div by 58.8% and 55.0%, respectively.
Takeaways & Limitations
Collapsed adapters can be repaired from within the generator without external diversity objectives, because post-training suppresses inherited modes rather than deleting them.
Takeaways & Limitations
Sensitivity to the anti-hub pool size and flipping threshold remains future work.
Abstract
from arXiv · showhide
Reward post-training of diffusion generators inevitably concentrates probability mass on a few reward-favored modes, a mode collapse that erases within-prompt diversity. Existing methods for mitigating collapse rely on external signals or interfaces, augmenting the reward with perceptual objectives, adjusting reference regularization, or modifying the text encoder, but none repairs an adapter that has already collapsed while preserving the acquired reward. We observe that online post-training primarily reallocates probability mass over capabilities inherited from pretraining rather than learning new visual content. Collapse is therefore suppression, not deletion, and can be reversed from within the generator. We propose ReNFT, which repairs a high-reward, low-diversity adapter through internal probability-mass recalibration. Unconditional probes first prioritize "anti-hub" prompts where the prompt-independent bias is easiest to expose. Two policy-dominated mixed routes then generate matched counterfactual proposals from the same prompt and initial noise, one probing the frozen base direction for suppressed alternatives and the other exposing the post-trained unconditional tendency. Reward ranking with an adaptive flipping guard assigns pull and push roles, and a joint-and-paired NFT update realizes the repair. On PickScore and GenEval, ReNFT retains 98.9% and 99.0% of NFT's reward while improving DreamSim-Div by 58.8% and 55.0%, respectively, offering a complementary alternative to external interventions.
1 Introduction
Reward post-training can collapse within-prompt diversity by concentrating probability on reward-favored modes, but ReNFT treats this as reversible suppression within the generator. It uses internal bias probes and matched counterfactual routes to recalibrate probability mass while preserving reward.
- Motivation: NFT can map different initial noises under one prompt to similar outputs, causing high-reward adapters to lose structural and stylistic diversity.The collapse is most severe under the fastest optimizer, NFT.
- Motivation: Existing mitigations use external perceptual rewards, base-model regularization, or text-representation changes, leaving open repair of an already-collapsed adapter without sacrificing reward.The stated question is whether internal generator distributions can provide such a repair signal.
- Motivation: Online post-training reallocates inherited probability mass toward reward-favored modes, starving alternatives rather than deleting visual capabilities.This makes collapsed modes potentially recoverable through internal generator routes.
- ReNFT: ReNFT prioritizes anti-hub prompts using the post-trained unconditional route, where prompt-independent reward bias is easiest to expose.The unconditional route is used diagnostically rather than optimized directly.
- ReNFT: Two matched internal routes share the prompt and initial noise: the frozen base probes suppressed alternatives, while the unconditional route exposes learned bias for reward-ranked pull and push updates.The route construction is training-only; repaired sampling uses the standard conditional forward pass.
- Results: ReNFT preserves 98.9–99.0% of NFT’s reward while improving DreamSim-Div by 58.8% and 55.0% on PickScore and GenEval, respectively.The result supports internal probability-mass recalibration as a repair of reward-induced collapse.
2 Related Work
Related work improves reward post-training through optimization, reward shaping, regularization, interfaces, and trajectory routing. ReNFT instead targets an already-collapsed adapter using internal routes of the generator.
- Reward post-training: DiffusionDPO and Flow-GRPO improve alignment by increasing probability assigned to reward-preferred outputs, while DiffusionNFT uses forward-process reconstruction for faster alignment.Later variants refine clipping, credit assignment, self-correction, and distillation.
- Collapse mitigation: Reward-shaping and interface-level methods primarily prevent collapse during training, whereas KL-style anchoring can restore suppressed modes but ties attainable reward to the base model.These limitations motivate a repair method based on internal generator components.
- ReNFT: ReNFT generates candidates from a frozen base route and an unconditional route of the same post-trained generator, rather than relying on external diversity signals or text-encoder modification.The repair signal comes from components already inside the post-trained generator.
3 Preliminaries
The paper formulates flow-matching generation through rectified-flow interpolation and uses NFT-style forward-process regression for reward post-training. High-reward endpoints are reconstructed, while low-reward endpoints receive mirrored updates around an EMA reference.
- Flow-Matching Generation: Rectified-flow interpolation forms a noisy state xt from clean latent x0 and initial noise ε at interpolation time t.The notation uses xt = (1 − t)x0 + tε with ε sampled from a standard normal.
- Flow-Matching Generation: The model predicts transport velocity vθ(xt, t, y), with routes for the frozen base, conditional policy, and current-policy unconditional behavior.An EMA-smoothed policy supplies a stable reference velocity vold.
- NFT-Style Post-Training: NFT re-noises generated endpoints at training timesteps and regresses the target velocity through forward-process reconstruction instead of likelihood-ratio policy gradients.The re-noised state uses fresh noise and the endpoint’s reward determines the training branch.
- NFT-Style Post-Training: High-reward endpoints use ordinary prediction, whereas low-reward endpoints are pushed away through mirrored prediction around the EMA reference.With β = 1, the mirrored prediction becomes 2vold − vθ.
- NFT-Style Post-Training: The combined MSE objective trains quickly without likelihood ratios or SDE rollouts, while the EMA reference anchors and stabilizes post-training.In ReNFT, pull and push branches operate on different endpoints rather than the same x0.
4 Methodology
ReNFT repairs high-reward, low-diversity adapters by recalibrating probability mass through internal generator routes rather than adding external diversity objectives or modifying text representations. It prioritizes anti-hub prompts, constructs matched route-based counterfactuals, and applies reward-ranked joint-and-paired NFT updates.
- ReNFT targets an already-collapsed adapter without adding an external diversity objective or modifying the generator’s text representations.
- 4.1 Repair View: Collapse is treated as probability-mass compression over inherited pretrained capabilities, leaving suppressed modes reachable through internal generator routes.
- 4.2 Reading Post-Training Bias: The unconditional route exposes prompt-independent side effects, and anti-hub prompts are selected where this bias is most distinguishable.
- 4.3 Constructing Counterfactual Proposals: Two mixed routes branch from the same prompt and initial noise, combining conditional, frozen-base, and unconditional directions to produce matched alternatives.
- 4.3 Constructing Counterfactual Proposals: Reward ranking assigns the higher-reward endpoint as pull and the lower-reward endpoint as push, while an adaptive guard preserves route-A coverage when route-B wins dominate.
- 4.4 Reward-Ranked Repair: The NFT repair uses shared perturbations for paired updates and a native rollout-noise joint update before paired lower-noise updates; mixed routing adds no inference overhead.
5 Experiments
ReNFT is evaluated on PickScore and GenEval using quantitative, qualitative, and ablation studies. It recovers within-prompt diversity while retaining nearly all of NFT’s reward, with anti-hub selection and internal route construction supporting the repair.
- 5.2 Quantitative Results: ReNFT attains the highest diversity among compared post-trained methods on all three diversity metrics under both evaluation protocols.The comparison includes LPIPS-Div, DreamSim-Div, and DINOv3-Div.
- 5.2 Quantitative Results: NFT’s diversity collapses as reward post-training increases PickScore and GenEval, while ReNFT repairs earlier NFT checkpoints under the same remaining 50-step budget.DreamSim-Div falls from 0.242 to 0.119 on PickScore and from 0.339 to 0.149 on GenEval during continued NFT training.
- 5.2 Quantitative Results: ReNFT retains 98.9% of NFT’s PickScore reward and 99.0% of GenEval reward while improving DreamSim-Div by 58.8% and 55.0%, respectively.Its reward cost is roughly 1%, compared with a 55–59% diversity gain.
- 5.3 Qualitative Results: Qualitative comparisons show ReNFT restoring scene, style, background, and text variation while preserving prompt-specific count, position, color, and text fidelity.The comparisons use shared noise sets across GenEval, OCR, and PickScore prompts.
- 5.4 Ablation Study: Anti-hub prompt prioritization preserves reward trajectories but produces substantially stronger diversity recovery than random prompt selection.The controlled ablation keeps prompt count, sampling budget, and update rule fixed.
- 5.4 Ablation Study: Appendix ablations support the default rollout length, joint-noise choice, and EMA decay, while anti-hub pool-size and flipping-threshold sensitivity remain future work.The EMA ablation reports a reward–diversity trade-off: larger decays preserve reward, whereas smaller decays recover more diversity at greater reward cost.
6 Conclusion
The paper frames reward-induced mode collapse as internal probability-mass reallocation and presents ReNFT as an internal repair procedure. Across PickScore and GenEval, it preserves nearly all NFT reward while substantially improving diversity without external diversity objectives or text-encoder modification.
- 6 Conclusion: ReNFT treats collapse as suppression of base-model capabilities rather than deletion, making suppressed alternatives reachable through internal generator routes.The repair uses anti-hub prompts, matched counterfactuals, reward-based pull and push roles, and joint-and-paired NFT updates.
- 6 Conclusion: ReNFT retains 98.9% and 99.0% of NFT’s reward while improving DreamSim-Div by 58.8% and 55.0% on PickScore and GenEval.The paper positions this internal repair perspective as complementary to external interventions.
A Experimental Details
The experiments use SD3.5-M LoRA adapters and implement ReNFT through anti-hub probing, matched routed rollouts, reward ranking, and paired NFT updates with EMA.
- A Experimental Details: ReNFT trains SD3.5-M LoRA adapters with learning rate 10^-4, 48 unique samples per epoch, group size 24, and seed 42.Repair uses off-policy training with EMA decay 0.25, 36 rollout steps, and 7 training timesteps per update.
- A Experimental Details: Algorithm 1 initializes repair from a hacked checkpoint, frozen base route, reward model, encoders, prompt bank, routing patterns, and optimization hyperparameters.The trainable and EMA parameters both initially equal the hacked checkpoint.
- A Experimental Details: Each optimizer step samples unconditional probes, selects anti-hub prompts, and generates matched route endpoints from shared initial noise.The endpoints are ranked by reward before entering the repair batch.
- A Experimental Details: The repair applies one self-noise joint update followed by fresh-noise paired updates, then performs an EMA parameter update.The minimum route-A pull ratio is enforced by GuardFlip before training updates.
B.1 Ablation of Anti-Hub Selection and Pool Size
The ablations test anti-hub prompt selection and candidate-pool size under a fixed SD3.5-M PickScore repair budget. Anti-hub selection is important for diversity recovery, while coverage saturates at pools of 1,000 or more candidates.
- B.1 Ablation of Anti-Hub Selection and Pool Size: Replacing anti-hub prompts with random candidates keeps reward trajectories similar but sharply reduces diversity recovery.The on/off ablation preserves prompt count, sampling budget, and update rule.
- B.1 Ablation of Anti-Hub Selection and Pool Size: Anti-hub selection scores 1,000 candidate prompts using 24 unconditional probe generations and SigLIP2 text–image embeddings, choosing the lowest mean similarities.These prompts are farthest from the exposed unconditional tendency.
- B.1 Ablation of Anti-Hub Selection and Pool Size: Pools of 1,000, 2,000, and 3,000 candidates yield nearly identical reward and diversity trajectories, indicating rapid coverage saturation.The 500-candidate pool retains reward best but recovers noticeably less diversity; 1,000 is sufficient by default.
- B.1 Ablation of Anti-Hub Selection and Pool Size: The EMA ablation shows that larger decays preserve reward but barely recover diversity, whereas smaller decays recover more diversity at greater reward cost.On-policy training with decay 0 reaches the highest diversity and lowest reward; decay 0.25 is the reported method.
B.2 Ablation of EMA Decay
The ablations identify a reward–diversity trade-off governed by EMA decay, rollout length, route ordering, and repair duration. The selected configuration uses decay 0.25, 36-step rollouts, and a 50-step repair budget.
- Rollout length: The 36-step rollout is selected because it ends marginally higher on reward than 24 steps while remaining close on diversity.Twelve steps raise diversity most but reduce reward quickly; the extra sampling overhead is modest for roughly 50 optimizer steps.
- Route ordering: Placing b first and u last injects base-like structure early and exposes prompt-independent bias near the formed image.Earlier u produces a pull ratio near 1.0 and degenerate outputs, while a base step at the end lacks a subsequent conditional correction.
- Route pattern: Avoiding consecutive bb or uu blocks prevents extended reversion toward the base distribution or amplification of prompt-independent bias.The default repeating pattern places b at the start, u at the end, and keeps most steps on the conditional route θ.
- Repair duration: The guarded pull ratio stays within 0.52–0.62 after step 10, while the raw ratio falls to 0.32–0.42 after step 50.The widening gap reflects guard flips on low-margin pairs and motivates terminating repair at 50 steps.
C.2 Unconditional Sampling Analysis
Unconditional sampling reveals that NFT contracts diverse base-model outputs into reward-specific hubs. ReNFT reopens the distribution across subjects, styles, and rendering modes on both SD3.5-M and FLUX.2-klein-base.
- SD3.5-M: Under PickScore and GenEval, checkpoint H concentrates fixed-noise samples into distinct protocol-specific hubs.PickScore favors warm-toned detailed female portraits, whereas GenEval favors isolated full-body figures on plain backgrounds.
- FLUX.2-klein-base: ReNFT reopens the unconditional distribution with varied subjects, styles, and rendering modes on both FLUX.2-klein-base backbones.The contraction-and-reopening pattern is observed on both the 4B and 9B models, rather than only on SD3.5-M.
D.1 Full Comparison with E2PO
Against E2PO, ReNFT achieves higher diversity while retaining more of NFT’s reward, but the comparison combines reproduced E2PO values with unavailable official checkpoints and evaluation code. The same repair pattern also transfers to larger backbones without retuning.
- Comparison basis: E2PO’s unavailable checkpoint and evaluation code prevent direct comparison, so the main paper reuses its reported rewards and omits diversity metrics.A local reproduction yields slightly lower rewards than the paper-reported values, likely because of implementation differences.
- E2PO comparison: ReNFT achieves higher diversity on all six metrics while retaining 98.9% vs. 98.6% PickScore and 99.0% vs. 97.4% GenEval relative to NFT.These figures compare ReNFT with reproduced E2PO under the common evaluation protocol.
- Cross-backbone transfer: On FLUX.2-klein-base, NFT lowers DreamSim-Div by 48% and 44%, while ReNFT retains 98.9% and 98.7% of NFT’s reward and improves diversity by 33% and 35%.The results cover the 4B and 9B backbones, respectively, using the same 50-step repair budget without retuning.
E Extended Qualitative Results
Qualitative comparisons show ReNFT recovering composition, color, rendering-style, and structural diversity while avoiding many NFT artifacts. The recovery remains backbone-specific rather than forcing both models toward one common output range.
- 4B backbone: On FLUX.2-klein-base-4B prompts, ReNFT preserves prompt structure while increasing background, texture, text accuracy, and style diversity over NFT.For the axolotl, storefront, and squirrel prompts, ReNFT avoids NFT’s artifacts and produces more coherent or varied scenes.
- 9B backbone: On the 9B backbone, ReNFT reopens multiple compositional styles and subject structures after NFT collapses to narrow modes.The effect appears on axolotl, pineapple-character, and sheep-sign prompts, although some prompt-relevant forms remain shared by NFT and ReNFT.
- Overall comparison: Across illustrated prompts, ReNFT is visibly more varied than NFT in composition, color, and rendering style while retaining most of NFT’s reward.The qualitative pattern agrees with the quantitative results in Table 4.