Alignement des modèles de diffusion texte-image par rétropropagation de récompense
Aligning Text-to-Image Diffusion Models with Reward Backpropagation
October 5, 2023
papers.authors: Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki
cs.AI
papers.abstract
Les modèles de diffusion texte-image ont récemment émergé à l'avant-garde de la génération d'images, alimentés par des ensembles de données d'entraînement texte-image non supervisés ou faiblement supervisés à très grande échelle. En raison de leur apprentissage non supervisé, contrôler leur comportement dans des tâches en aval, comme maximiser la qualité perçue des images par les humains, l'alignement texte-image ou la génération éthique d'images, est difficile. Des travaux récents affinent les modèles de diffusion pour des fonctions de récompense en aval en utilisant l'apprentissage par renforcement classique, connu pour la forte variance de ses estimateurs de gradient. Dans cet article, nous proposons AlignProp, une méthode qui aligne les modèles de diffusion sur des fonctions de récompense en aval en utilisant la rétropropagation end-to-end du gradient de récompense à travers le processus de débruitage. Bien qu'une implémentation naïve d'une telle rétropropagation nécessiterait des ressources mémoire prohibitives pour stocker les dérivées partielles des modèles texte-image modernes, AlignProp affine des modules de poids d'adaptation de faible rang et utilise la vérification de gradient pour rendre son utilisation de la mémoire viable. Nous testons AlignProp dans l'affinage de modèles de diffusion pour divers objectifs, tels que l'alignement sémantique texte-image, l'esthétique, la compressibilité et la contrôlabilité du nombre d'objets présents, ainsi que leurs combinaisons. Nous montrons qu'AlignProp atteint des récompenses plus élevées en moins d'étapes d'entraînement que les alternatives, tout en étant conceptuellement plus simple, ce qui en fait un choix direct pour optimiser les modèles de diffusion pour des fonctions de récompense différentielles d'intérêt. Le code et les résultats de visualisation sont disponibles à l'adresse https://align-prop.github.io/.
English
Text-to-image diffusion models have recently emerged at the forefront of
image generation, powered by very large-scale unsupervised or weakly supervised
text-to-image training datasets. Due to their unsupervised training,
controlling their behavior in downstream tasks, such as maximizing
human-perceived image quality, image-text alignment, or ethical image
generation, is difficult. Recent works finetune diffusion models to downstream
reward functions using vanilla reinforcement learning, notorious for the high
variance of the gradient estimators. In this paper, we propose AlignProp, a
method that aligns diffusion models to downstream reward functions using
end-to-end backpropagation of the reward gradient through the denoising
process. While naive implementation of such backpropagation would require
prohibitive memory resources for storing the partial derivatives of modern
text-to-image models, AlignProp finetunes low-rank adapter weight modules and
uses gradient checkpointing, to render its memory usage viable. We test
AlignProp in finetuning diffusion models to various objectives, such as
image-text semantic alignment, aesthetics, compressibility and controllability
of the number of objects present, as well as their combinations. We show
AlignProp achieves higher rewards in fewer training steps than alternatives,
while being conceptually simpler, making it a straightforward choice for
optimizing diffusion models for differentiable reward functions of interest.
Code and Visualization results are available at https://align-prop.github.io/.