Technical Summary: Stitched Value Model for Diffusion Alignment (StitchVM)

1. Problem Statement

Diffusion and flow-based generative models excel at capturing data distributions but often require alignment with task-specific rewards (e.g., prompt fidelity, aesthetic preference, human preference) to be practically useful. The core challenge in this alignment is that rewards are typically defined for clean output images ($x_0$), whereas alignment procedures (both at inference and training time) require estimating the expected reward of clean samples induced by noisy intermediate latents ($z_t$). This estimation is known as the value function.

Existing methods for approximating this value function suffer from significant trade-offs:

• Tweedie Approximation: Estimates the posterior mean of the clean sample and computes the reward on that proxy. While efficient, it introduces substantial bias in high-noise regimes and requires expensive additional evaluations of the denoiser and VAE decoder.
• Monte Carlo (MC) Approximation: Rolls out multiple denoising trajectories from $z_t$ and averages the rewards. While more accurate, it incurs prohibitive computational costs due to the need for multiple full rollouts.
• Learned Value Models: Directly training a value function for noisy latents is an open problem. Existing attempts often rely on small-scale training or narrow preference datasets, resulting in inferior accuracy and generalization compared to foundation-scale pixel-space reward models. Furthermore, training a new value model for every new diffusion backbone or reward model is computationally unsustainable.

2. Methodology: StitchVM

The authors propose StitchVM (Stitched Value Model), a framework that transfers the capabilities of pretrained, high-quality pixel-space reward models into the noisy latent regime with minimal computational cost. The approach leverages model stitching, a technique that recombines layers from different neural networks.

Core Architecture

StitchVM constructs a hybrid model by attaching a truncated, frozen diffusion backbone (the "head") to a sliced, pretrained reward model (the "tail") via a lightweight stitching layer.

• Head: A diffusion backbone (e.g., SD 3.5, FLUX) truncated at layer $i$. This component natively processes noisy latents and extracts noise-aware features.
• Tail: A pretrained reward model (e.g., CLIP, HPSv2, Aesthetic Predictor) starting from layer $j$. This component retains the robust, task-specific reward capability learned during large-scale pretraining.
• Stitching Layer ($s_\psi$): A lightweight module that aligns the feature space of the diffusion head with the reward tail.

Training Procedure

The training of StitchVM occurs in two stages:

1. Interface Selection (Closed-Form): The authors search for the optimal stitching indices $(i, j)$ where the diffusion features (from noisy latents) and reward features (from clean images) are most compatible. This is achieved by fitting a linear mapping $W$ between paired features extracted from a small probe set of clean images and their corresponding noisy latents. The pair minimizing the feature-matching loss is selected.
2. Lightweight Finetuning: Once the interface is fixed, the stitching layer and the truncated reward tail are finetuned using unlabeled clean images. The objective is to regress the value function $V(z_t)$ to the original reward score $r(z_0)$ of the clean image. This off-policy objective allows the model to learn the expected reward of clean samples conditioned on noisy latents without requiring full rollouts.

Application to Alignment

StitchVM serves as a direct value function estimator, replacing approximations in both inference-time and training-time alignment:

• Inference-Time:
  • DPS (Diffusion Posterior Sampling): Replaces the Tweedie-based gradient with the direct gradient of StitchVM ($\nabla_{z_t} V(z_t)$), eliminating the need for backpropagation through the denoiser and VAE.
  • FK Steering: Enables "proposal scaling," where each particle spawns multiple local proposals. These are scored cheaply by StitchVM to select the best candidate, a process that would be too costly with Tweedie approximations.
• Training-Time:
  • Direct Reward Finetuning (e.g., AlignProp, DRaFT): Allows training to stop at intermediate noisy latents $z_\tau$ and use the StitchVM prediction as the supervision target, avoiding full trajectory rollouts and long backpropagation chains.
  • DiffusionNFT: Replaces terminal clean rewards with intermediate value estimates, enabling efficient online RL on the forward process.

3. Key Contributions

• A New Paradigm for Value Modeling: StitchVM demonstrates that powerful pixel-space reward models can be effectively transferred to the noisy latent regime without large-scale retraining, simply by stitching them to a diffusion backbone.
• Efficiency and Accuracy: The method achieves state-of-the-art performance in value estimation while being exceptionally lightweight. For example, stitching CLIP ViT-L into an SD 3.5 Medium value model takes only $\approx 10$ GPU-hours.
• Performance Gains in Alignment:
  • Inference: DPS becomes 3.2$\times$ faster with 50% reduction in peak GPU memory. FK steering with StitchVM significantly outperforms standard FK steering and Best-of-$N$ sampling, particularly by enabling efficient local proposal selection.
  • Training: DiffusionNFT becomes 2.3$\times$ faster, and direct reward finetuning (DRaFT) becomes 1.3$\times$ faster while achieving higher generation quality (e.g., +30% GenEval) due to effective supervision at high-noise steps.
• Generalization: The approach works across diverse diffusion backbones (SD 3.5 Medium/Large, FLUX) and reward models (CLIP, DFN-CLIP, HPSv2, Aesthetic Predictor).

4. Experimental Results

The paper evaluates StitchVM across multiple benchmarks:

• Value Function Quality: On noisy latents with varying noise levels ($\sigma \in [0.1, 0.9]$), StitchVM closely matches the performance of the original clean reward models in zero-shot retrieval (MSCOCO, Flickr30K), preference prediction (HPDv2, ImageReward), and aesthetic scoring (AVA). It significantly outperforms baselines like VAE-stitching (which collapses at high noise) and models trained from scratch on noisy data (e.g., NoisyCLIP, DiNa-LRM).
• Inference Alignment: In DPS and FK steering experiments on DrawBench, StitchVM-enhanced methods consistently improve metrics like ImageReward, HPSv2, and GenEval while drastically reducing computational cost. Notably, increasing the number of local proposals ($M$) in FK steering using StitchVM proves more efficient than increasing the number of particles ($N$).
• Training Alignment: Training curves show that StitchVM variants reach higher final scores with fewer GPU-hours compared to standard DRaFT and DiffusionNFT. Qualitative results indicate that StitchVM-augmented models converge faster and produce sharper, higher-quality images earlier in training.

5. Significance and Claims

The paper claims that StitchVM provides a generic template for combining latent diffusion backbones with pixel-space feedforward models without sacrificing the extensive pretraining of either. By lifting powerful reward models to latent space, it enables a new style of diffusion alignment where the correct value function for noisy latents is constructed once and amortized over many samples and iterations.

The authors position StitchVM not as a replacement for foundation model pretraining, but as a practical, low-cost mechanism to adapt these models to specific tasks. They emphasize that this approach makes value-model-based alignment accessible and efficient, potentially accelerating research into safer and more reliable generative modeling. The work suggests that the "stitching" of pre-trained components is a viable alternative to the prohibitive costs of training value functions from scratch or relying on biased approximations.