Het afstemmen van tekst-naar-beeld diffusiemodellen met beloningsbackpropagatie

Aligning Text-to-Image Diffusion Models with Reward Backpropagation

October 5, 2023

Auteurs: Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki

cs.AI

Samenvatting

Text-to-image diffusiemodellen zijn recent naar voren gekomen als toonaangevend in beeldgeneratie, aangedreven door zeer grootschalige, onbewaakte of zwak begeleide tekst-naar-beeld trainingsdatasets. Vanwege hun onbewaakte training is het moeilijk om hun gedrag te beheersen in downstream taken, zoals het maximaliseren van door mensen waargenomen beeldkwaliteit, beeld-tekstuitlijning of ethische beeldgeneratie. Recente werken finetunen diffusiemodellen naar downstream beloningsfuncties met behulp van standaard reinforcement learning, berucht vanwege de hoge variantie van de gradientenschatters. In dit artikel stellen we AlignProp voor, een methode die diffusiemodellen uitlijnt met downstream beloningsfuncties door middel van end-to-end backpropagatie van de beloningsgradiënt door het denoisingsproces. Hoewel een naïeve implementatie van dergelijke backpropagatie een onhaalbaar geheugengebruik zou vereisen voor het opslaan van de partiële afgeleiden van moderne tekst-naar-beeldmodellen, finetunt AlignProp laag-rang adaptergewichtmodules en gebruikt het gradient checkpointing om het geheugengebruik haalbaar te maken. We testen AlignProp in het finetunen van diffusiemodellen naar verschillende doelen, zoals beeld-tekst semantische uitlijning, esthetiek, comprimeerbaarheid en beheersbaarheid van het aantal aanwezige objecten, evenals hun combinaties. We laten zien dat AlignProp hogere beloningen behaalt in minder trainingsstappen dan alternatieven, terwijl het conceptueel eenvoudiger is, waardoor het een rechttoe rechtaan keuze is voor het optimaliseren van diffusiemodellen voor differentieerbare beloningsfuncties van belang. Code en visualisatieresultaten zijn beschikbaar op 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/.