Apple zrýchľuje generovanie obrázkov na štyri kroky bez opustenia presnej pravdepodobnosti

Výskumníci z Apple, University of Pennsylvania a University of Illinois Urbana-Champaign predstavili metódu Normalizing Trajectory Models, skrátene NTM. Rieši známy problém generatívnych obrazových modelov: klasické difúzne modely dosahujú vysokú kvalitu postupným odstraňovaním šumu, no pri generovaní musia často vykonať veľa po sebe nasledujúcich výpočtov. NTM sa nesnaží iba preskočiť časť týchto krokov. Každý spätný prechod modeluje ako flexibilný podmienený normalizačný tok a zachováva pritom presne vyčísliteľnú vierohodnosť generatívnej trajektórie. Apple uvádza, že takto vytvorený model na benchmarkoch generovania obrazu z textu dorovnáva alebo prekonáva silné referenčné metódy už pri štyroch krokoch vzorkovania. Ide však o výsledok vedeckej práce, nie o oznámenie hotového produktu.

Prečo je málo krokov náročný problém

Difúzny model sa počas učenia oboznamuje s postupným pridávaním šumu do dát a pri generovaní sa učí tento proces obrátiť. Veľký počet malých zmien je výpočtovo priaznivý, pretože každý jednotlivý prechod sa dá približovať pomerne jednoduchým rozdelením. Nevýhodou je sekvenčnosť: ďalší krok závisí od predchádzajúceho, takže ani výkonný akcelerátor nedokáže všetku prácu spracovať naraz. Práca DDIM už ukázala, že vhodne navrhnutý nedeterministický alebo implicitný proces môže generovanie výrazne zrýchliť a pritom používať rovnaký tréningový cieľ ako pôvodné difúzne modely. Pri extrémnom stlačení na niekoľko hrubých prechodov však jednoduché gaussovské aproximácie prestávajú dostatočne vystihovať, kam sa má vzorka posunúť.

Doterajšie systémy tento rozdiel často riešia destiláciou, consistency tréningom alebo adversariálnym cieľom. Silný model či dlhá trajektória vytvoria učebný signál pre rýchlejší model, ktorý sa pokúša dosiahnuť podobný výsledok v malom počte krokov. Takéto postupy môžu byť prakticky úspešné, no podľa autorov NTM pri nich generovanie stráca časť čistého pravdepodobnostného rámca. Model potom môže vytvárať presvedčivé obrázky, ale nemá zároveň presnú vierohodnosť nad celou trajektóriou. To komplikuje analýzu, porovnávanie aj niektoré formy ďalšieho učenia. NTM je pokus spojiť obe vlastnosti: hrubé a expresívne kroky potrebné na nízku latenciu s explicitným modelom pravdepodobnosti.

Čo mení Normalizing Trajectory Model

Jadrom návrhu je podmienený normalizačný tok v každom spätnom kroku. Normalizačné toky používajú invertibilné transformácie, pri ktorých možno sledovať zmenu hustoty pravdepodobnosti. Jeden prechod preto nemusí byť obmedzený na jednoduchú gaussovskú zmenu a môže opísať zložitejšie rozdelenie možných nasledujúcich stavov. Architektúra podľa Apple kombinuje plytké invertibilné bloky vo vnútri jednotlivých krokov s hlbokým paralelným prediktorom, ktorý pracuje naprieč trajektóriou. Celú sieť možno trénovať od začiatku alebo inicializovať z modelu využívajúceho flow matching. Druhá možnosť je dôležitá prakticky: vývojár nemusí nevyhnutne zahodiť všetku investíciu do existujúceho generatívneho základu.

Flow matching poskytuje širší kontext, prečo je tento smer zaujímavý. Pôvodná práca o flow matchingu formulovala učenie spojitých normalizačných tokov ako regresiu vektorových polí nad zvolenými pravdepodobnostnými cestami. Tie môžu zahŕňať difúzne cesty, ale aj iné interpolácie, napríklad založené na optimálnom transporte. NTM túto líniu neposúva iba ďalším spojitým tokom. Sústreďuje sa na celú krátku diskrétnu trajektóriu a každý jej veľký krok vybavuje vlastnou expresívnejšou distribúciou. Výsledkom má byť model, ktorý zvládne veľké skoky medzi šumom a obrazom bez toho, aby sa pravdepodobnostný opis zmenil na nevyčísliteľnú aproximáciu.

Presná vierohodnosť trajektórie má aj druhé využitie. Autori z nej odvodzujú skórovú funkciu a používajú ju na samodestiláciu ľahšieho denoisera. Tento denoiser má následne vytvárať kvalitné vzorky v štyroch krokoch. Samodestilácia znamená, že učiteľský signál pochádza zo samotného NTM, nie z oddeleného proprietárneho modelu. V princípe to môže zjednodušiť tréningový reťazec a lepšie prepojiť rýchly generátor s rozdelením, ktoré sa pôvodný model skutočne naučil. Verejný opis však zatiaľ neposkytuje dostatok detailov na posúdenie nákladov tejto fázy ani toho, ako citlivý je výsledok na architektúru a dátový súbor.

Praktický význam rýchlejšieho vzorkovania

Ak sa výsledky potvrdia mimo experimentov autorov, štyri kroky by mohli znížiť latenciu interaktívnych obrazových nástrojov a počet výpočtov potrebných na jednu vzorku. To je relevantné pre návrhárske aplikácie, generovanie viacerých variantov, úpravy obrazu v reálnom čase aj nasadenie na zariadeniach s obmedzeným výkonom. Nižší počet krokov však automaticky neznamená štvornásobne či desaťnásobne nižšie náklady. Jeden krok NTM je expresívnejší než jednoduchý denoisingový krok a obsahuje invertibilné bloky, preto treba porovnávať reálny čas, pamäť, energetickú spotrebu a dávkovú priepustnosť. Samotný počet vyhodnotení siete je užitočná, ale nie postačujúca prevádzková metrika.

Pre menšie európske a slovenské tímy je podstatná najmä možnosť stavať rýchle generovanie na už natrénovanom flow-matching modeli. Ak bude metóda reprodukovateľná a implementácia dostupná, mohla by obmedziť potrebu trénovať úplne nový základný model. Potenciálny prínos sa týka aj lokálneho spracovania citlivých vizuálnych dát, kde odosielanie obrázkov do cudzej služby nie je prijateľné. Zatiaľ však nemožno tvrdiť, že NTM je pripravený na bežné lokálne nasadenie. Apple na výskumnej stránke neoznamuje verejné váhy, produkčné API ani podporovaný hardvér a neuvádza výsledky merané na spotrebiteľských zariadeniach.

Čo výsledky zatiaľ nedokazujú

Oznámenie hovorí o dorovnaní alebo prekonaní silných obrazových základov pri štyroch krokoch, ale z krátkeho verejného súhrnu nie je možné posúdiť všetky experimentálne podmienky. Dôležité budú konkrétne dátové súbory, rozlíšenie, veľkosť modelu, výpočtový rozpočet tréningu a metriky kvality. Automatické skóre navyše nemusí zachytiť čitateľnosť textu v obraze, presné počítanie objektov, zložité priestorové pokyny ani konzistentnosť identity. Nezávislá reprodukcia by mala porovnať NTM nielen s klasickou difúziou, ale aj so súčasnými destilovanými a flow-matching systémami pri rovnakej latencii a podobnom počte parametrov.

Otvorená zostáva aj všeobecnosť prístupu. Apple opisuje výsledky pre generovanie obrazu z textu, no princíp trajektórie by mohol byť relevantný aj pre video, zvuk alebo iné modality. Pri videu sa však dramaticky zvyšujú pamäťové nároky a nároky na časovú konzistentnosť, takže úspech na statických obrázkoch nemožno automaticky preniesť. NTM je preto najlepšie čítať ako nový výskumný kompromis medzi rýchlosťou, expresivitou a pravdepodobnostnou kontrolou. Jeho význam sa ukáže až po zverejnení úplných experimentov, implementácie a nezávislých meraní na reálnom hardvéri.