将清晰图像的奖励模型迁移到噪声潜空间,实现高效扩散模型对齐。
Stitched Value Model for Diffusion Alignment

- 用冻结的扩散主干连接预训练的像素级奖励模型,构建混合价值函数。
- 在多个下游任务中提升速度3倍以上,峰值显存减半。
- 适合需要快速、低资源扩散对齐的研究与应用开发者。
实际应用中,基于扩散或流的生成模型必须与特定任务奖励对齐,如提示保真度或美学偏好。这种对齐具有挑战性,因为奖励定义在干净输出图像上,但对齐过程需要在噪声中间潜变量上估计价值函数。现有方法依赖Tweedie或蒙特卡洛近似,在估计偏差与计算成本间权衡:Tweedie估计效率高但有偏差,蒙特卡洛更准确但代价昂贵。一个自然选择是学习价值函数,但如何有效训练适用于噪声潜变量的强大通用价值模型仍是开放问题。本文提出StitchVM,一种模型拼接框架,可高效将为清晰图像预训练的奖励模型迁移至噪声潜空间。StitchVM从已有、截断的像素空间奖励模型出发,附加一个冻结的扩散主干作为其头部。由此产生的混合模型既保留了预训练良好的鲁棒奖励能力,又继承了扩散主干处理噪声潜变量的天然优势。拼接过程极为轻量,例如拼接并微调CLIP ViT-L和SD 3.5 Medium仅需10 GPU小时。通过将强大像素空间奖励模型提升至潜空间,StitchVM开启了一种新的扩散对齐范式:不再进行粗略且昂贵的逐样本价值函数近似,而是一次性构建适用于真实噪声潜变量的正确函数,并在大量样本和迭代中复用。实验表明,该方法在多种下游控制与后训练方法中均取得显著提升:DPS速度提升3.2倍,峰值显存减半;DiffusionNFT速度提升2.3倍。
原文摘要 · Abstract (English)
For practical use, diffusion- or flow-based generative models must be aligned with task-specific rewards, such as prompt fidelity or aesthetic preference. That alignment is challenging because the reward is defined for clean output images, but the alignment procedure requires value function estimates at noisy intermediate latents. Existing methods resort to Tweedie-style or Monte Carlo approximations, trading off estimator bias against computational cost: Tweedie estimates are efficient but biased, while Monte Carlo estimates are more accurate but require expensive rollouts. A natural alternative would be a learned value function, but it remains an open question how to effectively train a strong and general value model specifically for noisy latents. Here, we propose StitchVM, a model stitching framework that efficiently transfers reward models pretrained for clean images to the noisy latent regime. StitchVM starts from an existing, truncated pixel-space reward model and attaches a frozen diffusion backbone to it as its head. From the pixel-space model, the resulting hybrid retains a carefully pretrained, robust reward capability; from the diffusion backbone, it inherits its native ability to handle noisy latents. The stitching procedure is exceptionally lightweight, e.g., stitching and finetuning CLIP ViT-L and SD 3.5 Medium takes only 10 GPU-hours. By lifting powerful pixel-space reward models to latent space, StitchVM opens up a new style of diffusion alignment: instead of rough, yet costly per-sample approximation of the value function, the correct function for the actual, noisy latents is constructed once and then amortized over many samples and iterations. We show that this approach yields improvements across a broad range of downstream steering and post-training methods: DPS becomes $3.2\times$ faster while halving peak GPU memory, and DiffusionNFT becomes $2.3\times$ faster.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。