← 研踪 PaperTrailarXiv:2605.19804 · v2
Stitched Value Model for Diffusion Alignment
首次提交 2026-05-19 21:02
版本更新 2026-11-08 04:50
雷达观测 2026-11-10 12:16
摘要
Reward-based alignment of diffusion and flow models requires estimating the value function at noisy intermediate latents, although pretrained reward models operate on clean images. Common approaches use Tweedie-style clean-image proxies or Monte Carlo rollouts, introducing approximation bias or substantial computational cost. Directly learning the value function offers an alternative, but matching the performance and generalization of foundation-scale pixel-space reward models on noisy latents remains challenging. We address this gap by transferring pretrained reward capabilities to noisy latents. We propose \textbf{StitchVM} (Stitched Value Model), a model stitching framework that connects a frozen diffusion-backbone prefix to a pretrained reward-model suffix through a compatible feature interface. A closed-form feature fit selects the interface, and lightweight distillation adapts the stitched model, combining noise-aware diffusion features with pretrained reward capability. For example, adapting CLIP ViT-L with SD~3.5 Medium at $512\times512$ resolution takes approximately 10 GPU-hours on a single GH200. Once trained, StitchVM provides an amortized alternative to Tweedie-style and Monte Carlo approximations of the value function: it estimates value directly from noisy latents, avoiding clean-image reconstruction and additional denoising rollouts for each evaluation. We demonstrate benefits across inference-time steering and post-training. StitchVM improves the quality and efficiency of FK steering and DRaFT, accelerates DPS by $3.7\times$ with half the peak GPU memory, and accelerates DiffusionNFT by $2.3\times$.
作者备注：Project page: https://gohyojun15.github.io/StitchVM/
观测到的版本
相邻记录可能跳过中间版本；首次发现 v2 及以上也会标为版本更新。
v2 · 2026-11-08 04:50 · 最新观测版本
Stitched Value Model for Diffusion Alignment
Reward-based alignment of diffusion and flow models requires estimating the value function at noisy intermediate latents, although pretrained reward models operate on clean images. Common approaches use Tweedie-style clean-image proxies or Monte Carlo rollouts, introducing approximation bias or substantial computational cost. Directly learning the value function offers an alternative, but matching the performance and generalization of foundation-scale pixel-space reward models on noisy latents remains challenging. We address this gap by transferring pretrained reward capabilities to noisy latents. We propose \textbf{StitchVM} (Stitched Value Model), a model stitching framework that connects a frozen diffusion-backbone prefix to a pretrained reward-model suffix through a compatible feature interface. A closed-form feature fit selects the interface, and lightweight distillation adapts the stitched model, combining noise-aware diffusion features with pretrained reward capability. For example, adapting CLIP ViT-L with SD~3.5 Medium at $512\times512$ resolution takes approximately 10 GPU-hours on a single GH200. Once trained, StitchVM provides an amortized alternative to Tweedie-style and Monte Carlo approximations of the value function: it estimates value directly from noisy latents, avoiding clean-image reconstruction and additional denoising rollouts for each evaluation. We demonstrate benefits across inference-time steering and post-training. StitchVM improves the quality and efficiency of FK steering and DRaFT, accelerates DPS by $3.7\times$ with half the peak GPU memory, and accelerates DiffusionNFT by $2.3\times$.
查看此版本解读 →查看此版本原文 ↗