Schaalbaarheid van inferentie bij diffusiemodellen: een nieuw paradigma

Inleiding tot de schaalbaarheid van inferentie bij diffusiemodellen

Recente ontwikkelingen in Large Language Models (LLMs) hebben de effectiviteit van schaalvergroting tijdens inferentie aangetoond. Modellen zoals o1, o3, DeepSeek R1, QwQ en Step Reasoner mini hebben laten zien dat een toename van de rekenkracht tijdens inferentie de prestaties aanzienlijk kan verbeteren. Dit roept de vraag op: kan dit principe ook worden toegepast op diffusiemodellen? Een team onder leiding van Xie Saining aan de New York University heeft deze vraag onderzocht. Door middel van een systematisch onderzoek met behulp van een algemeen zoekframework, ontdekten ze dat inferentie-time schaling inderdaad effectief is voor diffusiemodellen. Het verhogen van de rekenkracht tijdens inferentie leidt tot een aanzienlijke verbetering van de kwaliteit van gegenereerde samples. Bovendien maakt de complexiteit van afbeeldingen verschillende combinaties van componenten binnen het framework mogelijk, afgestemd op verschillende toepassingsscenario's.

Belangrijkste bevindingen van inferentie-time schaling

Inferentie-time schaling is effectief voor diffusiemodellen: Het toewijzen van meer computationele middelen tijdens inferentie resulteert in samples van hogere kwaliteit.
Flexibiliteit in componentcombinaties: Het framework maakt verschillende componentconfiguraties mogelijk, geschikt voor diverse toepassingen.
Verder dan ruisverwijderingsstappen: Het onderzoek suggereert dat het zoeken naar betere ruis tijdens sampling een andere dimensie is voor het schalen van NFE, in plaats van simpelweg het verhogen van ruisverwijderingsstappen.
Twee ontwerpassen: Het framework richt zich op twee belangrijke ontwerpassen:
- Verifiers: Het geven van feedback tijdens het zoekproces.
- Algoritmen: Het vinden van betere ruiskandidaten.

Onderzoeksmethodologie voor de diffusiemodellen

Het team onderzocht drie verschillende scenario's voor verifiers, die verschillende use-cases simuleerden:

Scenario's waarin bevoorrechte informatie over de uiteindelijke evaluatie beschikbaar is.
Scenario's waarin voorwaardelijke informatie beschikbaar is om de generatie te begeleiden.
Scenario's zonder aanvullende informatie beschikbaar.

Voor algoritmen onderzochten ze:

Random Search: Het selecteren van de beste uit een vaste set kandidaten.
Zero-Order Search: Het iteratief verbeteren van ruiskandidaten met behulp van verifier feedback.
Path Search: Het iteratief verbeteren van diffusie sampling trajecten met behulp van verifier feedback.

De studie onderzocht deze ontwerpen aanvankelijk in een relatief eenvoudige ImageNet klasse-conditionele generatie setup. Vervolgens pasten ze deze ontwerpen toe op grootschalige tekst-conditionele generatie en evalueerden ze hun voorgestelde framework.

Het schalen van inferentietijd bij diffusiemodellen

Het paper stelt een framework voor voor het schalen van inferentietijd in diffusiemodellen, waarbij de uitdaging wordt gezien als een zoektocht naar optimale sampling ruis. Het proces omvat twee kerncomponenten:

Verifiers: Dit zijn vooraf getrainde modellen die de kwaliteit van gegenereerde samples beoordelen. Ze nemen gegenereerde samples in, samen met optionele voorwaarden, en geven een scalaire score.
Algoritmen: Deze algoritmen gebruiken de verifier scores om betere kandidaat samples te vinden. De functie neemt een verifier (V), een vooraf getraind diffusiemodel (_θ) en een set gegenereerde samples en voorwaarden in, en geeft de beste initiële ruis.

Het totale inferentiebudget wordt gemeten door het totale aantal functie-evaluaties (NFE), inclusief zowel ruisverwijderingsstappen als zoekkosten.

Zoekverifiers en hun rol bij diffusiemodellen

De onderzoekers begonnen met een Oracle verifier, die volledige informatie heeft over de uiteindelijke evaluatie van geselecteerde samples. Voor ImageNet omvatte dit statistieken zoals FID en IS. Vervolgens onderzochten ze meer toegankelijke vooraf getrainde modellen als supervised verifiers, zoals CLIP en DINO. Deze modellen werden gebruikt om samples te classificeren, waarbij de sample met de hoogste logit die overeenkomt met het klassenlabel werd geselecteerd.

Ze observeerden echter dat deze classificatoren, die puntgewijs werken, slechts gedeeltelijk overeenkomen met de doelstellingen van de FID-score. Dit leidde tot een vermindering van sample variantie en mode collapse naarmate de berekening toenam. Dit fenomeen, dat 'verifier hacking' wordt genoemd, werd versneld door de onbeperkte zoekruimte van het random search algoritme. Interessant is dat de studie ontdekte dat verifiers niet noodzakelijkerwijs voorwaardelijke informatie nodig hebben om de zoektocht effectief te begeleiden. Ze observeerden een sterke correlatie tussen de logits van DINO/CLIP classificatoren en de cosinus similariteit van de feature space tussen de x voorspelling op een laag ruisniveau en de uiteindelijke schone sample. Dit leidde tot het gebruik van zelf-supervised verifiers, die geen aanvullende voorwaardelijke informatie nodig hebben en toch effectief schaalgedrag vertoonden.

Zoekalgoritmen voor betere sampling bij diffusiemodellen

Om verifier hacking te beperken, onderzochten de onderzoekers meer verfijnde zoekalgoritmen die kandidaat samples geleidelijk optimaliseren. Dit omvatte een zero-order zoekmethode:

Beginnen met een willekeurige Gaussiaanse ruis als een draaipunt.
Het vinden van N kandidaten in de buurt van het draaipunt.
Het uitvoeren van kandidaten via de ODE solver om samples en verifier scores te verkrijgen.
Het bijwerken van het draaipunt met de beste kandidaat en het herhalen van stappen 1-3.

Ze onderzochten ook een path search algoritme, dat de mogelijkheid onderzoekt om langs het sampling traject te zoeken:

Het samplen van N initiële ruis samples en het uitvoeren van de ODE solver naar een ruisniveau σ.
Het toevoegen van ruis aan elke sample en het simuleren van een forward noising proces.
Het uitvoeren van een ODE solver op elke noisy sample en het behouden van de top N kandidaten op basis van verifier scores, dit herhalen totdat de ODE solver σ=0 bereikt.
Willekeurig zoeken naar de overige N samples en de beste behouden.

Zowel zero-order als path search algoritmen behouden een sterke localiteit in vergelijking met random search.

Schalen in text-to-image scenario's bij diffusiemodellen

Het team onderzocht de schaalmogelijkheden van het zoekframework in grootschalige text-to-image taken. Ze gebruikten de DrawBench en T2I-CompBench datasets voor evaluatie, met het FLUX.1-dev model als basis. Ze breidden ook de selectie van supervised verifiers uit, waaronder Aesthetic Score Predictor, CLIPScore en ImageReward. Daarnaast creëerden ze een Verifier Ensemble door deze drie verifiers te combineren.

Analyse: Verifier-taak afstemming bij diffusiemodellen

De studie vergeleek de resultaten van verschillende verifier-algoritme combinaties op verschillende datasets. Op DrawBench ontdekten ze dat het gebruik van alle verifiers de sample kwaliteit over het algemeen verbeterde. Ze observeerden echter dat het gebruik van Aesthetic en CLIP verifiers in isolatie kon leiden tot overaanpassing van hun biases, wat resulteerde in negatieve effecten op elkaar. Dit komt voort uit een mismatch in hun evaluatiefocus: Aesthetic Score richt zich op visuele kwaliteit, waarbij vaak sterk gestileerde afbeeldingen de voorkeur krijgen, terwijl CLIP visueel-tekst afstemming prioriteit geeft, soms ten koste van de visuele kwaliteit. Ze merkten op dat sommige verifiers meer geschikt zijn voor specifieke taken, en de effectiviteit van een verifier hangt af van de afstemming met de taakeisen.

Algoritme prestaties bij diffusiemodellen

De drie zoekalgoritmen (Random, Zero-Order en Path) verbeterden allemaal effectief de sampling kwaliteit op DrawBench. Random Search presteerde echter in sommige aspecten beter vanwege de lokale aard van de andere twee methoden. De random search converteerde sneller naar verifier bias, terwijl de andere twee algoritmen verbetering vereisen op minder dan optimale kandidaten.

Compatibiliteit met fine-tuning bij diffusiemodellen

Het team onderzocht de compatibiliteit van hun zoekmethode met fine-tuned modellen. Ze gebruikten een DPO-fine-tuned Stable Diffusion XL model en ontdekten dat de zoekmethode kon worden gegeneraliseerd naar verschillende modellen en de prestaties van reeds afgestemde modellen kon verbeteren.

Effecten van verschillende dimensies van inferentie berekening bij diffusiemodellen

De studie onderzocht hoe verschillende aspecten van inferentie berekening de resultaten beïnvloeden:

Aantal zoekiteraties: Het verhogen van iteraties brengt de ruis dichter bij het optimum.
Berekening per zoekiteratie: Het aanpassen van het aantal ruisverwijderingsstappen per iteratie onthult verschillende computationeel optimale gebieden.
Finale generatie berekening: Het team gebruikte optimale instellingen voor de uiteindelijke ruisverwijderingsstappen om de hoogste uiteindelijke sample kwaliteit te waarborgen.

Effectiviteit van investeringen in berekeningen bij diffusiemodellen

De onderzoekers onderzochten de effectiviteit van inferentie-time schaling op kleinere diffusiemodellen. Ze ontdekten dat, voor ImageNet, het schalen van kleinere modellen zeer efficiënt kan zijn. In bepaalde gevallen kan zoeken op een kleiner model beter presteren dan grotere modellen zonder zoekopdracht. De effectiviteit hangt echter af van de basisprestaties van het kleinere model. In tekstgebaseerde instellingen presteerde PixArt-Σ, met slechts een fractie van de berekening, beter dan FLUX-1.dev. Deze resultaten tonen aan dat aanzienlijke computationele middelen die tijdens de training worden besteed, kunnen worden gecompenseerd door kleinere hoeveelheden berekening tijdens de generatie, wat resulteert in efficiënter samples van hogere kwaliteit.