PaperlayerView exact revision on arXiv

Source-linked AI summary

Stitched Value Model for Diffusion Alignment

Hyojun Go, Hyungjin Chung, Prune Truong, Goutam Bhat, Li Mi, Zhaochong An, Zixiang Zhao, Dominik Narnhofer, Serge Belongie, Federico Tombari, Konrad Schindler

arXiv:2605.19804v1cs.CVcs.AIcs.LG

TL;DR

Diffusion alignment needs accurate value estimates for noisy latents, but existing approximations trade bias against expensive computation. StitchVM stitches pretrained clean-image reward models to diffusion backbones, improving efficiency across alignment methods while maintaining or improving quality.

  • Problem

    Strong, general value models for noisy latents remain difficult to train, although alignment requires such estimates for rewards defined on clean images.

  • Method

    StitchVM stitches a truncated frozen diffusion backbone to a sliced pretrained reward model through a lightweight layer, then lightly finetunes the hybrid.

  • Results

    Across inference- and training-time alignment, StitchVM improves efficiency while maintaining or improving quality; DPS becomes 3.2× faster while halving peak GPU memory.

  • Takeaways & Limitations

    StitchVM enables direct evaluation of noisy-latent value functions in place of expensive per-sample approximations while preserving or improving alignment quality.

  • Takeaways & Limitations

    Directly training noisy-latent value models requires substantial data and compute and must be repeated for each new diffusion backbone or reward model.

Abstract

from arXiv · show

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.

1. Introduction

Diffusion alignment methods need repeated value estimates for noisy latents, but clean-image rewards make existing Tweedie and Monte Carlo approximations biased or expensive. StitchVM instead stitches a truncated diffusion backbone to a pretrained reward-model tail, transferring clean-image reward capabilities into an efficient noisy-latent value model.

  • Motivation: Alignment methods repeatedly assess noisy latents z_t with a value function measuring expected rewards of the clean samples they induce.This requirement spans both training-time and inference-time alignment methods.
  • Limitations of Existing Estimators: Tweedie approximation can be biased at high noise and requires extra denoiser and VAE-decoder evaluations, whereas Monte Carlo approximation incurs costly denoising rollouts.Both methods are workarounds because rewards are normally defined for clean images x_0.
  • StitchVM: StitchVM combines a truncated frozen diffusion backbone as a noisy-latent-capable head with a sliced pretrained reward model as a reward-capable tail through a lightweight stitching layer.The framework transfers pretrained reward-model capabilities into the noisy latent regime with only a small finetuning cost.
  • Results: Stitched models retain the benchmark performance of their underlying clean reward models while directly ingesting noisy latents across multiple diffusion backbones and reward models.The tested backbones include SD 3.5 Medium, SD 3.5 Large, and FLUX; reward models include DFN-CLIP, CLIP, Aesthetic Score Predictor, and HPSv2.
  • Results: ≈10 hours on a single GH200 GPU is sufficient to transfer ViT-L/14@336px CLIP into an SD 3.5 Medium value function.The stitching uses only a few unlabeled images and lightweight finetuning.
  • Downstream Alignment: 3.2× faster and halving peak GPU memory are achieved when the stitched value model replaces DPS’s long gradient paths with direct value-model gradients.The low-cost value estimator also supports more efficient local proposal selection in FK steering.

2. Related Work

Prior diffusion-alignment methods approximate value functions on noisy latents using Tweedie or Monte Carlo methods, while newer approaches learn noisy-latent value models directly. Related training and stitching work motivates combining noise awareness with broader reward generalization through hybrid models.

  • Alignment methods and value function: Tweedie and Monte Carlo approximations evaluate pixel-level rewards indirectly on noisy latents for guidance, particle weighting, or related alignment procedures.Tweedie uses an estimated clean sample, whereas Monte Carlo evaluates noisy-latent values through sampling.
  • Alignment methods and value function: Learned noisy-latent value models support credit assignment, high-noise reward feedback, and lower rollout cost, but typically generalize narrowly because training uses limited preference data or task-specific labels.These models have been applied to PPO-style post-training, direct reward finetuning, and search-based inference.
  • Training value models and noisy latent reward models: Diffusion-feature predictors are naturally noise-aware, yet their prediction heads commonly lack broad generalization because they are trained on narrow preference data.Existing approaches attach prediction heads or LLM interfaces to diffusion features.
  • Model stitching: Model stitching recombines early layers from one network with later layers from another, enabling hybrid models across architectures and applications including resource-constrained reassembly.Recent generative-model applications stitch 3D reconstruction networks onto clean latents; the paper extends this idea to noisy latents.

3. Preliminary

The preliminary section formulates flow matching in latent space and defines reward-based alignment through a terminal soft value function. It highlights that practical value estimation remains difficult because Tweedie approximations are biased, whereas Monte Carlo estimates are costly, high-variance, or unstable.

  • Flow matching formulation: Flow matching operates in latent space, interpolating from clean latents at t = 0 to a reference Gaussian at t = 1.The Gaussian conditional path uses α_t = 1 − t and σ_t = t, inducing a marginal path between the two distributions.
  • Flow matching formulation: The conditional path induces a marginal velocity field that can be used for sampling through ODEs, SDEs, or discrete transition kernels.For flow matching, the conditional velocity is u_t(z_t|z_0) = ϵ − z_0.
  • Reward-based alignment: Alignment seeks samples that score highly under task-specific rewards, including prompt alignment, aesthetics, human preference, and physical consistency.The reward-tilted target distribution is defined relative to the pretrained base prior, with partition functions for pixel and latent spaces.
  • Value estimation challenges: Tweedie guidance approximates the value as r(E[z_0|z_t]), but this introduces Jensen-gap bias.The approximation is used because the exact value function is intractable.
  • Value estimation challenges: Monte Carlo value estimates avoid the former bias but have high variance and require substantial computation, while RL rollouts are unstable and high-variance.These limitations affect particle-based inference and KL-regularized reinforcement-learning post-training.

4. Methodology

StitchVM stitches diffusion-backbone features to pretrained clean-image reward models, then lightly finetunes the interface to estimate values directly for noisy latents. This enables cheaper inference-time steering and training-time alignment without full denoising rollouts.

  • Stitching framework: StitchVM aligns diffusion features with a pretrained reward model’s feature space through a lightweight stitching layer, combining noisy-latent processing with precise clean-image rewards.The diffusion backbone is truncated at one layer, while the reward model begins at another selected layer.
  • Stage 1: Selecting the stitching interface: The method exhaustively searches candidate layer pairs by closed-form linear feature matching, then selects the pair with the lowest matching loss.Paired clean and noisy-latent features are extracted for each candidate interface before fitting the linear mapping.
  • Stage 2: Finetuning StitchVM: StitchVM finetunes the stitching layer and subsequent reward-model components on noisy latents, using original reward scores from clean images as supervision for the standard value function.The procedure uses unlabeled clean images, forward-process noise sampling, an off-policy objective, and the standard rather than soft value function for training stability.
  • Inference-time alignment: StitchVM evaluates values directly in noisy latent space, avoiding Tweedie’s denoiser and decoder evaluations and long backpropagation chains for more accurate, efficient steering.For FK steering, a single forward pass makes value estimates cheaper and allows more local proposals per particle without substantially increasing compute.
  • Training-time alignment: StitchVM enables training to stop at intermediate noisy latents while supervising with their value estimates, improving and accelerating direct reward finetuning and DiffusionNFT.This avoids full denoising rollouts and long gradient chains while retaining reward-weighted forward-process regression for DiffusionNFT.

5. Experiments

Across noisy-latent evaluation, StitchVM preserves pretrained reward-model capability, outperforms alternative transfer and retraining baselines, and improves inference- and training-time alignment efficiency. Its lightweight transfer procedure costs ≈10 GPU-hours at 512×512 resolution and scales to multiple diffusion backbones and reward models.

  • 5.1 Main results: At low noise (σ≤0.5), StitchVM closely matches clean reward models across CLIP retrieval, HPSv2 preference prediction, and aesthetic prediction, while remaining robust as noise increases.Performance declines gradually with noise, but remains substantially more robust than the baselines.
  • 5.1 Main results: Diffusion-feature stitching substantially outperforms VAE stitching across all noise levels, especially at high noise where the VAE baseline collapses.The comparison attributes this robustness to diffusion features trained to process noisy latents rather than only clean latent space.
  • 5.1 Main results: StitchVM outperforms NoisyCLIP despite NoisyCLIP’s larger-scale LAION-400M training and outperforms DiNa-LRM on HPSv2 using unlabeled images instead of preference data.These results favor transferring pretrained reward models through stitching over retraining on noisy latents or preference data.
  • 5.2 Inference-time alignment: 33% lower cost is achieved by FK steering with (N=8, M=6), matching standard FKS at N=14 and showing that local proposal scaling is more efficient than increasing N alone.Across the compute range, the StitchVM curve lies above the standard N-scaling curve.
  • 5.3 Training-time alignment: 22–26% lower GPU-hours for DRaFT and over 55% lower GPU-hours for DiffusionNFT are obtained with StitchVM by stopping rollouts at intermediate noisy latents and evaluating the value function directly.StitchVM also improves DRaFT generation quality through direct value-function supervision.
  • Training cost: ≈10 GPU-hours is the StitchVM training cost at 512×512 resolution, increasing to 24–32 GPU-hours at 1024×1024 for stitching-layer search and subsequent finetuning.The procedure is reported as a lightweight, one-time transfer compared with large-scale reward-model retraining.

6. Conclusion … A.3. Reinforcement Post-training of Diffusion

StitchVM transfers pretrained pixel-space reward models into value functions for noisy diffusion or flow latents, enabling efficient direct evaluation while preserving alignment quality. The appendices justify the sampling, reparametrization, gradient-guidance, and reinforcement-learning foundations underlying this approach.

  • 6. Conclusion: StitchVM stitches a frozen diffusion backbone to a pretrained pixel-space reward model, then applies lightweight self-supervised finetuning to construct noisy-latent value functions.The framework combines the backbone’s noisy-latent handling with the reward model’s pretrained capability.
  • 6. Conclusion: Direct StitchVM evaluation replaces expensive per-sample approximations, improving efficiency while maintaining or improving alignment quality.The method is presented as a generic way to combine latent diffusion backbones with pixel-space feedforward models without sacrificing either model’s pretraining.
  • A. Proofs: Flow-based models can be sampled with either ODE or SDE procedures, with the SDE formulation expressed using the score function and a diffusion coefficient controlling injected noise.For the stated FM schedule, a natural diffusion-coefficient choice preserves the probability-flow ODE marginals.
  • A.1. Reparametrization in Diffusion and Flow-based Models: The velocity field in flow matching can be reparameterized to recover both the posterior-mean denoiser and the score function through Tweedie’s formula.The appendix specializes these relations to the FM schedule α_t = 1 − t and σ_t = t.
  • A.2. From score reparametrization to gradient guidance: The score of the reward-tilted marginal equals the pretrained score shifted by the gradient of the value function at each intermediate latent.This connects the tilted distribution to the value function defined over noisy states.
  • A.2. From score reparametrization to gradient guidance: Gradient guidance samples the reward-tilted distribution by adding a drift term proportional to ∇z_tV_t(z_t) to the pretrained sampling SDE.Following the log-value gradient steers trajectories toward high-reward regions in the manner required by the tilted target.
  • A.3. Reinforcement Post-training of Diffusion: KL-regularized reinforcement learning over the discrete denoising MDP has the reward-tilted distribution as its unique optimum.The reward is defined at the final denoising step, and optimizing the objective is equivalent to minimizing reverse KL to the tilted target.

A.4. Off-policy Value Model Training · B. Extended Related Work · C. Additional Methodological Details

The off-policy objective learns the posterior-mean reward at each noisy latent, while this conditional-mean value remains useful for gradient guidance and particle methods. Related work addresses noisy-latent value estimation through Tweedie approximations, Monte Carlo rollouts, and learned intermediate-reward models.

  • A.4. Off-policy Value Model Training: The standard and soft value functions coincide in the noiseless limit t→0, when the posterior concentrates on the corresponding clean latent.At finite noise, Jensen’s inequality gives equality only when the reward is constant on the posterior support.
  • A.4. Off-policy Value Model Training: To leading order in reward scale, gradient guidance using the regressed conditional-mean value samples from the reward-tilted distribution, with temperature absorbed into the guidance coefficient.The temperature is λ, and the implicit scale is c_t.
  • A.4. Off-policy Value Model Training: O(λ^2) corrections distinguish the standard and conditional-mean soft values in particle methods, controlled by the conditional reward variance σ_t^2.The correction is largest in the high-noise regime, where σ_t^2 is largest, while the conditional mean retains the dominant ranking signal in experiments.
  • B. Extended Related Work: Inference-time alignment approximates a soft optimal denoising policy because noisy latents cannot directly receive rewards defined in clean pixel space.Tweedie methods estimate a one-step clean image, decode it through the VAE, and evaluate the pixel-space reward.
  • B. Extended Related Work: Training-time RL alignment commonly estimates trajectory values with Monte Carlo rollouts, including terminal-reward backpropagation, PPO-style policy gradients, and DiffusionNFT.These methods use sampled or complete denoising trajectories to obtain reward signals.
  • B. Extended Related Work: Learned value and noisy-latent reward models reduce rollout costs or improve credit assignment by evaluating intermediate latents and providing per-step denoising feedback.LatSearch avoids fully rolling out every candidate trajectory, while PPO-style post-training uses intermediate feedback.

C.1. StitchVM Training … D.2. Inference-Time Alignment Experiments

StitchVM selects and trains a lightweight bridge from frozen diffusion features to pretrained reward-model representations, then applies the resulting noisy-latent value estimates across inference-time steering and reward-finetuning methods. The experiments specify implementation choices for stitching, baselines, and alignment procedures.

  • C.1. StitchVM Training: StitchVM ranks diffusion/reward layer pairs using a closed-form feature-matching loss estimated from 200 held-out clean images with uniformly sampled noise levels.The probe set caches paired features for N_probe = 200 clean images, drawing t∼Unif[0, 1] along the forward process.
  • C.1. StitchVM Training: Restricting the reward-model search to early blocks avoids sharply higher fitting loss and substantially worse predictions after value-loss finetuning.The selected pair (i★, j★) is therefore searched within the early-block region.
  • C.2. FK steering with StitchVM: FK steering with StitchVM preserves standard FK weighting and resampling, while each proposal-scaling step scores M local transition-kernel proposals and keeps the highest-scoring one.The modification is confined to proposal-scaling steps; standard FK potential and resampling remain unchanged.
  • C.3. AlignProp & DRaFT with StitchVM: AlignProp and DRaFT replace terminal-reward backpropagation through the full denoising chain with StitchVM evaluation at a stopped intermediate latent.Stopping at τ provides supervision in high-noise regions and avoids complete reverse rollouts, reducing computation.
  • C.4. DiffusionNFT with StitchVM: DiffusionNFT retains weighted forward regression but replaces the clean terminal reward with V(i★,j★)_ω(z_τ) and anchors regression at z_τ.The intermediate value estimate supplies a training signal without completing the reverse process to z_0.
  • D.1. Stitched Value Model Experiments: The experiments fit the stitching layer and reward suffix while freezing the diffusion backbone, using a 1×10^-5 base learning rate and a 5× stitching-layer multiplier.Layer search uses N_probe = 200 HPDv2 images, restricts reward-side candidates to j≤4, and sweeps all diffusion DiT-block indices.
  • D.1. Stitched Value Model Experiments: The VAE-stitching baseline uses the same noisy-latent training recipe but omits diffusion DiT features, making its conceptual difference from StitchVM the front-end representation.It replaces the diffusion head with the identity on the noisy VAE latent and finetunes for 15 epochs.
  • D.2. Inference-Time Alignment Experiments: Inference-time evaluations cover DPS and FK steering across SD 3.5 Medium, SD 3.5 Large, and FLUX.1-dev at 1024 × 1024 resolution, with reward and compositional-alignment benchmarks.DPS uses 100 denoising steps, while FK steering uses generator-specific default configurations and separate early high-noise proposal scaling with StitchVM.

D.3. Training-Time Alignment Experiments · E. Additional Experimental Results

The training-time alignment experiments use SD3.5 Medium with joint DFN-CLIP and HPSv2 rewards, evaluating direct reward finetuning, DiffusionNFT, and Flow-GRPO-Fast under shared compute and sampling protocols. StitchVM variants replace terminal clean-image rewards with predictions on noisy latents, while final evaluation uses original clean-image reward models.

  • D.3. Training-Time Alignment Experiments: SD3.5 Medium is trained on HPDv2 prompts with an equal-weight sum of DFN-CLIP and HPSv2 rewards.The shared setup applies to direct reward finetuning and DiffusionNFT.
  • D.3. Training-Time Alignment Experiments: StitchVM variants replace the terminal clean reward with StitchVM predictions, while standard variants evaluate rewards on decoded clean images.The shorthand is R(z0) = R(D(z0)).
  • D.3. Training-Time Alignment Experiments: 16 GPUs are used for all training runs at 512×512 resolution, arranged as 4 nodes × 4 NVIDIA GH200 GPUs.This shared setup covers the training-time alignment experiments.
  • D.3. Training-Time Alignment Experiments: Direct reward finetuning uses LoRA with rank r = 32 and α = 64 on all attention projections of the joint MM-DiT blocks.AdamW uses weight decay 0, learning rate 5×10^-5, gradient clipping at norm 1.0, and EMA on trainable parameters.
  • D.3. Training-Time Alignment Experiments: DiffusionNFT follows its official multi-reward setup, replacing the standard rewards with the joint DFN-CLIP and HPSv2 reward.The passage specifies this as an exception to the official implementation protocol.
  • D.3. Training-Time Alignment Experiments: Flow-GRPO-Fast follows its official PickScore setup but uses the joint DFN-CLIP and HPSv2 reward on the same 16-GPU configuration.This setup is part of the training-time alignment experiments.
  • D.3. Training-Time Alignment Experiments: 40 denoising steps at cfg = 1.0 are used for sample generation, and all methods are evaluated on fully denoised samples with original clean-image reward models.GPU-hours are reported as wall-clock training cost, with timings measured on the same 16-GH200 layout.

E.1. Full Numerical Results of StitchVM Performance · E.2. Training Curves in Training-Time Alignment

Table 5 provides the full numerical results for StitchVM across retrieval, preference, and aesthetic evaluations, while Figures 4–5 show its training-time efficiency benefits for DRaFT and DiffusionNFT. During DRaFT finetuning, StitchVM reaches higher final scores with substantially less compute across metrics.

  • E.1. Full Numerical Results of StitchVM Performance: Table 5 reports full numerical results corresponding to Figure 2’s line plots.The table covers three evaluation settings: zero-shot image-text retrieval, preference accuracy, and aesthetic score correlation.
  • E.1. Full Numerical Results of StitchVM Performance: Zero-shot retrieval is evaluated on MSCOCO and Flickr30K using Recall@1 in Image→Text and Text→Image directions.The evaluated models are CLIP ViT-L/14 and DFN-CLIP.
  • E.1. Full Numerical Results of StitchVM Performance: Preference accuracy is evaluated on HPDv2 and ImageReward for HPSv2.Table 5 also evaluates aesthetic score correlation on the AVA test split using SR.
  • E.2. Training Curves in Training-Time Alignment: Figure 5 reports DiffusionNFT and DiffusionNFT with StitchVM training curves, showing StitchVM accelerates DiffusionNFT training.The comparison uses the same five plotted metrics against GPU-hours.
  • E.2. Training Curves in Training-Time Alignment: Figures 4 and 5 plot GenEval, HPSv2, DFN-CLIP, ImageReward, and PickScore against GPU-hours during joint DFN-CLIP and HPSv2 reward finetuning.The curves assess both training-reward metrics and held-out metrics for DRaFT and DiffusionNFT.
  • E.2. Training Curves in Training-Time Alignment: DRaFT with StitchVM reaches higher final scores with substantially less compute across metrics, improving the quality–efficiency trade-off.Figure 4 compares DRaFT and DRaFT with StitchVM during finetuning on the joint reward.

E.3. Analysis of Stitching Interface Search · E.4. Cross-Backbone Generalization of StitchVM · E.5. Stopping-Step Distribution for RL Finetuning with StitchVM

StitchVM’s closed-form interface loss efficiently excludes catastrophic stitch points, while cross-backbone transfer preserves performance with lower value-model cost. For RL finetuning, intermediate stopping-step windows outperform high-noise stopping, with Unif{12, . . . , 17} offering the best quality–efficiency trade-off.

  • E.3. Analysis of Stitching Interface Search: The closed-form feature-matching loss provides a cheap protocol for sweeping diffusion-block and CLIP-block stitch interfaces before end-to-end training.The sweep fits W★ in closed form and compares the loss with trained-model MSCOCO Recall@1.
  • E.3. Analysis of Stitching Interface Search: Once the CLIP cut moves beyond early blocks, the loss rises by roughly an order of magnitude and Recall@1 falls from around 49 to below 5.Stage-2 finetuning cannot recover from these poor interfaces.
  • E.3. Analysis of Stitching Interface Search: Within the low-loss region j≤4, the lowest-loss interface (i, j) = (4, 1) achieves Recall@1 of 48.3, about 1.5 points below the within-region maximum.The method restricts the reward-model cut to early CLIP blocks and selects the lowest-loss cell, enabling a much cheaper search.
  • E.4. Cross-Backbone Generalization of StitchVM: Across fifteen reward–metric cells, an SD 3.5 Medium StitchVM closely matches an SD 3.5 Large StitchVM when steering SD 3.5 Large.Their HPSv2 results differ by at most 0.002, while other metrics remain similar.
  • E.4. Cross-Backbone Generalization of StitchVM: 0.72 vs. 0.70: under HPSv2 reward, the SD 3.5 Medium StitchVM achieves higher GenEval than the SD 3.5 Large StitchVM.The backbones share the same SD3 16-channel VAE, making their noisy latents dimensionally compatible at every noise level.
  • E.4. Cross-Backbone Generalization of StitchVM: Using the smaller SD 3.5 Medium backbone reduces per-step value-model cost because StitchVM uses only the diffusion backbone’s early blocks.This supports guiding a larger generator with a StitchVM trained on a smaller backbone.
  • E.5. Stopping-Step Distribution for RL Finetuning with StitchVM: Unif{2, . . . , 12} consistently underperforms across GenEval, HPSv2, DFN-CLIP, ImageReward, and PickScore, indicating that stopping too early yields less useful value targets.Smaller step indices represent earlier, higher-noise latents; larger indices represent later, cleaner latents.
  • E.5. Stopping-Step Distribution for RL Finetuning with StitchVM: Unif{12, . . . , 17} provides the best quality–efficiency trade-off, converges faster than Unif{12, . . . , 25}, and is used by default for DiffusionNFT and DRaFT.Unif{20, . . . , 25} remains competitive on some metrics, such as PickScore, but is less stable overall.

E.6. Qualitative Results on RL Finetuning with StitchVM

Qualitative comparisons show that adding StitchVM improves RL-finetuning outcomes across DRaFT-1, DRaFT-3, and DiffusionNFT, enabling earlier target-prompt achievement and higher visual quality throughout training.

  • Qualitative Results: Across DRaFT-1, DRaFT-3, and DiffusionNFT, StitchVM-augmented variants reach the target prompt earlier and produce visually higher-quality samples throughout training.Examples include sharper details and more saturated colors characteristic of HPSv2-tuned outputs.
  • Qualitative Results: Figure 10 qualitatively compares DiffusionNFT with and without StitchVM across training GPU-hours.

F. Limitation · G. Broader Impacts

StitchVM is limited to feedforward-model-based rewards, while timestep-aware training is identified as a promising future improvement. Its benefits may improve controllability and alignment, but could also facilitate deceptive or harmful reward optimization.

  • F. Limitation: Limitation: StitchVM does not directly apply to rewards that are not implemented as feedforward models.The authors suggest training surrogate reward models as a possible remedy but leave this direction for future work.
  • F. Limitation: Future direction: Timestep-aware training methods are proposed as a promising way to improve StitchVM.The work focuses on a simple method rather than more complex alternatives that may further improve performance.
  • G. Broader Impacts: Positive impact: StitchVM makes training noisy latent value models from existing pretrained reward models substantially cheaper.This can improve and accelerate reward-based alignment methods for diffusion and flow models.
  • G. Broader Impacts: Positive impact: The framework may make generative models more controllable, more aligned with human preferences, and easier to adapt to downstream tasks.These benefits are especially relevant where clean-image reward models are already available.
  • G. Broader Impacts: Negative impact: Stronger alignment and steering methods could generate more persuasive synthetic images, including misleading or deceptive visual content.The authors note that these risks are not unique to StitchVM.
  • G. Broader Impacts: Negative impact: StitchVM may lower the cost of optimizing diffusion-based generation toward poorly specified or harmful objectives.The same improvements in controllability and reward optimization may therefore be misused.
Loading 2605.19804v1…