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$.