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

cs.LGAug 30, 2026

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

Authors: Yuchen Bao, Chao Wen, Haowei Wang, Ruoxin Chen, Donghao Luo, Jiahui Zhan, Wenjian Huang, Shen Chen,

Organizations: Southern University of Science and Technology · Tencent Youtu Lab

Abstract

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.

Figures & tables

Figure 1: Conceptual view of probability-mass recalibration. Pretraining spreads probability across the modes of each prompt, while unconditional behavior spans their marginal range. NFT collapses each prompt onto a reward-near mode and unconditional behavior onto a fixed mode, suppressing alternatives. ReNFT restores unconditional coverage and reward-boundary modes while leaving low-quality modes outside suppressed. Gray dashed lines compare the prompts’ highest-reward mode probabilities.

Figure 2: Training dynamics and ablations. (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 after 50 repair steps from the same $H$ ; the dashed line marks NFT at $H$ .

Figure 3: ReNFT pipeline. (a) Unconditional probes prioritize anti-hub prompts; (b) two policy-dominated mixed routes generate matched proposals from the same prompt and initial noise; (c) reward ranking and the adaptive flipping guard assign pull/push targets for joint-and-paired NFT repair. The VAE decoder is omitted for clarity.

Figure 4: Qualitative comparison on GenEval and OCR prompts. Each prompt group shows four samples from the same method under different initial noises.

Figure 5: Qualitative comparison on PickScore prompts. Each prompt group shows four samples from the same method under different initial noises.

Method	Model-Based Reward				Diversity
Method	PickScore $\uparrow$	Aesthetic $\uparrow$	ImageReward $\uparrow$	HPSv2 $\uparrow$	LPIPS-Div $\uparrow$	DreamSim-Div $\uparrow$	DINO-Div $\uparrow$
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
+ E2PO †	23.38	6.538	1.29	0.325	–	–	–

Table 1: Quantitative results on the PickScore test set. Shaded : in-domain training reward. Parenthesized row: reference only, excluded from ranking. Bold : best; underline : second best. † See Section 5.1 .

Method	Rule-Based Reward	Diversity
Method	GenEval $\uparrow$	LPIPS-Div $\uparrow$	DreamSim-Div $\uparrow$	DINO-Div $\uparrow$
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
+ E2PO †	0.932	–	–	–

Table 2: Quantitative results on the GenEval test set. Shaded : in-domain training reward. Parenthesized row: reference only, excluded from ranking. Bold : best; underline : second best. † See Section 5.1 .

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Ablation of Anti-Hub Candidate Pool Size on SD3.5-M (PickScore). The green trajectory (1,000 candidates) is identical to the default configuration in all other ablation figures.

Figure 7: Ablation of EMA Decay on SD3.5-M (PickScore). Larger decays preserve reward but barely recover diversity; smaller decays recover more diversity at larger reward cost, with on-policy training (decay

0

) reaching the highest diversity but the lowest reward. The

0.25

curve is Ours.

Figure 8: Ablation of Rollout Length and Joint-Step Self-Noise on SD3.5-M (PickScore). The 36-step curve is Ours; all variants branch from the same checkpoint

H

Figure 9: Dual-Route Pull Ratios during SD3.5-M PickScore Repair. (a) The raw route-

A

(

b\theta\theta

, base-probing) pull ratio, which equals route

A

’s win rate in the per-pair reward ranking before the adaptive flipping guard; (b) the guarded ratio after the guard enforces the

\geq 0.5

per-group floor. The vertical dotted line marks the reported 50-step budget.

Figure 10: Unconditional Sample Display across NFT Training and ReNFT Repair on SD3.5-M. Empty-prompt (CFG=0) samples from a fixed noise grid. The base generator branches into PickScore (top) and GenEval (bottom) repair trajectories, shown at the hacked checkpoint $H$ and repair steps 20 and 40.

Figure 11: Unconditional Sample Display on FLUX.2-klein-base. Empty-prompt samples from a fixed noise grid for the base model, the NFT checkpoint, and ReNFT (4B backbone on top, 9B at the bottom).

	PickScore Test Set							GenEval Test Set
Method	PickScore $\uparrow$	Aesthetic $\uparrow$	ImageReward $\uparrow$	HPSv2 $\uparrow$	LPIPS-Div $\uparrow$	DreamSim-Div $\uparrow$	DINO-Div $\uparrow$	GenEval $\uparrow$	LPIPS-Div $\uparrow$	DreamSim-Div $\uparrow$	DINO-Div $\uparrow$
SD3.5-M	(21.73)	(6.026)	(1.06)	(0.295)	(0.599)	(0.242)	(0.251)	(0.654)	(0.687)	(0.339)	(0.379)
+ DiffusionNFT	23.51	6.592	1.44	0.327	0.430	0.119	0.112	0.946	0.292	0.149	0.186
+ E 2 PO (reproduced)	23.17	6.312	1.34	0.327	0.498	0.166	0.159	0.921	0.463	0.204	0.258
+ Ours	23.26	6.344	1.41	0.323	0.565	0.189	0.182	0.937	0.496	0.231	0.295

Appendix

Table 3: Comparison with Locally Reproduced E 2 PO on SD3.5-M (PickScore and GenEval Test Sets). All metrics are computed under the same evaluation protocol. Unlike the main paper, which reuses E 2 PO’s paper-reported reward values, this table uses our local reproduction throughout. Shaded : in-domain training reward.

FLUX.2-klein-base-4B
Method	PickScore $\uparrow$	Aesthetic $\uparrow$	ImageReward $\uparrow$	HPSv2 $\uparrow$	LPIPS-Div $\uparrow$	DreamSim-Div $\uparrow$	DINO-Div $\uparrow$
Base	(21.26)	(5.995)	(0.99)	(0.289)	(0.533)	(0.246)	(0.274)
+ DiffusionNFT	23.38	6.723	1.35	0.320	0.433	0.127	0.122
+ Ours	23.13	6.489	1.39	0.318	0.485	0.169	0.157
FLUX.2-klein-base-9B
Base	(22.04)	(5.921)	(1.24)	(0.298)	(0.526)	(0.232)	(0.253)

Appendix

Table 4: Quantitative Results on Two FLUX.2-klein-base Backbones (PickScore Test Set). Shaded : in-domain training reward. Parenthesized rows: reference only, excluded from comparison.

Figure 12: Qualitative Comparison on FLUX.2-klein-base-4B. Each row is one PickScore prompt; each cell shows four samples from the same method under different initial noises.

Figure 13: Qualitative Comparison on FLUX.2-klein-base-9B. Same layout as Figure 12 .

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,192 papers

GuidanceNFT: NFT-Based Reward Fine-Tuning of Guided Diffusion Models
Oct 11, 2026Jiayin Liu, Zeyu Michael Li, Bingshuo Qian +1Diffusion Model Fine-Tuning T2I Diffusion Models
Tsinghua University · Duke University
TILT: Model-Intrinsic Reward Alignment For Compositional Diffusion
May 16, 2026Debottam Dutta, Jianchong Chen, Jaehoon Hahm +1Diffusion Model Alignment Diffusion Model Guidance
Electrical and Computer Engineering, University of Illinois Urbana-Champaign, Urbana, IL, USA · Zhejiang University, China · Physics, University of Illinois Urbana-Champaign, Urbana, IL, USA
NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment
Jun 16, 2026Jisung Hwang, Yunhong Min, Jaihoon Kim +2Diffusion Model Alignment Diffusion Model Guidance
1KAIST · †This work was done in part while the author was a visiting researcher at The University of Tokyo. · 2The University of Tokyo

Oct 11, 2026J/K move · Enter open · S save

GuidanceNFT: NFT-Based Reward Fine-Tuning of Guided Diffusion Models

Jiayin Liu, Zeyu Michael Li, Bingshuo Qian, Xiang Cheng

Tsinghua University · Duke University

Diffusion Model Fine-Tuning T2I Diffusion Models Diffusion Model Alignment Diffusion Model Guidance Classifier-Free Guidance

Open paper AskPDF arXiv

DiffusionNFT (NFT) fine-tunes diffusion models toward task-specific rewards, but optimizing these rewards does not guarantee fidelity to the full prompt. Although classifier-free guidance (CFG) strengthens prompt alignment, NFT is formulated as a CFG-free method, and applying CFG to the fine-tuned velocity field can reduce the optimized reward. This motivates optimizing rewards over the guided output distribution while seeking to retain its prompt alignment. However, applying standard NFT directly to guided samples can cause reward degradation or training collapse. We trace this failure to a mismatch between the reference velocity used in the NFT loss and the score of the guided sample distribution, which biases the improvement direction. We propose GuidanceNFT, which estimates the matching reference velocity through online flow matching on guided samples and uses it to correct both regression terms. With an exact reference, the corrected objective recovers NFT's intended update. In our experiments, the learned correction sustains reward improvement under both explicit and built-in guidance. On SD3.5-M, GuidanceNFT achieves higher optimized rewards than NFT with and without CFG across GenEval, text-rendering, and multi-reward objectives, together with better overall prompt alignment as measured by off-domain evaluation metrics. Experiments on models with built-in guidance, including Model Guidance and ELF, further demonstrate improved ImageNet generation, a better perplexity-entropy tradeoff on OpenWebText, and higher accuracy on GSM8K.

Abstract

Diffusion Models

Explore similar work

GuidanceNFT: NFT-Based Reward Fine-Tuning of Guided Diffusion Models

TILT: Model-Intrinsic Reward Alignment For Compositional Diffusion

NoiseTilt: Noise-Tilted Reverse Kernels for Diffusion Reward Alignment