Novel method makes tools like Stable Diffusion and DALL-E-3 faster by simplifying the image-generating process to a single step while maintaining or enhancing image quality.

I vår tidsalder med kunstig intelligens kan datamaskiner generere sin egen “kunst” ved hjelp av diffusjonsmodeller, som iterativt legger til struktur i en støyende utgangstilstand til et tydelig bilde eller en video oppstår. Diffusjonsmodeller har plutselig fått en plass ved alles bord: Skriv inn noen få ord og opplev øyeblikkelige, dopaminstimulerende drømmelandskap i skjæringspunktet mellom virkelighet og fantasi. Bak kulissene er det en kompleks og tidkrevende prosess som krever mange iterasjoner for at algoritmen skal perfeksjonere bildet.

Forskere ved MIT Computer Science and Artificial Intelligence Laboratory (CSAIL) har introdusert et nytt rammeverk som forenkler flerstegsprosessen i tradisjonelle diffusjonsmodeller til ett enkelt trinn, og som dermed fjerner tidligere begrensninger. Dette gjøres ved hjelp av en slags lærer-elev-modell: å lære en ny datamodell å etterligne oppførselen til mer kompliserte, originale modeller som genererer bilder. Tilnærmingen, kjent som DMD (Distribution Matching Distillation), bevarer kvaliteten på de genererte bildene og gjør det mulig å generere dem mye raskere.

“Det vi har utviklet, er en ny metode som gjør dagens diffusjonsmodeller som Stable Diffusion og DALLE-3 30 ganger raskere”, sier Tianwei Yin, doktorgradsstudent i elektroteknikk og informatikk ved MIT, tilknyttet CSAIL og hovedansvarlig for DMD-rammeverket. “Dette fremskrittet reduserer ikke bare beregningstiden betydelig, men opprettholder også, om ikke overgår, kvaliteten på det genererte visuelle innholdet. Teoretisk sett kombinerer tilnærmingen prinsippene for generative adversarial networks (GAN) med prinsippene for diffusjonsmodeller, noe som gjør det mulig å generere visuelt innhold i ett enkelt trinn – i sterk kontrast til de hundre trinnene med iterativ forbedring som dagens diffusjonsmodeller krever. Det kan potensielt bli en ny generativ modelleringsmetode som utmerker seg både når det gjelder hastighet og kvalitet.”

Denne ett-trinns diffusjonsmodellen kan forbedre designverktøyene, gjøre det mulig å skape innhold raskere og potensielt bidra til fremskritt innen legemiddelforskning og 3D-modellering, der hurtighet og effektivitet er avgjørende.

Distribusjonsdrømmer

DMD har to smarte komponenter. Først bruker den et regresjonstap, som forankrer tilordningen for å sikre en grov organisering av bilderommet for å gjøre treningen mer stabil. Deretter bruker den et distribusjonstilpasningstap, som sikrer at sannsynligheten for å generere et gitt bilde med studentmodellen tilsvarer frekvensen i den virkelige verden. For å gjøre dette bruker den to diffusjonsmodeller som fungerer som guider, noe som hjelper systemet med å forstå forskjellen mellom virkelige og genererte bilder og gjør det mulig å trene opp den raske ett-trinnsgeneratoren.

Systemet oppnår raskere generering ved å trene opp et nytt nettverk for å minimere avviket mellom de genererte bildene og bildene fra treningsdatasettet som brukes av tradisjonelle diffusjonsmodeller. “Den viktigste innsikten vår er å tilnærme gradienter som styrer forbedringen av den nye modellen ved hjelp av to diffusjonsmodeller”, sier Yin.

“På denne måten destillerer vi kunnskapen fra den opprinnelige, mer komplekse modellen til den enklere og raskere modellen, samtidig som vi omgår de beryktede problemene med ustabilitet og moduskollaps i GAN-modeller.”

Yin og kollegene brukte forhåndstrente nettverk til den nye studentmodellen, noe som forenklet prosessen. Ved å kopiere og finjustere parametere fra de opprinnelige modellene oppnådde teamet rask konvergens i treningen av den nye modellen, som er i stand til å produsere bilder av høy kvalitet med samme arkitektoniske grunnlag. “Dette gjør det mulig å kombinere den nye modellen med andre systemoptimaliseringer basert på den opprinnelige arkitekturen for å akselerere skapelsesprosessen ytterligere”, legger Yin til.

Når DMD ble satt på prøve mot de vanlige metodene ved hjelp av et bredt spekter av benchmarks, viste DMD konsistent ytelse. I den populære benchmarken for generering av bilder basert på spesifikke klasser på ImageNet er DMD den første ett-trinns diffusjonsteknikken som produserer bilder som er omtrent på nivå med de opprinnelige, mer komplekse modellene, med en supernær Fréchet inception distance (FID)-poengsum på bare 0,3, noe som er imponerende, siden FID handler om å bedømme kvaliteten og mangfoldet i de genererte bildene. I tillegg utmerker DMD seg innen tekst-til-bilde-generering i industriell skala og oppnår toppmoderne ytelse innen ett-trinns generering. Det er fortsatt et lite kvalitetsgap når det gjelder vanskeligere tekst-til-bilde-applikasjoner, noe som tyder på at det er rom for forbedringer på sikt.

I tillegg er ytelsen til de DMD-genererte bildene nært knyttet til egenskapene til lærermodellen som brukes under destillasjonsprosessen. I den nåværende formen, som bruker Stable Diffusion v1.5 som lærermodell, arver studenten begrensninger som å gjengi detaljerte avbildninger av tekst og små ansikter, noe som tyder på at DMD-genererte bilder kan forbedres ytterligere med mer avanserte lærermodeller.

“Å redusere antall iterasjoner har vært den hellige gral for diffusjonsmodeller siden de ble utviklet”, sier Fredo Durand, professor i elektroteknikk og informatikk ved MIT, CSAILs hovedforsker og en av hovedforfatterne av artikkelen. “Vi er veldig glade for endelig å kunne generere bilder i ett trinn, noe som vil redusere beregningskostnadene dramatisk og akselerere prosessen.”

 

“Endelig en artikkel som kombinerer allsidigheten og den høye visuelle kvaliteten til diffusjonsmodeller med sanntidsytelsen til GAN,” sier Alexei Efros, professor i elektroteknikk og informatikk ved University of California i Berkeley, som ikke var involvert i denne studien. “Jeg forventer at dette arbeidet vil åpne for fantastiske muligheter for visuell redigering av høy kvalitet i sanntid.”

 

Yin og Durands medforfattere er professor i elektroteknikk og datavitenskap ved MIT og CSAILs hovedforsker William T. Freeman, samt Adobes forskere Michaël Gharbi SM ’15, PhD ’18, Richard Zhang, Eli Shechtman og Taesung Park. Arbeidet ble delvis støttet av U.S. National Science Foundation (blant annet Institute for Artificial Intelligence and Fundamental Interactions), Singapore Defense Science and Technology Agency, Gwangju Institute of Science and Technology og Amazon. Arbeidet deres vil bli presentert på Conference on Computer Vision and Pattern Recognition i juni.