Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Each language version is independently generated for its own context, not a direct translation.

这篇论文提出了一种让 AI“画得更好”的新方法。为了让你轻松理解，我们可以把训练 AI 生成图像（比如 Stable Diffusion）的过程想象成教一个新手画家画画。

1. 核心问题：新手画家需要“微调”

现在的 AI 模型（预训练模型）就像是一个已经看过几百万张画、有一定基础的新手画家。他画的大概像样，但如果你让他画“一只戴着墨镜的猫在冲浪”，他可能画得不够精准，或者猫看起来不像猫。

我们需要对他进行微调（Fine-tuning），让他更听指挥，画得更好。

2. 传统方法：笨拙的“试错法”

以前的微调方法（比如 PPO 或 DDPO）有点像让画家一边画一边改。

过程：画家画一笔，你告诉他“这里不对，改一下”，他再画一笔，你再告诉“那里也不对”。
缺点：这就像在黑暗中摸索，画家需要记住整个作画过程的每一步（从第一笔到最后一笔），计算量巨大，而且很容易“走火入魔”（训练不稳定），导致画出来的东西更奇怪。

3. 论文的新方法：聪明的“筛选法” (GRAFT & P-GRAFT)

这篇论文提出了两个核心概念，我们可以用**“选照片”和“中途截停”**的比喻来理解。

概念一：GRAFT（广义拒绝采样微调）—— “只挑最好的照片”

想象一下，你让画家画了 100 张“戴墨镜冲浪的猫”。

传统做法：把这 100 张图都拿给画家看，告诉他“这张稍微好点，那张差一点”，让他慢慢学。
GRAFT 做法：你直接只挑出那 10 张画得最好的，把剩下的 90 张扔掉（拒绝采样）。然后，你只让画家照着这 10 张最好的图重新练习。
好处：画家不再被那些画得烂的图干扰，只专注于学习“什么是好画”。这种方法在数学上被证明比“边画边改”更稳定、更高效。

概念二：P-GRAFT（部分 GRAFT）—— “中途截停，只练最难的部分”

这是论文最精彩的创新点。

画画的阶段：AI 画画是从一团模糊的噪点（像电视雪花）开始，一步步去噪，最后变成清晰的图像。
- 早期：从雪花变成模糊的轮廓（这时候画得像什么，其实很难判断，因为太模糊了）。
- 中期：轮廓清晰，能看出是猫还是狗了。
- 晚期：细节完善，加上墨镜和冲浪板。
传统微调的痛点：如果让 AI 从头（雪花）开始学，它要处理的信息太复杂，就像让一个刚学画画的人直接去画复杂的细节，容易学歪（偏差大）。
P-GRAFT 的妙招：
1. 让画家（预训练模型）先自己画，画到中途（比如已经能看出是猫了，但还没加墨镜）。
2. 这时候，你暂停一下，看看最终结果（加完墨镜的猫）好不好。
3. 如果最终结果好，你就把**中途那个“猫的轮廓”**挑出来，告诉画家：“记住这个轮廓，以后你就从这种状态开始画！”
4. 如果最终结果不好，就把那个中途的轮廓扔掉。
为什么有效？
- 比喻：这就好比教人开车。传统方法是让人从“怎么发动引擎”一直练到“怎么在高速上超车”，中间容易出错。P-GRAFT 是让人先练好“发动引擎和起步”（这是预训练模型擅长的），然后只教他“怎么在高速上超车”（这是微调要学的）。
- 科学原理：论文发现，在画画的中途，AI 面临的“不确定性”比刚开始时小，但比结束时大。在这个中间阶段进行微调，既不会太难（像刚开始那样混乱），也不会太简单（像最后那样已经定型了），达到了**“难度适中，进步最快”**的平衡点。

4. 另一个大招：逆噪修正 (Inverse Noise Correction) —— “给颜料桶换个配方”

这部分主要针对另一种 AI 模型（Flow Models）。

比喻：想象 AI 生成图像的过程，是从一个装满随机颜料的桶（噪声）开始，通过一个管道（模型），流出来变成一幅画。
问题：有时候流出来的画颜色不对，可能是因为颜料桶里的颜料配方（初始噪声分布）本身就有问题，而不是管道（模型）坏了。
解决方法：
1. 我们有一堆画得好的画（目标数据）。
2. 我们把这些画倒着通过管道，推回颜料桶里。
3. 这时候我们会发现，推回去的“颜料”和原来的“随机颜料”不一样。
4. 于是，我们训练一个小模型，专门负责把普通的随机颜料变成这种“特制颜料”。
5. 以后 AI 画画时，先用这个小模型把颜料调好，再流进大管道。
好处：不需要给 AI 奖励（比如“画得好给你糖吃”），只需要调整一下“原材料”，就能让画出来的东西质量大幅提升，而且速度更快。

5. 总结：这篇论文带来了什么？

更聪明地学习：不再让 AI 从头到尾死磕，而是只挑最好的结果，并且只在中途最关键的阶段进行强化训练。
效果更好：在文字生成图片（Text-to-Image）的任务上，比以前的方法画得更准、更听话（比如 VQAScore 分数更高）。
更省资源：不需要像以前那样算得那么累，就能达到更好的效果。
通用性强：不仅适用于画画，还能用来生成分子结构（制药）、设计房间布局等。

一句话总结：
这篇论文教 AI 画画时，不再让它“笨拙地从头改到尾”，而是教它**“只盯着最好的结果，并在最关键的中途阶段进行特训”，甚至还能“调整原材料的配方”**，从而用更少的力气，画出更完美的画。

Each language version is independently generated for its own context, not a direct translation.

1. 研究背景与问题 (Problem)

核心挑战：
预训练的生成模型（特别是扩散模型和流模型）通常需要根据特定任务或奖励反馈进行微调（Fine-tuning）。在语言模型中，常用的策略是基于强化学习（RL）的算法（如 PPO），结合 KL 散度正则化来防止模型偏离原始分布。然而，将此类方法直接应用于扩散模型面临以下困难：

边际似然不可计算： 扩散模型的边际似然（Marginal Likelihood）难以计算，导致 PPO 中所需的 KL 正则化项无法直接实施。
现有替代方案的缺陷：
- 忽略 KL 项会导致大规模训练不稳定。
- 使用轨迹 KL（Trajectory KL）约束虽然可行，但会导致初始值函数偏差（Initial Value Function Bias）问题，且效果往往不如预期。
采样效率与学习难度： 现有的基于拒绝采样（Rejection Sampling）的微调方法（如 RAFT）通常针对最终生成的数据分布，忽略了扩散过程中间状态的信息，且在高噪声阶段的学习可能面临较大的方差。

研究目标：
探索如何通过塑造扩散模型在中间噪声水平（Intermediate Noise Levels）的分布，来更有效地进行微调，并解决预训练流模型中的学习误差，而无需显式的奖励信号。

2. 方法论 (Methodology)

论文提出了两个核心框架：P-GRAFT 和 Inverse Noise Correction。

2.1 理论基础：广义拒绝采样微调 (GRAFT)

作者首先统一了现有的基于拒绝采样的微调方法（如 RAFT, RSO, Best-of-N），提出了广义拒绝采样微调 (Generalized Rejection sAmpling Fine-Tuning, GRAFT) 框架。

核心发现： GRAFT 本质上是在执行带有重塑奖励（Reshaped Rewards） 的 KL 正则化奖励最大化。
数学意义： 尽管扩散模型的边际似然不可计算，但 GRAFT 通过拒绝采样策略，隐式地实现了对边际 KL 约束的满足，从而避免了轨迹 KL 带来的偏差问题。

2.2 核心创新一：P-GRAFT (Partial-GRAFT)

基于 GRAFT，作者提出了P-GRAFT，旨在利用扩散模型的中间状态特性。

机制：
1. 从参考模型生成完整轨迹，计算最终样本的奖励。
2. 在中间去噪步骤 $t$ （而非最终步骤 $t=0$ ）对样本进行拒绝采样。
3. 使用被接受的中间状态 $X_t$ 作为训练数据，仅微调模型从 $T$ 到 $t$ 的部分。
4. 推理时，前 $T \to t$ 步使用微调模型， $t \to 0$ 步切换回参考模型。
理论依据（偏差 - 方差权衡）：
- 方差： 随着 $t$ 增大（噪声越多），给定 $X_t$ 时奖励 $r(X_0)$ 的条件方差增加（奖励信号更嘈杂）。
- 偏差/学习难度： 随着 $t$ 增大，去噪任务越简单（分数函数 $\nabla \log q_t$ 越接近高斯分布的简单形式），模型更容易学习，偏差减小。
- 结论： 选择一个“适当”的中间时间 $t$ 可以平衡方差（奖励噪声）和偏差（学习难度），从而获得比仅针对最终分布微调更好的效果。

2.3 核心创新二：逆噪声校正 (Inverse Noise Correction)

针对流模型（Flow Models），作者提出了一种无需显式奖励即可校正预训练模型误差的方法。

动机： 流模型的最终生成样本完全由初始噪声分布决定。如果预训练模型生成的分布与真实数据分布存在偏差，这种偏差源于初始噪声分布的偏差。
算法流程：
1. 利用预训练流模型的可逆性（Reversibility），通过反向欧拉法（Backward Euler） 将真实数据样本逆向映射回“噪声空间”，得到逆噪声分布 (Inverse Noise Distribution)。
2. 训练一个轻量级的“噪声校正器（Noise Corrector）”模型，学习从标准高斯分布映射到该逆噪声分布。
3. 推理阶段： 先生成校正后的噪声，再输入预训练的主模型生成图像。
优势： 无需奖励函数，仅通过校正初始噪声分布即可修正预训练模型的生成质量，且计算成本更低。

3. 主要贡献 (Key Contributions)

理论统一与扩展： 提出了 GRAFT 框架，证明了基于拒绝采样的微调等价于 KL 正则化的奖励最大化，解决了扩散模型边际 KL 不可计算的难题。
P-GRAFT 框架： 提出了部分去噪微调策略，通过中间分布塑造，利用偏差 - 方差权衡理论，显著提升了微调效率和质量。
逆噪声校正： 首次提出针对流模型的无奖励校正方法，利用流模型的可逆性修正初始噪声分布，提升了生成质量并降低了计算量。
广泛的实证验证： 在文本到图像（T2I）、布局生成、分子生成和无条件图像生成等多个任务上验证了方法的有效性。

4. 实验结果 (Results)

4.1 文本到图像生成 (Text-to-Image)

基准： 在 Stable Diffusion v2 (SDv2) 上进行微调。
指标： VQAScore（提示词对齐度）。
表现：
- P-GRAFT 在所有基准（GenAI-Bench, T2ICompBench++, GenEval）上均优于基线 SDv2、SDXL-Base 以及基于策略梯度的 DDPO 方法。
- 最佳配置： 在中间步骤 $t = 0.25N$ 处进行微调（P-GRAFT(0.25N)）效果最佳，验证了偏差 - 方差权衡理论。
- 对比 DDPO： P-GRAFT 在更少的采样轮次和梯度调用下，取得了显著更高的 VQAScore（例如在 GenAI-Bench 上提升了 8.81% 相对基线）。

4.2 布局与分子生成 (Layout & Molecule Generation)

模型： 基于 IGD (Interleaved Gibbs Diffusion) 框架。
布局生成： P-GRAFT 在无条件及类别条件生成中均提升了元素对齐度（Alignment），且 FID 分数优于 GRAFT，说明其更好地保留了预训练分布特性。
分子生成： 目标是生成稳定分子。P-GRAFT 在仅使用 1 倍采样轮次的情况下，达到了与 GRAFT（需 9 倍采样轮次）相当甚至更好的稳定性，且保持了分子多样性（通过去重策略）。

4.3 无条件图像生成与逆噪声校正

数据集： CelebA-HQ 和 LSUN-Church。
结果： 逆噪声校正显著降低了 FID 分数。
效率： 使用噪声校正器（100 步）+ 预训练模型（100 步）的组合，在生成质量和 FLOPs（浮点运算次数）上均优于仅使用预训练模型运行 1000 步的效果。校正器参数量仅为原模型的 1/4，大幅降低了延迟。

5. 意义与影响 (Significance)

理论突破： 为扩散模型的强化学习微调提供了坚实的理论基础，证明了拒绝采样可以隐式处理 KL 约束，无需计算不可行的边际似然。
效率提升： P-GRAFT 通过利用中间状态，证明了在去噪过程的早期进行微调比全程微调更高效，减少了训练成本并提高了收敛速度。
无奖励微调： 逆噪声校正提供了一种全新的视角，即通过修正输入端的噪声分布来优化生成模型，这为那些缺乏高质量奖励信号的场景（如科学发现、复杂结构生成）提供了新的微调思路。
通用性： 该方法不仅适用于连续扩散模型，还成功扩展到了离散 - 连续混合模型（如分子和布局生成），展示了强大的泛化能力。

总结： 该论文通过深入分析扩散过程的中间分布特性，提出了一套高效、理论完备的微调框架，显著提升了生成模型在多种任务上的性能，同时降低了计算和训练成本。