扩散模型对齐的缝合价值模型 Stitched Value Model for Diffusion Alignment
通过缝合技术将像素空间奖励模型高效转移到噪声潜在空间，实现快速准确的扩散对齐
前置知识
扩散模型（Diffusion Model）
扩散模型是一类生成模型，通过向数据中逐步添加噪声（前向过程），然后学习逐步去除噪声（反向过程）来生成新样本。以流匹配（Flow Matching）框架为例，定义了从干净潜在 z_0 到高斯噪声 z_1 的概率路径 z_t = alpha_t z_0 + sigma_t epsilon，其中 alpha_t = 1-t，sigma_t = t。模型学习预测速度场 u_t(z_t) 来指导去噪过程。
本文提出的StitchVM需要处理扩散模型的中间噪声潜在，理解扩散模型的数学形式和去噪轨迹是理解价值函数如何在中间状态工作的基础。
价值函数（Value Function）
价值函数 V_t(z_t) 衡量从噪声潜在 z_t 出发，完整去噪后获得的期望奖励。数学定义为 V_t(z_t) := log E[exp(r(z_0)) | z_t]，其中 z_0 ~ p_{0|t}(z_0|z_t)。在扩散对齐中，价值函数用于指导采样轨迹或提供训练信号，但直接评估困难，因为奖励 r(x_0) 通常只定义在干净图像上。
本文核心问题就是如何在噪声潜在空间准确高效地估计价值函数，理解价值函数的定义和作用是理解整个对齐框架的关键。
模型缝合（Model Stitching）
模型缝合是一种将两个预训练神经网络的某些层重新组合成新网络的技术。通过添加轻量级的缝合层 s_psi，将模型A的早期层 phi_{<= i} 与模型B的后期层 phi_{>= j} 连接，形成混合模型。缝合层可以是线性变换 W，通过最小化特征匹配损失 ||W u_{<= i}^theta(z_t) - r_{<= j-1}^phi(z_0)||_2^2 学习。
StitchVM的核心技术就是模型缝合，将扩散backbone和奖励模型缝合在一起，理解缝合原理有助于理解本文如何高效地将像素空间知识转移到潜在空间。
奖励倾斜（Reward Tilting）
奖励倾斜是扩散模型对齐的数学框架，通过奖励函数 r(x_0) 调整采样分布，使生成样本在保持真实性的同时最大化奖励。目标分布为 p^*(z_0) = (1/Z) p(z_0) exp(r(z_0))。实现方式包括推理时的梯度引导（修改速度场 u_r^t(z_t) = u_t(z_t) + c_t grad_{z_t} V_t(z_t)）和训练时的RL优化（最大化 E_{z_0 ~ p_theta}[r(z_0)] - D_KL(p_theta || p)）。
本文的StitchVM为奖励倾斜框架提供了更高效的价值函数估计，理解奖励倾斜的数学形式有助于理解StitchVM如何改进推理和训练时对齐方法。
研究动机
现有扩散模型对齐方法面临核心矛盾：奖励函数 r(x_0) 定义在干净图像上，但对齐过程需要反复评估中间噪声潜在 z_t 的价值函数。现有解决方案各有严重缺陷。Tweedie近似先估计 E[z_0|z_t] 再计算奖励，虽然高效（仅需一次去噪器和VAE解码），但在高噪声区域存在严重偏差（Jensen gap），且需要额外的反传链。蒙特卡洛近似从 z_t 展开多条去噪轨迹并平均奖励，理论上无偏但计算成本极高，需要多次完整去噪和VAE解码，在实际应用中往往不可行。直接学习噪声潜在价值函数的方案虽然理论上理想，但需要大规模数据训练，每次新的扩散backbone或改进的奖励模型都需要重新训练，成本不可持续。现有小规模训练的价值模型要么复用扩散特征导致泛化能力有限，要么从干净图像奖励模型投影到噪声空间但存在分布偏移，都无法匹配基础级像素空间奖励模型的广泛零样本能力。
本文的目标是本文提出StitchVM框架，目标是高效地将预训练的像素空间奖励模型转移到噪声潜在空间，使其成为可以直接在 z_t 上评估的价值函数。具体目标包括：仅用少量无标签图像和轻量级微调（约10 GPU小时），保持原始奖励模型在干净图像上的基准性能；直接处理噪声潜在，避免Tweedie的偏差和MC的高方差；适用于多种扩散backbone（SD 3.5 Medium/Large、FLUX）和奖励模型（CLIP、DFN-CLIP、HPSv2、Aesthetic Predictor）；同时提升推理时（DPS、FK steering）和训练时（DRaFT、DiffusionNFT）对齐方法的效率和效果。
与已有工作不同的是，本文的独特切入角度是将模型缝合技术首次系统性应用于噪声潜在价值函数学习。与传统方法不同，StitchVM不是从头训练价值函数，也不是用Tweedie/MC近似，而是通过找到扩散特征与奖励特征兼容的缝合接口 (i*, j*)，将冻结的扩散backbone作为头部，截断的奖励模型作为尾部，用轻量级缝合层连接。关键洞察是：扩散backbone天然处理噪声潜在，而奖励模型在像素空间有强大预训练能力，缝合可以在保留两者优势的同时避免各自缺点。这种缝合方案比大规模重新训练高效得多，又比Tweedie/MC近似准确得多，为扩散对齐提供了新的技术路径。
核心方法
StitchVM的整体思路是借鉴模型缝合技术，将两个已经预训练好的模型——扩散模型backbone和像素空间奖励模型——组合成一个能够在噪声潜在空间直接输出价值的混合模型。直觉上，扩散模型的中间层已经在处理噪声潜在，学习到了丰富的噪声感知特征；而奖励模型在大规模图像上预训练，掌握了精准的任务相关奖励预测能力。如果能找到两者表示兼容的接口，就可以将奖励模型的理解能力嫁接到扩散backbone的噪声处理能力上。技术路线分为两阶段：第一阶段通过穷举搜索找到最佳的缝合接口 (i*, j*)，即扩散模型的第 i* 层和奖励模型的第 j* 层，使得它们的特征可以用线性变换 W 对齐；第二阶段用无标签图像微调缝合层和截断的奖励模型，以原始奖励模型的分数作为监督目标，学习真正的价值函数 V_t(z_t) = E[r(z_0) | z_t]。
StitchVM的核心创新点是首次将模型缝合成功应用于噪声潜在价值函数学习。与已有方法的本质区别在于：现有工作要么通过Tweedie近似间接估计价值函数（存在偏差），要么通过MC采样直接估计（方差高、成本大），要么从头训练价值模型（成本不可持续）。StitchVM则利用了预训练模型的现有知识，通过一次性轻量级缝合和微调，就将像素空间奖励模型的能力转移到潜在空间。这种方法的关键洞察是：在合适的缝合点，扩散特征和奖励特征高度兼容，只需线性变换就能对齐，剩余的差距可以通过短期微调弥补，且不会损害奖励模型的预测能力。这与传统微调不同，StitchVM只训练缝合层和奖励模型尾部，完全冻结扩散backbone头部，既保留了扩散模型的噪声处理能力，又避免了破坏其预训练知识。
方法步骤详情
StitchVM包含两个主要步骤。步骤1：选择缝合接口。给定候选缝合索引 (i, j)，从干净图像 x_0 及其潜在 z_0 = E(x_0) 出发，采样噪声潜在 z_t ~ p_t(z_t|z_0)，提取配对特征 u_{<= i}^theta(z_t) 和 r_{<= j-1}^phi(z_0)。对每个候选对 (i, j)，通过特征匹配学习线性映射 W，目标是最小化 ||W u_{<= i}^theta(z_t) - r_{<= j-1}^phi(z_0)||_2^2。这个优化可以闭式求解，因此可以高效评估大量候选对。选择特征匹配损失最小的 (i*, j*) 作为最佳缝合接口。步骤2：微调StitchVM。由于 (i*, j*) 已选择为表示最兼容的点，扩散特征经线性变换后已接近奖励模型期望的特征，只存在小的不匹配。短期微调缝合层 s_psi 和截断奖励模型 r_{>= j*}^phi 就能补偿这个差距。微调使用无标签干净图像 z_0，对每个 z_0 采样噪声潜在 z_t ~ p_t(z_t|z_0)，以原始奖励模型的分数 r_phi(z_0) 作为监督目标，最小化损失 L_value(omega) = E_{z_0, t, epsilon}[||V_(i*, j*)^omega(z_t) - r_phi(z_0)||_2^2]。可以证明这个损失的最小化器满足 V_(i*, j*)^omega*(z_t) = E[r_phi(z_0) | z_t]，即真正的价值函数。
技术新颖性
StitchVM的技术新颖性体现在多个方面。首先是首次系统性将模型缝合技术应用于生成模型对齐任务，特别是将像素空间模型缝合到噪声潜在空间。与现有缝合工作不同（如VIST3A缝合到VAE解码器），StitchVM缝合的是扩散backbone，能够直接处理高噪声潜在，在高噪声区域表现显著优于VAE缝合基线。其次，StitchVM提出了在扩散-奖励模型对之间寻找兼容缝合接口的实用方法，通过闭式求解的特征匹配来评估候选对，避免了昂贵的端到端搜索。第三，StitchVM采用off-policy回归目标，不需要在策略采样，进一步节省计算。第四，StitchVM的创新还体现在推理和训练时对齐方法的改进上：在推理时，StitchVM可以直接提供价值函数梯度 grad_{z_t} V_t(z_t)，避免长反传链；在训练时，可以在中间噪声潜在停止并直接获得监督，避免完整展开。这些创新使StitchVM在保持效果的同时大幅提升了效率。
实验结果
论文通过一系列实验验证了StitchVM的有效性。在主要结果中，StitchVM在低噪声区域（sigma <= 0.5）与原始干净奖励模型性能接近，在CLIP检索（MSCOCO和Flickr30K的Avg. Recall@1）、HPSv2偏好预测（HPDv2和ImageReward的准确率）和美学预测（AVA测试集的SRCC）等任务上都保持了基准性能。随着噪声水平增加，StitchVM表现出渐进的性能下降，但仍然显著优于基线方法。具体来说，StitchVM在高噪声（sigma = 0.9）时仍保持CLIP检索约50%的干净性能，而VAE缝合基线几乎崩溃。与NoisyCLIP相比，StitchVM尽管训练规模小得多（AVA+HPDv2 vs LAION-400M），在所有噪声水平上都表现更好。与DiNa-LRM相比，StitchVM使用无标签图像而非大规模偏好数据（HPDv3），在HPDv2和ImageReward上都取得更高准确率。在推理时方法上，StitchVM使DPS在SD 3.5 Medium上ImageReward从0.96提升到1.22，HPSv2从0.335提升到0.348，同时峰值GPU内存从56.4GB降到26.0GB（下降54%），采样时间从52.8秒降到16.5秒（提速3.2倍）。在SD 3.5 Large上也有类似改进。FK steering方面，StitchVM在SD 3.5 Medium上使ImageReward从0.93（FKS）和0.91（BoN）提升到1.10，GenEval从0.62提升到0.69。实验还展示了新的缩放轴：增加每个粒子的局部提案数 M 比增加粒子数 N 更高效，(N=8, M=6) 以33%更低的成本达到标准FKS N=14 的性能。在训练时方法上，StitchVM使DRaFT-1的GPU小时从128.1降到94.8（下降26%），GenEval从0.53提升到0.69；DiffusionNFT的GPU小时从191.5降到84.7（下降56%），而GenEval从0.67提升到0.68。StitchVM训练成本方面，每个模型在512x512分辨率仅需约10 GPU小时（GH200），在1024x1024分辨率需24-32 GPU小时，包括缝合层搜索和微调两个阶段。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| CLIP零样本图像-文本检索 | Avg. Recall@1 | StitchVM保持接近干净奖励模型性能（sigma=0.1时约48-50%） | VAE缝合在高噪声（sigma=0.9）崩溃至接近0 | 在高噪声区域显著优于所有基线 |
| HPSv2偏好预测 | 准确率 | StitchVM在HPDv2上保持约80%（sigma=0.1），在ImageReward上保持约68% | DiNa-LRM在HPDv2上约75%，在ImageReward上约60% | 在两个数据集上均优于DiNa-LRM |
| DPS推理时对齐 | ImageReward / GPU内存 / 采样时间 | 1.22 / 26.0GB / 16.5秒（SD3.5-Medium） | 0.96 / 56.4GB / 52.8秒 | ImageReward提升27%，内存下降54%，速度提升3.2倍 |
| FK steering推理时对齐 | ImageReward / GenEval | 1.10 / 0.69（SD3.5-Medium） | 0.93 / 0.62（FKS） | ImageReward提升18%，GenEval提升11% |
| DRaFT训练时对齐 | GenEval / GPU小时 | 0.69 / 94.8（DRaFT-1+StitchVM） | 0.53 / 128.1（DRaFT-1） | GenEval提升30%，计算成本下降26% |
| DiffusionNFT训练时对齐 | GenEval / GPU小时 | 0.68 / 84.7 | 0.67 / 191.5 | 计算成本下降56%，GenEval略提升 |
局限与改进
作者承认的局限性包括：StitchVM的成功依赖于找到兼容的缝合接口，虽然论文显示多个扩散backbone和奖励模型对都存在这样的接口，但理论上某些模型对可能不存在高度兼容的表示，此时可能需要更复杂的缝合层架构。此外，StitchVM虽然高效，但仍需要一次性10-32 GPU小时的训练成本，对于资源受限的研究者可能仍有一定门槛。论文主要在流匹配框架上验证，虽然提到可以扩展到传统DDPM，但未在DDPM上系统评估。我的观察是：StitchVM在极高噪声（sigma > 0.9）时性能下降显著，虽然比基线好但仍远低于干净性能，这可能限制在早期高噪声阶段的应用；StitchVM的微调使用无标签图像，但不同任务可能需要领域特定的数据来获得最佳迁移；论文主要在图像生成上验证，在视频、3D生成等更复杂的模态上的效果未知；StitchVM与新的奖励模型或扩散架构的兼容性需要进一步验证。
独立分析的弱点
StitchVM的几个潜在弱点值得指出。首先，缝合接口的搜索是穷举的，对于层数很深的模型可能搜索成本较高，虽然论文中搜索成本相对训练成本很小，但未来可以考虑更智能的接口选择策略，比如基于表示相似度的启发式方法或可微分搜索。其次，StitchVM使用线性缝合层，虽然论文显示线性变换足够，但对于某些表示差距较大的模型对，可能需要更复杂的非线性缝合层或引入适配器模块。第三，StitchVM的微调是离策略的，使用无标签图像，但不同任务可能需要不同的噪声分布或采样策略，可以考虑引入任务特定的微调数据或强化学习式的在线微调。第四，StitchVM目前主要关注单一奖励模型，但在实际应用中可能需要同时优化多个奖励（如文本对齐、美学、安全性等），可以考虑多任务学习或奖励加权策略。改进方向包括：自适应缝合层架构选择、基于梯度流的接口优化、在线微调策略、多奖励融合机制等。
未来方向
作者提出的未来工作方向包括：探索StitchVM在其他模态（如视频、3D生成）上的应用，验证缝合框架的泛化能力；研究更复杂的缝合层架构，以处理表示差距更大的模型对；探索动态缝合策略，根据噪声水平或任务阶段调整缝合接口或缝合层参数；将StitchVM与新兴的对齐方法结合，如基于Transformer的扩散模型、多模态对齐等。基于本文成果可以延伸的研究方向包括：将缝合框架扩展到其他类型的对齐，如RLHF、指令微调等；研究缝合技术在模型压缩、知识迁移等其他任务中的应用；探索端到端的缝合学习，同时优化接口和缝合层；研究StitchVM的可解释性，分析缝合层学到了什么表示转换；将StitchVM与自动化机器学习结合，自动搜索最优的缝合配置；探索StitchVM在安全对齐中的应用，如毒性检测、偏见缓解等。
复现评估
论文提供了项目页面（https://gohyojun15.github.io/StitchVM），但截至论文发表时代码和数据集的开源情况未明确说明。从复现难度评估，主要挑战包括：需要访问预训练的扩散backbone（SD 3.5 Medium/Large、FLUX）和奖励模型（CLIP、DFN-CLIP、HPSv2、Aesthetic Predictor），其中SD 3.5和FLUX可能需要特定许可；训练数据AVA和HPDv2需要申请或通过官方渠道获取；论文使用GH200 GPU进行训练，虽然可以用其他GPU替代但可能需要调整批量大小和训练时间；缝合接口搜索需要实现特征匹配和闭式求解，虽然理论上简单但实现细节可能影响结果。总体而言，如果代码和模型权重开源，复现难度中等；如果需要从头实现模型缝合逻辑，复现难度较高。论文提供了详细的实验设置和超参数，有利于复现。建议作者提供完整代码、预训练的StitchVM权重、缝合接口搜索脚本和评估基准，以促进社区复现和进一步研究。
论文图表
