Stitched Value Model for Diffusion Alignment

• URL: http://arxiv.org/abs/2605.19804v1
• Date: Tue, 19 May 2026 13:02:50 GMT
• Title: Stitched Value Model for Diffusion Alignment
• Abstract summary: Diffusion- or flow-based generative models must be aligned with task-specific rewards, such as prompt fidelity or aesthetic preference.<n>Here, we propose StitchVM, a model stitching framework that efficiently transfers reward models pretrained for clean images to the noisy latent regime.<n>We show that this approach yields improvements across a broad range of downstream steering and post-training methods.
• Score: 76.52330634059528
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Sample-Adaptive Latent Rewards for Uncertainty-Guided Diffusion Post-Training [51.18425726646089]
  We propose textscSURE, a unified latent-space framework for image and video diffusion models.<n>It learns reward distributions and directly uses their reliability to guide dense post-training.<n>textscSURE-REFL achieves the highest VBench quality, semantic, and total scores among the evaluated methods.
  arXiv  Detail & Related papers  (2026-08-06T14:55:42Z)
• Meta Flow Maps enable scalable reward alignment [22.109809498990362]
  Controlling generative models is computationally expensive.<n>This is because optimal alignment with a reward function--whether via inference-time steering or fine-tuning--requires estimating the value function.<n>We introduce Meta Flow Maps (MFMs), a framework consistency models and flow maps into extending the regime.
  arXiv  Detail & Related papers  (2026-01-20T19:39:56Z)
• Which Layer Causes Distribution Deviation? Entropy-Guided Adaptive Pruning for Diffusion and Flow Models [77.55829017952728]
  EntPruner is an entropy-guided automatic progressive pruning framework for diffusion and flow models.<n>Experiments on DiT and SiT models demonstrate the effectiveness of EntPruner, achieving up to 2.22$times$ inference speedup.
  arXiv  Detail & Related papers  (2025-11-26T07:20:48Z)
• MIRA: Towards Mitigating Reward Hacking in Inference-Time Alignment of T2I Diffusion Models [86.07486858219137]
  Diffusion models excel at generating images conditioned on text prompts.<n>The resulting images often do not satisfy user-specific criteria measured by scalar rewards such as Aesthetic Scores.<n>Recently, inference-time alignment via noise optimization has emerged as an efficient alternative.<n>We show that this approach suffers from reward hacking, where the model produces images that score highly, yet deviate significantly from the original prompt.
  arXiv  Detail & Related papers  (2025-10-02T00:47:36Z)
• Learn to Guide Your Diffusion Model [84.82855046749657]
  We study a technique for improving quality of samples from conditional diffusion models.<n>We learn guidance weights $omega_c,(s,t)$, which are functions of the conditioning $c$, the time $t$ from which we denoise, and the time $s$ towards which we denoise.<n>We extend our framework to reward guided sampling, enabling the model to target distributions tilted by a reward function.
  arXiv  Detail & Related papers  (2025-10-01T12:21:48Z)
• DeepCache: Accelerating Diffusion Models for Free [65.02607075556742]
  DeepCache is a training-free paradigm that accelerates diffusion models from the perspective of model architecture. DeepCache capitalizes on the inherent temporal redundancy observed in the sequential denoising steps of diffusion models. Under the same throughput, DeepCache effectively achieves comparable or even marginally improved results with DDIM or PLMS.
  arXiv  Detail & Related papers  (2023-12-01T17:01:06Z)
• DiffSmooth: Certifiably Robust Learning via Diffusion Models and Local Smoothing [39.962024242809136]
  We propose DiffSmooth, which first performs adversarial purification via diffusion models and then maps the purified instances to a common region via a simple yet effective local smoothing strategy. For instance, DiffSmooth improves the SOTA-certified accuracy from $36.0%$ to $53.0%$ under $ell$ $1.5$ on ImageNet.
  arXiv  Detail & Related papers  (2023-08-28T06:22:43Z)
• Generative Modeling with Flow-Guided Density Ratio Learning [12.192867460641835]
  Flow-Guided Density Ratio Learning (FDRL) is a simple and scalable approach to generative modeling. We show that FDRL can generate images of dimensions as high as $128times128$, as well as outperform existing gradient flow baselines on quantitative benchmarks.
  arXiv  Detail & Related papers  (2023-03-07T07:55:52Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.