新智元报道
新智元报道
【新智元导读】MAPLE实验室提出通过强化学习优化图像生成模型的去噪过程，使其能以更少的步骤生成高质量图像，在多个图像生成模型上实现了减少推理步骤，还能提高图像质量。
无独有偶，齐国君教授领导的MAPLE实验室在OpenAI发布会一周前公布的工作中也发现了图像生成领域的主打方法扩散模型和流模型中也存在类似的过程：模型从高斯噪声开始的多步去噪过程也类似一个思维链，逐步「思考」怎样生成一张高质量图像，是一种图像生成领域的「图像链CoT」。
与OpenAI不谋而和的是，机器学习与感知（MAPLE）实验室认为强化学习微调方法同样可以用于优化多步去噪的图像生成过程，论文指出利用与人类奖励对齐的强化学习监督训练，能够让扩散模型和流匹配模型自适应地调整推理过程中噪声强度，用更少的步数生成高质量图像内容。
论文地址：https://arxiv.org/abs/2412.01243
研究背景
这一过程存在两个问题：
1. 经典的扩散模型训练方法只能保证每一步去噪能尽可能恢复出原始图像，不能保证整个去噪过程得到的图像符合人类的偏好；
2. 经典的扩散模型所有的图片都采用了同样的去噪策略和步数；而显然不同复杂度的图像对于人类来说生成难度是不一样的。
如下图所示，当输入不同长度的prompt的时候，对应的生成任务难度自然有所区别。那些仅包含简单的单个主体前景的图像较为简单，只需要少量几步就能生成不错的效果，而带有精细细节的图像则需要更多步数，即经过强化微调训练后的图像生成模型就能自适应地推理模型去噪过程，用尽可能少的步数生成更高质量的图像。
值得注意的是，类似于LLM对思维链进行的动态优化，对扩散模型时间进行优化的时候也需要动态地进行，而非仅仅依据输入的prompt；换言之，优化过程需要根据推理过程生成的「图像链」来动态一步步预测图像链下一步的最优去噪时间，从而保证图像的生成质量满足reward指标。
方法
OpenAI的O1通过在输出最终结果之前生成额外的token让LLM能进行额外的思考和推理，模型所需要做的最基本的决策是生成下一个token；而扩散和流匹配模型的「思考」过程则是在生成最终图像前，在不同噪声强度对应的扩散时间（diffusion time）执行多个额外的去噪步骤。为此，模型需要知道额外的「思考」步骤应该在反向扩散过程推进到哪一个diffusion time的时候进行。
为了实现这一目的，在网络中引入了一个即插即用的时间预测模块（Time Prediction Module, TPM）。这一模块会预测在当前这一个去噪步骤执行完毕之后，模型应当在哪一个diffusion time下进行下一步去噪。
具体而言，该模块会同时取出去噪网络第一层和最后一层的图像特征，预测下一个去噪步骤时的噪声强度会下降多少。模型的输出策略是一个参数化的beta分布。
由于单峰的Beta分布要求α>1且β>1，研究人员对输出进行了重参数化，使其预测两个实数a和b，并通过如下公式确定对应的Beta分布，并采样下一步的扩散时间。
在强化微调的训练过程中，模型会在每一步按输出的Beta分布随机采样下一个扩散时间，并在对应时间执行下一个去噪步骤。直到扩散时间非常接近0时，可以认为此时模型已经可以近乎得到了干净图像，便终止去噪过程并输出最终图像结果。
通过上述过程，即可采样到用于强化微调训练的一个决策轨迹样本。而在推理过程中，模型会在每一个去噪步骤输出的Beta分布中直接采样众数作为下一步对应的扩散时间，以确保一个确定性的推理策略。
设计奖励函数时，为了鼓励模型用更少的步数生成高质量图像，在奖励中综合考虑了生成图像质量和去噪步数这两个因素，研究人员选用了与人类偏好对齐的图像评分模型ImageReward（IR）用以评价图像质量，并将这一奖励随步数衰减至之前的去噪结果，并取平均作为整个去噪过程的奖励。这样，生成所用的步数越多，最终奖励就越低。模型会在保持图像质量的前提下，尽可能地减少生成步数。
将整个多步去噪过程当作一个动作进行整体优化，并采用了无需值模型的强化学习优化算法RLOO [1]更新TPM模块参数，训练损失如下所示：
在这一公式中，s代表强化学习中的状态，在扩散模型的强化微调中是输入的文本提词和初始噪声；y代表决策动作，也即模型采样的扩散时间；
代表决策器，即网络中A是由奖励归一化之后的优势函数，采用LEAVE-One-Out策略，基于一个Batch内的样本间奖励的差值计算优势函数。
通过强化微调训练，模型能根据输入图像自适应地调节扩散时间的衰减速度，在面对不同的生成任务时推理不同数量的去噪步数。对于简单的生成任务（较短的文本提词、生成图像物体少），推理过程能够很快生成高质量的图像，噪声强度衰减较快，模型只需要思考较少的额外步数，就能得到满意的结果；对于复杂的生成任务（长文本提词，图像结构复杂）则需要在扩散时间上密集地进行多步思考，用一个较长的图像链COT来生成符合用户要求的图片。
通过调节不同的γ值，模型能在图像生成质量和去噪推理的步数之间取得更好的平衡，仅需要更少的平均步数就能达到与原模型相同的性能。
同时，强化微调的训练效率也十分惊人。正如OpenAI最少仅仅用几十个例子就能让LLM学会在自定义领域中推理一样，强化微调图像生成模型对数据的需求也很少。不需要真实图像，只需要文本提词就可以训练，利用不到10,000条文本提词就能取得不错的明显的模型提升。
经强化微调后，模型的图像生成质量也比原模型提高了很多。可以看出，在仅仅用了原模型一半生成步数的情况下，无论是图C中的笔记本键盘，图D中的球棒还是图F中的遥控器，该模型生成的结果都比原模型更加自然。
针对Stable Diffusion 3、Flux-dev等一系列最先进的开源图像生成模型进行了强化微调训练，发现训练后的模型普遍能减少平均约50%的模型推理步数，而图像质量评价指标总体保持不变，这说明对于图像生成模型而言，强化微调训练是一种通用的后训练（Post Training）方法。
结论
在多个开源图像生成模型上的实验结果表明，这种强化微调方法能在保持图像质量的同时显著减少约50%推理步数，微调后模型生成的图像在视觉效果上也更加自然。可以看出，强化微调技术在图像生成模型中仍有进一步应用和提升的潜力，值得进一步挖掘。
