Modelo de Valor Cosido para la Alineación de Difusión

Resumen

Para uso práctico, los modelos generativos basados en difusión o flujo deben alinearse con recompensas específicas de la tarea, como la fidelidad al prompt o la preferencia estética. Dicha alineación es un desafío porque la recompensa se define para imágenes de salida limpias, pero el procedimiento de alineación requiere estimaciones de la función de valor en latentes intermedios ruidosos. Los métodos existentes recurren a aproximaciones de estilo Tweedie o de Monte Carlo, equilibrando el sesgo del estimador con el coste computacional: las estimaciones de Tweedie son eficientes pero sesgadas, mientras que las de Monte Carlo son más precisas pero requieren simulaciones costosas. Una alternativa natural sería una función de valor aprendida, pero sigue siendo una cuestión abierta cómo entrenar eficazmente un modelo de valor robusto y general específicamente para latentes ruidosos. Aquí proponemos StitchVM, un marco de cosido de modelos que transfiere eficientemente modelos de recompensa preentrenados para imágenes limpias al régimen de latentes ruidosos. StitchVM parte de un modelo de recompensa truncado en el espacio de píxeles y acopla un backbone de difusión congelado como su cabeza. Del modelo en espacio de píxeles, el híbrido resultante conserva una capacidad de recompensa robusta y cuidadosamente preentrenada; del backbone de difusión, hereda su capacidad nativa para manejar latentes ruidosos. El procedimiento de cosido es excepcionalmente ligero; por ejemplo, coser y ajustar CLIP ViT-L y SD 3.5 Medium solo requiere 10 horas de GPU. Al elevar los potentes modelos de recompensa del espacio de píxeles al espacio latente, StitchVM abre un nuevo estilo de alineación de difusión: en lugar de una aproximación aproximada pero costosa por muestra de la función de valor, la función correcta para los latentes ruidosos reales se construye una vez y luego se amortiza en múltiples muestras e iteraciones. Demostramos que este enfoque produce mejoras en un amplio rango de métodos de guía y post-entrenamiento posteriores: DPS se vuelve 3,2 veces más rápido mientras reduce a la mitad la memoria máxima de GPU, y DiffusionNFT se vuelve 2,3 veces más 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.