Modelo de Valor Costurado para Alinhamento de Difusão

Resumo

Para uso prático, modelos generativos baseados em difusão ou fluxo devem ser alinhados com recompensas específicas da tarefa, como fidelidade ao prompt ou preferência estética. Esse alinhamento é desafiador porque a recompensa é definida para imagens de saída limpas, mas o procedimento de alinhamento requer estimativas da função de valor em latentes intermediários ruidosos. Métodos existentes recorrem a aproximações do tipo Tweedie ou de Monte Carlo, equilibrando viés do estimador versus custo computacional: as estimativas de Tweedie são eficientes, mas enviesadas, enquanto as de Monte Carlo são mais precisas, porém exigem *rollouts* dispendiosos. Uma alternativa natural seria uma função de valor aprendida, mas ainda é uma questão em aberto como treinar efetivamente um modelo de valor robusto e geral especificamente para latentes ruidosos. Aqui, propomos o StitchVM, uma estrutura de costura de modelos que transfere eficientemente modelos de recompensa pré-treinados para imagens limpas ao regime de latentes ruidosos. O StitchVM parte de um modelo de recompensa truncado existente no espaço de pixels e anexa um *backbone* de difusão congelado como sua cabeça. Do modelo no espaço de pixels, o híbrido resultante retém uma capacidade de recompensa robusta e cuidadosamente pré-treinada; do *backbone* de difusão, herda sua habilidade nativa de lidar com latentes ruidosos. O procedimento de costura é excepcionalmente leve; por exemplo, costurar e ajustar o CLIP ViT-L e o SD 3.5 *Medium* leva apenas 10 horas de GPU. Ao elevar modelos de recompensa poderosos do espaço de pixels para o espaço latente, o StitchVM abre um novo estilo de alinhamento por difusão: em vez de uma aproximação grosseira, porém custosa por amostra da função de valor, a função correta para os latentes ruidosos reais é construída uma vez e depois amortizada ao longo de muitas amostras e iterações. Mostramos que essa abordagem produz melhorias em uma ampla gama de métodos de orientação e pós-treinamento *downstream*: o DPS se torna 3,2 vezes mais rápido, reduzindo pela metade o pico de memória da GPU, e o DiffusionNFT se torna 2,3 vezes mais rápido.

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.