Unconditional Priors Matter! Improving Conditional Generation of Fine-Tuned Diffusion Models

Technische Samenvatting: Unconditional Priors Matter!

1. Het Probleem: Degradatie van Unconditional Priors

Bij het trainen van conditionele diffusiemodellen (zoals modellen voor tekst-naar-beeld, beeldbewerking of video-generatie) wordt vaak gebruikgemaakt van Classifier-Free Guidance (CFG). De standaardpraktijk is om één enkel netwerk te trainen dat zowel de conditionele ruis ($\epsilon_\theta(x_t, c)$) als de onconditionele ruis ($\epsilon_\theta(x_t, \emptyset)$) voorspelt. Dit gebeurt door de conditionering tijdens de training met een lage waarschijnlijkheid (bijv. 5-20%) weg te laten (dropout).

De auteurs stellen vast dat dit proces een fundamenteel probleem heeft: wanneer een basismodel (zoals Stable Diffusion) wordt fine-tuned voor een specifieke taak, raakt de "unconditional prior" (het vermogen van het model om realistische beelden te genereren zonder specifieke instructies) ernstig aangetast. Omdat het netwerk beperkte capaciteit heeft en de onconditionele data tijdens de training minder vaak voorkomt, verliest het model zijn rijke kennis over de algemene beeldstructuur. Dit leidt tot een paradox: hoewel het model beter wordt in de specifieke taak, zorgt de slechte onconditionele ruisvoorspelling voor een verslechtering van de uiteindelijke kwaliteit van de conditionele generatie via de CFG-formule.

2. Methodologie: Noise Replacement

De kernoplossing van het paper is een eenvoudige, training-vrije methode om de onconditionele ruis te vervangen door een kwalitatief betere bron.

In plaats van de onconditionele ruis van het fine-tuned model te gebruiken in de CFG-stap, stelt de methode voor om de onconditionele ruis van een ander, krachtig pre-trained model te gebruiken. De aangepaste CFG-formule wordt als volgt:

$$\epsilon^{(\gamma)}_{\theta, \psi}(x_t, c) = \epsilon_\psi(x_t, \emptyset) + \gamma(\epsilon_\theta(x_t, c) - \epsilon_\psi(x_t, \emptyset))$$

Hierbij is:

• $\epsilon_\theta$: De conditionele ruis voorspeld door het fine-tuned model.
• $\epsilon_\psi$: De onconditionele ruis voorspeld door een basismodel (of een ander model met sterke priors, zoals SD2.1 of PixArt-$\alpha$).
• $\gamma$: De guidance scale.

Belangrijke ontdekking: Het model $\psi$ hoeft niet noodzakelijkerwijs het originele basismodel te zijn waarop het fine-tuned model is gebouwd. Zelfs modellen met een andere architectuur (bijv. een DiT in plaats van een UNet) kunnen effectief dienen als bron voor de onconditionele prior.

3. Belangrijkste Bijdragen

• Identificatie van een nieuw fenomeen: Het aantonen dat de degradatie van onconditionele priors tijdens fine-tuning een directe oorzaak is van verminderde kwaliteit in conditionele generatie.
• Training-vrije oplossing: Een methode die direct kan worden toegepast op bestaande modellen zonder dat er extra training of architecturale wijzigingen nodig zijn.
• Architectuur-agnostisch: De methode werkt ongeacht of de modellen UNet- of Transformer-gebaseerd (DiT) zijn, en werkt zowel voor afbeeldingen als voor video.
• Cross-modale effectiviteit: Het laat zien dat onconditionele ruis van verschillende modellen effectief gecombineerd kan worden om de generatiekwaliteit te verbeteren.

4. Resultaten

De auteurs hebben de methode getest op diverse state-of-the-art modellen:

• Zero-1-to-3 (Novel View Synthesis): Verbeterde de nauwkeurigheid van belichting en vorm, wat resulteerde in betere PSNR en SSIM scores.
• Versatile Diffusion (Image Variations): Verminderde kleurverzadiging en objectvervormingen, met significante verbeteringen in FID en DINOv2 scores.
• DiT (Class-conditional): Toonde aan dat de methode ook werkt voor Transformer-gebaseerde modellen, waarbij de FID-score verbeterde bij grotere datasets.
• DynamiCrafter (Image-to-Video): Resulteerde in video's met een hogere esthetische kwaliteit en betere temporele consistentie (minder flikkering en vervorming).
• InstructPix2Pix (Image Editing): Verbeterde de trouw aan de tekstuele instructie terwijl de identiteit van het originele beeld beter behouden bleef.

5. Betekenis en Conclusie

Dit onderzoek is significant omdat het een praktische en efficiënte manier biedt om de prestaties van gespecialiseerde AI-modellen te maximaliseren zonder de enorme kosten van extra training. Het verschuift de focus van "hoe trainen we een model om alles te leren" naar "hoe combineren we de specifieke kennis van een model met de algemene kennis van een ander model tijdens de inferentie".

Hoewel de methode een iets hoger geheugengebruik en een minimale toename in inferentietijd met zich meebrengt (omdat er twee modellen in het geheugen moeten worden geladen), wegen de grote kwaliteitsverbeteringen in beeld- en videogeneratie hier ruimschoots tegenop.