DiScoFormer od Ai2 odhaduje hustotu aj skóre rozdelenia jedným transformerom

Ai2 zverejnilo cez Hugging Face výskumný článok o modeli DiScoFormer, ktorý sa pozerá na menej viditeľnú, ale prakticky dôležitú vrstvu moderného strojového učenia: ako z konečnej vzorky bodov odhadnúť celé pravdepodobnostné rozdelenie. Namiesto toho, aby model riešil iba jednu úzko definovanú dátovú sadu, má zobrať množinu bodov ako kontext a v jednom prechode odhadnúť dve veličiny naraz — hustotu rozdelenia a takzvané skóre, teda smer najrýchlejšieho rastu logaritmu hustoty.

Pre nešpecialistu to znie abstraktne, no ide o základnú operáciu za veľkou časťou dnešnej AI. Hustota hovorí, kde sú dáta pravdepodobné a kde sú zriedkavé. Skóre zasa ukazuje, ktorým smerom sa má bod posunúť, aby sa dostal do pravdepodobnejšej oblasti. Práve tento princíp používajú difúzne generatívne modely, ktoré začínajú od šumu a postupne ho posúvajú k realistickému obrázku, zvuku alebo inému výstupu. Podobná matematika sa objavuje aj v Bayesovskom vzorkovaní a vo vedeckých simuláciách, napríklad pri modelovaní fyzikálnych systémov.

Klasická metóda, s ktorou sa DiScoFormer porovnáva, je kernel density estimation. Tá pri novom bode sleduje, koľko trénovacích bodov je v jeho okolí a ako ďaleko sú. Je intuitívna a nepotrebuje tréning, ale vo vysokých dimenziách rýchlo naráža na pamäťové a presnostné limity. Ai2 tvrdí, že transformerová pozornosť sa dá chápať ako všeobecnejší variant tejto klasickej metódy: jeden attention head sa správa podobne ako adaptívne jadro a viac vrstiev vie zachytiť viac škál naraz.

Architektúra DiScoFormeru preto nie je len ďalší veľký model nasadený na starý problém. Model používa spoločnú chrbticu a dve výstupné hlavy, jednu pre hustotu a druhú pre skóre. Keďže skóre má byť gradientom logaritmu hustoty, medzi oboma výstupmi existuje matematická väzba. Autori ju používajú aj ako konzistenčnú stratu bez dodatočných anotácií: pri inferencii môže model niekoľkými krokmi doladiť výstup tak, aby sa hustotná a skórová hlava navzájom lepšie zhodovali.

Tréning je postavený na zmesiach Gaussových rozdelení. Tie sú praktické, pretože vedia aproximovať širokú triedu hladkých rozdelení a zároveň poskytujú presné cieľové hodnoty hustoty aj skóre. Model teda môže vidieť prakticky neobmedzené množstvo syntetických distribúcií a učiť sa všeobecný postup, nie len memorovať jednu dátovú tabuľku. To je dôležité najmä vtedy, ak má byť výsledný model opakovane použiteľný v rôznych vedeckých a generatívnych úlohách.

Podľa blogu Ai2 DiScoFormer prekonáva ručne ladené klasické odhady najmä tam, kde tradičné metódy slabnú. V storočlenných priestoroch má podľa autorov znížiť chybu skóre približne 6,5-násobne a chybu hustoty viac než 37-násobne oproti najlepšie nastavenej kernelovej metóde v ich experimentoch. Zároveň sa má správať lepšie, keď pribúdajú vzorky, zatiaľ čo klasické riešenie naráža na pamäť.

Najzaujímavejší dopad však nie je samotné číslo v benchmarku, ale možnosť mať predtrénovaný modul na odhad skóre a hustoty, ktorý sa nemusí nanovo trénovať pre každý problém. Ak sa tento smer potvrdí aj mimo syntetických zmesí a výskumných testov, môže znížiť cenu experimentovania pri difúznych modeloch, Bayesovských metódach a simuláciách, kde sa dnes často budujú špecializované odhady od nuly.

Zároveň nejde o okamžitý produktový prelom. Ai2 samo priznáva, že klasické metódy môžu byť rýchlejšie na menších dátach a že generalizácia mimo trénované rodiny rozdelení je presne tá časť, ktorú bude treba ďalej overovať. Výskumný význam DiScoFormeru je skôr v tom, že ukazuje most medzi starou neparametrickou štatistikou a transformerovou pozornosťou. Ak sa z takýchto mostov stanú spoľahlivé stavebné bloky, časť infraštruktúry generatívnej AI sa môže posunúť od ad hoc trénovania k znovupoužiteľným štatistickým komponentom.

Pre slovenské firmy a výskumné tímy je na tom zaujímavý aj širší trend: open-source ekosystém sa už nezaoberá iba veľkými jazykovými modelmi, ale aj špecializovanými komponentmi, ktoré môžu zlepšiť numerické a vedecké workflow. Ak bude DiScoFormer dostupný ako overiteľný nástroj s jasnými limitmi, môže byť užitočný aj tam, kde sa dnes generatívna AI nepoužíva na text, ale na urýchlenie odhadu pravdepodobnostných štruktúr v dátach.