Сшитая модель ценности для согласования диффузионных моделей

Аннотация

Для практического использования диффузионные или основанные на потоке генеративные модели должны быть согласованы с целевыми вознаграждениями, такими как точность соответствия запросу или эстетические предпочтения. Такое согласование представляет сложность, поскольку вознаграждение определяется для чистых выходных изображений, но процедура согласования требует оценок функции ценности на зашумленных промежуточных скрытых представлениях. Существующие методы прибегают к приближениям стиля Твиди или Монте-Карло, балансируя смещение оценки с вычислительными затратами: оценки Твиди эффективны, но смещены, тогда как оценки Монте-Карло более точны, но требуют дорогостоящих прогонов. Естественной альтернативой могла бы быть обучаемая функция ценности, однако остается открытым вопрос, как эффективно обучить сильную и общую модель ценности, предназначенную именно для зашумленных скрытых представлений. Здесь мы предлагаем StitchVM — фреймворк сшивания моделей, который эффективно переносит модели вознаграждения, предварительно обученные на чистых изображениях, в режим зашумленных скрытых представлений. StitchVM начинается с существующей усеченной модели вознаграждения в пиксельном пространстве и присоединяет к ней в качестве головы замороженную магистраль диффузии. От модели в пиксельном пространстве полученный гибрид наследует тщательно предварительно обученную устойчивую способность к вознаграждению; от магистрали диффузии он наследует ее естественную способность обрабатывать зашумленные скрытые представления. Процедура сшивания исключительно легковесна: например, сшивание и дообучение CLIP ViT-L и SD 3.5 Medium занимает всего 10 часов на GPU. Поднимая мощные модели вознаграждения из пиксельного пространства в пространство скрытых представлений, StitchVM открывает новый стиль согласования диффузии: вместо грубого, но дорогого пообразного приближения функции ценности, правильная функция для реальных зашумленных скрытых представлений строится один раз, а затем амортизируется на множество выборок и итераций. Мы показываем, что этот подход дает улучшения в широком спектре методов последующего управления и пост-обучения: DPS становится в 3,2 раза быстрее, при этом пиковое использование памяти GPU сокращается вдвое, а DiffusionNFT — в 2,3 раза быстрее.

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.