Gestikt Waardemodel voor Diffusie-uitlijning

Samenvatting

Voor praktisch gebruik moeten diffusie- of stromingsgebaseerde generatieve modellen worden afgestemd op taakspecifieke beloningen, zoals promptgetrouwheid of esthetische voorkeur. Die afstemming is uitdagend omdat de beloning is gedefinieerd voor schone uitvoerbeelden, maar de afstemmingsprocedure vereist waardefunctieschattingen op ruizige tussentijdse latenten. Bestaande methoden gebruiken Tweedie-achtige of Monte Carlo-benaderingen, waarbij een afweging wordt gemaakt tussen schattingsbias en computationele kosten: Tweedie-schattingen zijn efficiënt maar vertekend, terwijl Monte Carlo-schattingen nauwkeuriger zijn maar dure rollouts vereisen. Een natuurlijk alternatief zou een geleerde waardefunctie zijn, maar het blijft een open vraag hoe effectief een sterk en algemeen waardemodel te trainen dat specifiek is voor ruizige latenten. Hier stellen we StitchVM voor, een modelstitching-framework dat beloningsmodellen die zijn voorgetraind voor schone beelden efficiënt overbrengt naar het ruizige latentenregime. StitchVM vertrekt van een bestaand, afgekapt pixelruimte-beloningsmodel en bevestigt er een bevroren diffusie-backbone aan als hoofd. Van het pixelruimtemodel behoudt de resulterende hybride een zorgvuldig voorgetraind, robuust beloningsvermogen; van de diffusie-backbone erft het zijn natuurlijke vermogen om met ruizige latenten om te gaan. De stitchprocedure is uitzonderlijk lichtgewicht; het stitchen en finetunen van CLIP ViT-L en SD 3.5 Medium duurt bijvoorbeeld slechts 10 GPU-uren. Door krachtige pixelruimte-beloningsmodellen over te brengen naar de latente ruimte, opent StitchVM een nieuwe stijl van diffusie-afstemming: in plaats van een ruwe, maar kostbare per-sample benadering van de waardefunctie, wordt de juiste functie voor de daadwerkelijke, ruizige latenten eenmalig geconstrueerd en vervolgens geamortiseerd over vele monsters en iteraties. We tonen aan dat deze aanpak verbeteringen oplevert over een breed scala aan downstream stuur- en natrainingsmethoden: DPS wordt 3,2 keer sneller terwijl het piek-GPU-geheugen halveert, en DiffusionNFT wordt 2,3 keer sneller.

English

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.2times faster while halving peak GPU memory, and DiffusionNFT becomes 2.3times faster.