[Submitted on 19 May 2026 (v1), last revised 7 Nov 2026 (this version, v2)]
Authors:Hyojun Go, Hyungjin Chung, Prune Truong, Goutam Bhat, Li Mi, Zhaochong An, Zixiang Zhao, Dominik Narnhofer, Serge Belongie, Federico Tombari, Konrad Schindler
Abstract: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 $512times512$ 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.7times$ with half the peak GPU memory, and accelerates DiffusionNFT by $2.3times$.
Submission history
From: Hyojun Go [view email]
[v1] Tue, 19 May 2026 13:02:50 UTC (2,713 KB)
[v2] Sat, 7 Nov 2026 20:50:14 UTC (2,584 KB)
