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

Yuchen Bao^1,2, Chao Wen², Haowei Wang², Ruoxin Chen², Donghao Luo², Jiahui Zhan², Wenjian Huang¹, Shen Chen², Yiting Wang², Taiping Yao², Chengjie Wang², Shouhong Ding², Jianguo Zhang^1,*

¹Southern University of Science and Technology · ²Tencent Youtu Lab
^*Corresponding author

arXiv Code (Coming Soon) BibTeX

💡Overview

“Can an already-collapsed adapter be repaired using only distributions already inside the generator, without external diversity objectives or text-encoder modification, and without sacrificing its acquired reward?”

Online reward post-training sharply raises a diffusion generator’s reward, but probability mass concentrates on a few reward-favored modes: different initial noises under the same prompt produce highly similar outputs — a within-prompt mode collapse in structure and style, visible in the qualitative comparisons. Existing mitigations rely on external signals or interfaces, and, to our knowledge, none repairs an adapter that has already collapsed while preserving the acquired reward.

We observe that online reward post-training primarily reallocates probability mass over capabilities inherited from pretraining rather than learning new visual content — mode collapse is therefore suppression, not deletion, and can be reversed from within the generator (see the unconditional diagnostic). ReNFT realizes this repair through internal probability-mass recalibration: (a) unconditional probes prioritize anti-hub prompts, (b) matched internal routes generate counterfactual proposals, and (c) reward-guided pairing assigns positive/negative roles for joint-and-paired NFT repair. Relative to DiffusionNFT (NFT), ReNFT retains 98.9% and 99.0% of NFT’s endpoint reward on PickScore and GenEval while improving DreamSim-Div by 58.8% and 55.0%, respectively (Tables 1–2).

Conceptual view of probability-mass recalibration

Paper Fig. 1 — Pretraining spreads probability mass across the modes of each prompt. NFT collapses every prompt onto a reward-near mode and unconditional behavior onto a fixed mode, suppressing — not deleting — the alternatives. ReNFT restores unconditional coverage and reward-boundary modes, while low-quality modes remain suppressed.

🧩Method

Post-hoc repair of a collapsed adapter through internal routes of the generator

Paper Fig. 3 — The ReNFT pipeline: (a) unconditional probes prioritize anti-hub prompts; (b) matched internal routes generate counterfactual proposals; (c) reward-guided pairing assigns positive/negative roles for joint-and-paired NFT repair.

Design Highlights

Suppression, Not Deletion

Online reward post-training reallocates probability mass over capabilities inherited from pretraining rather than learning new visual content, so collapse is reversible from within the generator.

Anti-Hub Prompt Prioritization

Unconditional probes prioritize “anti-hub” prompts where the prompt-independent bias is easiest to expose and most likely to yield distinguishable proposals.

Matched Counterfactual Routes

Two matched internal routes branch from the same prompt and initial noise, exposing a base-directed alternative and the post-trained unconditional tendency.

Reward-Guided Pairing

The original reward compares each matched pair and assigns the positive and negative roles, so roles are decided per pair by the reward gap rather than hard-coded by route.

Joint-and-Paired NFT Repair

DiffusionNFT-style pull and mirrored-push updates realize the repair: one joint update at the shared trajectory origin, then paired updates across noise levels.

No External Diversity Objective

ReNFT repairs the collapsed adapter without adding a separate diversity objective or modifying the text encoder.

📈Training Dynamics

The tracked baselines gain reward at the cost of within-prompt diversity; ReNFT repairs the collapsed NFT checkpoint in 50 steps

Panels (a–d) track Flow-GRPO, DiverseGRPO, and NFT on both protocols. Across these baselines, reward gains are accompanied by an overall decline in within-prompt DreamSim-Div. ReNFT (green) branches from the hacked checkpoint H and recovers diversity while keeping reward close to the NFT endpoint. Panels (e–f) ablate the mixed-route pattern and anti-hub prompt prioritization.

Paper Fig. 2 — (a–b) PickScore reward and DreamSim-Div; (c–d) GenEval reward and DreamSim-Div. ReNFT branches from the hacked checkpoint H (star) and repairs for 50 steps; the dashed line marks the base generator. (e–f) Mixed-route pattern and anti-hub ablations over a 50-step repair; the dashed line marks NFT at H.

🖼️Qualitative Results

Base / NFT / DiverseGRPO / ReNFT on shared prompts, each with samples from different initial noises

PickScore Prompts

In the shown PickScore grids, NFT and DiverseGRPO exhibit recognizable hacking tendencies — a recurring color tone, densely packed elements, and a preference for single female subjects. ReNFT restores scene- and style-level variation across seeds without reverting to the base generator.

Paper Fig. 5 — Qualitative comparison on PickScore prompts.

GenEval & OCR Prompts

In the shown GenEval examples, NFT favors plain, highly uniform backgrounds and reduced variation across seeds. ReNFT restores background and compositional variation while preserving the requested objects and attributes. On the shown OCR prompts outside the PickScore training distribution, ReNFT produces more legible and varied text than NFT and DiverseGRPO.

Paper Fig. 4 — Qualitative comparison on GenEval and OCR prompts.

📊Quantitative Results

Reward retention and diversity recovery on the PickScore and GenEval protocols

Method	Model-Based Reward				Diversity
Method	PickScore ↑	Aesthetic ↑	ImageReward ↑	HPSv2 ↑	LPIPS-Div ↑	DreamSim-Div ↑	DINO-Div ↑
SD3.5-M	21.73	6.026	1.06	0.295	0.599	0.242	0.251
+ Flow-GRPO	22.90	6.277	1.31	0.305	0.463	0.122	0.134
+ DiffusionNFT	23.51	6.592	1.44	0.327	0.430	0.119	0.112
+ DiverseGRPO	22.99	6.234	1.27	0.332	0.506	0.163	0.145
+ E²PO^†	23.38	6.538	1.29	0.325	–	–	–
+ ReNFT (Ours)	23.26	6.344	1.41	0.323	0.565	0.189	0.182

Paper Tab. 1 — PickScore protocol at 750 optimizer steps; ReNFT branches from the NFT checkpoint at step 700 and repairs for 50 steps. Bold: best; underline: second best; gray row: untrained base reference; shaded column: in-domain training reward. ^†E²PO is not open-sourced; we match its official training budget and reuse its paper-reported reward values; its paper reports no results on our diversity metrics.

Method	Rule-Based Reward	Diversity
Method	GenEval ↑	LPIPS-Div ↑	DreamSim-Div ↑	DINO-Div ↑
SD3.5-M	0.654	0.687	0.339	0.379
+ Flow-GRPO	0.802	0.343	0.176	0.205
+ DiffusionNFT	0.946	0.292	0.149	0.186
+ DiverseGRPO	0.869	0.454	0.198	0.246
+ E²PO^†	0.932	–	–	–
+ ReNFT (Ours)	0.937	0.496	0.231	0.295

Paper Tab. 2 — GenEval protocol at 300 optimizer steps; ReNFT branches from the NFT checkpoint at step 250 and repairs for 50 steps. Marking conventions as above.

Starting from the collapsed NFT adapter, ReNFT retains 98.9% and 99.0% of NFT’s reward on PickScore and GenEval, while improving DreamSim-Div by 58.8% and 55.0%, respectively. All three diversity metrics rank best among the post-trained methods evaluated for diversity under both protocols, while reward remains close to the NFT endpoint. The same repair is visible in the training curves and the qualitative grids.

🔍Appendix Highlights

Diagnostic unconditional samples and cross-backbone qualitative results

Unconditional Samples across Repair (SD3.5-M)

Empty-prompt samples from a fixed noise grid, held constant across training and repair stages, expose the prompt-independent bias. Online reward post-training contracts these samples toward protocol-specific hubs, while repair progressively broadens them without simply reverting to the base generator. This diagnostic pattern supports the paper’s suppression-not-deletion hypothesis and complements the conditional diversity collapse tracked in the training dynamics and qualitative comparisons.

Paper Fig. 10 — Unconditional sample display across NFT training and ReNFT repair on SD3.5-M.

Cross-Backbone Generalization (FLUX.2-klein-9B)

ReNFT also transfers to the FLUX.2-klein-base-9B backbone, restoring visible variation relative to NFT while retaining most of NFT’s PickScore reward.

Paper Fig. 13 — Qualitative comparison on FLUX.2-klein-base-9B.

📚Citation

If you find ReNFT useful in your research, please cite our work

@misc{bao2026renft,
      title={ReNFT: Repairing Mode Collapse in Reward Post-Training via Internal Probability-Mass Recalibration},
      author={Yuchen Bao and Chao Wen and Haowei Wang and Ruoxin Chen and Donghao Luo and Jiahui Zhan and Wenjian Huang and Shen Chen and Yiting Wang and Taiping Yao and Chengjie Wang and Shouhong Ding and Jianguo Zhang},
      year={2026},
      eprint={2609.00061},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.00061},
}

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