Apple skúma, ako skrotiť problematické tokeny v difúznych transformeroch

Výskumníci z Apple a Rice University sa zamerali na menej viditeľný problém difúznych transformátorov používaných na generovanie obrázkov: malú skupinu tokenov s neobvykle vysokou normou. Takéto tokeny môžu pri výpočte pozornosti získať neprimeraný vplyv, hoci nesú iba obmedzenú informáciu o konkrétnej časti obrazu. Práca Taming Outlier Tokens in Diffusion Transformers tvrdí, že tento jav nevzniká iba vo vizuálnych transformeroch určených na porozumenie obrazu, ale aj v moderných generatívnych pipeline založených na difúznych transformeroch.

Autori skúmali architektúru, v ktorej reprezentáciu vstupu vytvára predtrénovaný enkodér typu Vision Transformer a samotné generovanie zabezpečuje difúzny transformer, označovaný ako DiT. Ide o usporiadanie spájané s reprezentančnými autoenkodérmi, teda RAE. Výskum identifikoval dva možné zdroje problému: odľahlé reprezentácie môžu prichádzať už z enkodéra a ďalšie vysokonormové tokeny sa môžu vytvoriť vo vnútri denoizéra, najmä v jeho stredných vrstvách. Nejde teda nevyhnutne o jednu izolovanú chybu na vstupe modelu.

Vysoká norma tokenu znamená, že jeho vektorová reprezentácia má výrazne väčšiu veľkosť než väčšina ostatných. V mechanizme pozornosti môže takýto token priťahovať neprimerane veľkú časť výpočtovej pozornosti modelu. Problémom nie je iba samotná číselná extrémnosť. Token môže súčasne strácať väzbu na lokálny obrazový obsah, ktorý mal reprezentovať. Dôsledkom môže byť poškodenie priestorovej informácie a vznik artefaktov počas postupného odšumovania obrazu.

Jednoduché odstránenie alebo maskovanie tokenov s najvyššou normou podľa autorov kvalitu nezlepšilo. Tento negatívny výsledok je dôležitý, pretože ukazuje, že riešením nemusí byť obyčajný prahový filter. Extrémna hodnota je skôr príznakom hlbšieho narušenia lokálnej sémantiky obrazových segmentov. Ak model problematický token už využíva v ďalších vrstvách alebo sa okolo neho vytvorili závislosti, jeho mechanické vymazanie môže odstrániť aj užitočnú informáciu bez toho, aby napravilo pôvodnú reprezentáciu.

Navrhnutá metóda Dual-Stage Registers preto zasahuje do oboch hlavných stupňov pipeline. Pre enkodér využíva trénované registre, ak ich architektúra podporuje. V ostatných prípadoch autori zavádzajú rekurzívne registre aplikované počas inferencie. Druhú časť tvoria difúzne registre určené pre denoizér. Register možno chápať ako pomocný token, ktorý modelu poskytne miesto na zachytenie globálnych alebo nešpecifických informácií, aby sa nemuseli hromadiť v tokenoch reprezentujúcich konkrétne časti obrazu.

Dvojstupňový prístup je podstatný aj z praktického hľadiska. Úprava iba enkodéra by nemusela zachytiť tokeny, ktoré vzniknú až pri denoizácii, zatiaľ čo zásah iba do DiT by ponechal poškodené reprezentácie prichádzajúce zo vstupnej časti systému. DSR sa snaží tieto dva prípady oddeliť a riešiť na mieste ich vzniku. Rekurzívny variant navyše naznačuje možnosť pracovať aj s predtrénovanými enkodérmi, pri ktorých nie sú k dispozícii pôvodné trénované registre alebo nie je praktické celý model znovu trénovať.

Autori uvádzajú, že zásahy konzistentne obmedzili artefakty spojené s odľahlými tokenmi a zlepšili kvalitu generovania na ImageNete aj pri rozsiahlejšom generovaní obrázkov z textu. Verejný abstrakt však neposkytuje dostatok údajov na samostatné posúdenie veľkosti zlepšenia pri každej architektúre, rozlíšení či výpočtovom rozpočte. Z výsledku preto zatiaľ nemožno vyvodiť, že registre automaticky zlepšia každý difúzny model alebo že ich prínos preváži dodatočnú zložitosť vo všetkých produkčných nasadeniach.

Práca je relevantná pre vývojárov generatívnych modelov, pretože presúva časť pozornosti od bežného zväčšovania modelov k vnútornému toku reprezentácií. Ak malý počet tokenov dokáže skresliť pozornosť naprieč vrstvami, lepšie riadenie týchto tokenov môže zvýšiť kvalitu bez potreby zásadne väčšieho modelu. Potenciálnym prínosom je aj stabilnejšia tvorba detailov a lepšie zachovanie lokálnej štruktúry, hoci konkrétny dopad na rýchlosť, pamäť a náklady musí byť vyhodnotený pri každej implementácii.

Výsledky zároveň nadväzujú na pozorovania z klasických vizuálnych transformátorov, kde sa registre používajú na oddelenie globálnych výpočtových informácií od tokenov obrazových polí. Novinkou je prenesenie tejto otázky do generatívnej pipeline a rozlíšenie problémov v enkodéri a denoizéri. To môže byť užitočné aj diagnosticky: pri obrazových artefaktoch nebude stačiť sledovať iba finálny výstup, ale aj normy tokenov, rozdelenie pozornosti a lokálnu sémantiku v rôznych hĺbkach modelu.

Opatrnosť je naďalej namieste. Preprint bol na arXiv odoslaný 6. mája 2026 a v čase zverejnenia je označený ako práca v recenznom konaní. Zistenia preto ešte nemajú váhu nezávisle potvrdeného štandardu a bude dôležité overiť ich na ďalších dátových súboroch, architektúrach a tréningových režimoch. Štúdia napriek tomu ponúka konkrétnu hypotézu aj technický zásah: kvalitu difúznych transformátorov môže ovplyvňovať nielen počet parametrov či tréningové dáta, ale aj to, kam model ukladá informácie, ktoré nepatria žiadnemu konkrétnemu obrazovému segmentu.