Sekvenciu upravujúce modely

Chcem opísať niečo, čo nazývam sekvenciu upravujúci model. Na rozdiel od súčasných jazykových modelov, ktoré generujú ďalší token, hypotetické sekvenciu upravujúce modely generujú celú sekvenciu paralelne, pravdepodobne iteratívnym vylepšovaním návrhov, vďaka čomu sú vhodnejšie pre súčasný hardvér. Difúzne modely sú skorým a nedokonalým príkladom. Potrebujeme však niečo, čo podporuje vkladanie a mazanie. Nižšie opíšem nápad inšpirovaný metódou Gaussian splatting.

Hardvér diktuje architektúru

GPU majú v porovnaní s rýchlosťou pamäte obrovský výpočtový výkon. Dokážu vykonať stovky výpočtových operácií na jednu operáciu prenosu z pamäte. Aby sa tento výpočtový výkon využil, architektúry neurónových sietí sa snažia formulovať väčšinu výpočtov ako násobenie matice maticou, ktoré si vyžaduje O(N³) výpočtov, ale prenesenie len O(N²) dát z pamäte. Pre porovnanie, násobenie matice vektorom má zložitosť O(N²) tak pre výpočty, ako aj pre prenos dát, takže ho nakoniec limituje rýchlosť pamäte.

Dôvod, prečo súčasné jazykové modely vyzerajú tak, ako vyzerajú, má korene v hardvéri. Aby sa trénovali efektívne, architektúra transformátora vyhodnocuje všetky tokeny v sekvencii paralelne. Sekvenčné závislosti sú len na hraniciach vrstiev. Vrstiev je len ~100, takže architektúra je takmer dokonale paralelná. Architektúra transformátora je navrhnutá tak, aby sa dala skompilovať do série efektívnych maticových násobení v GPU. Vďaka tomu je trénovanie modelov transformátora oveľa efektívnejšie napríklad v porovnaní s rekurentnými modelmi.

Transformátory sú však paralelné iba počas trénovania, keď sú všetky tokeny známe vopred. Pri inferencii musia jazykové modely založené na transformátoroch generovať jeden token po druhom, pretože generovanie tokenov vyžaduje znalosť všetkých predchádzajúcich tokenov. Súčasné jazykové modely generujúce ďalší token sú teda pri inferencii extrémne neefektívne, čo má trochu prekvapivo dopad aj na spracovanie promptu.

Pamäť sa nezrýchli. Fyzika obmedzuje latenciu prístupu do vzdialenej pamäte na tretiu odmocninu kapacity pamäte (pozri tiež diskusiu na SE). Zovšeobecnil by som to aj na priepustnosť (bandwidth) pamäte, pretože spotreba energie rastie so vzdialenosťou. Dá sa to trochu oklamať širokou zbernicou do RAM a ešte širšou a kratšou zbernicou do HBM, ale nakoniec narazíte na limity spotreby energie a odvodu tepla. Detaily sa líšia v závislosti od hardvéru, ale aj prístup do L2 cache potrebuje rádovo viac energie než jedno násobenie čísiel s plávajúcou desatinnou čiarkou.

Ak rýchlosť pamäte zostane obmedzená, máme len dve možnosti: buď robiť výpočty tam, kde je pamäť, aby sa prístup do pamäte úplne eliminoval, alebo prejsť na modely s paralelizovateľnou inferenciou, aby sa minimalizoval podiel prenosu dát na celkových výpočtoch. Ľudský mozog má výpočty umiestnené tam, kde je aj pamäť, ale replikovať to v GPU by si vyžadovalo radikálne zmeny hardvéru. Krátkodobo je najlepšie staviť na modely s paralelnou inferenciou, pretože si vyžadujú iba softvérové zmeny.

Difúzne modely sú príliš neflexibilné

Jazykové modely odšumenia (alebo difúzie) na úrovni tokenov, ako napríklad nedávno predvedený Gemini Diffusion, pridajú nejaký druh šumu na úrovni tokenov a potom trénujú model, aby šum odstránil. Tento šum má mnoho podôb: maskovanie tokenov, zmena a preusporiadanie tokenov, šum vo vektorovej reprezentácii (embedding) alebo šum v distribúcii pravdepodobnosti tokenov. Model môže kontrolovať dĺžku sekvencie tak, že nevyužitú časť tokenového buffera vyplní špeciálnymi tokenmi. Modely odšumujúce tokeny sú najznámejšou paralelnou alternatívou k modelom ďalšieho tokenu, tak si preberme ich pozitíva.

Odšumovanie sa zdá byť úplne cudzím prístupom k produkcii jazyka. Žiaden človek takto nerozmýšľa. Funguje to však, pretože text zdieľa určité vlastnosti s obrázkami, odkiaľ odšumovacie modely pochádzajú (slávny Stable Diffusion). Tokeny predpovedajú okolitý text rovnako ako pixely predpovedajú okolitý obraz. Ak zamaskujete 90 % slov v článku, stále viete, o čom článok je. Okrem toho je jazyk elastický ako obrázky. Rovnakú informáciu môžete vtesnať do rôzneho počtu tokenov. To do istej miery zmierňuje potrebu vkladania a mazania. Odšumovací model funguje ako novinár, ktorého požiadali, aby napísal článok dlhý 5 000 slov. Podobne ako novinár, model rozdeľuje tokenový rozpočet sekciám, odsekom a vetám a prispôsobuje celkovú štruktúru aj detailné formulácie tak, aby sa do rozpočtu zmestil.

Ja ale nie som veľkým fanúšikom odšumovacích modelov, pretože odšumovanie na úrovni tokenov sa nevyhnutne bráni vkladaniu a mazaniu. Vloženie alebo odstránenie tokenu by si vyžadovalo posúvanie všetkých nasledujúcich tokenov. Model by to teoreticky mohol urobiť so správne nastaveným trénovaním, ale príde mi to zložité a zatiaľ som nevidel odšumovací model, ktorý by nadšene vykonával vkladanie a mazanie. Skracovanie alebo predlžovanie viet je v novinárskom štýle tolerovateľné, ale neviem si predstaviť, že by jazykové modely generovali zmysluplný zdrojový kód s pevným rozpočtom tokenov.

Inšpirácia metódou Gaussian splatting

Tokeny s pevnou pozíciou, či už v modeloch generujúcich ďalší token alebo v difúznych modeloch, sú podobné voxelom v 3D modelovaní. Podobne ako voxely, tokeny môžu meniť vnútorné vlastnosti (token samotný alebo reprezentáciu za ním), ale ich poloha a veľkosť sú pevné. Gaussian splatting je zaujímavou inšpiráciou, pretože nahrádza rigidné voxely pohyblivými, veľkosť meniacimi a prekrývajúcimi sa škvrnami.

Textovým ekvivalentom Gaussových škvŕn by bolo modelovanie sekvencie ako neusporiadanej množiny nezávislých textových reprezentácií. Každá reprezentácia môže predstavovať rozsah textu od jedného tokenu po celú sekvenciu. Počas inferencie sa môžu všetky reprezentácie v rámci sekvencie pohybovať, meniť veľkosť a meniť svoj vektor. Nepoznám žiadne existujúce implementácie. Toto je len orientačný nápad.

Grafy namiesto sekvencií

Na uľahčenie vkladania a mazania by bolo lepšie použiť na kódovanie pozície v sekvencii adresovanie obsahom (content addressing) namiesto číselnej vsdialenosti. Každá reprezentácia by niesla informáciu o okolitom kontexte, čo by jej implicitne dalo pozíciu v sekvencii voči ostatným reprezentáciám. Adresovanie obsahom by umožnilo reprezentáciám tvoriť orientovaný graf alternatívnych riešení, zatiaľ čo číselné vzdialenosti by reprezentácie uväznili v jednej úplne usporiadanej sekvencii.

Graf reprezentácií by umožnil niektorým reprezentáciám tvoriť tiché lešenie, ktoré poskytuje oporu, ukotvenie a priestor pre experimenty bez priameho vplyvu na výstup. Takéto lešenie by mohlo fungovať ako paralelná alternatíva k premýšľaniu v súčasných modeloch. Reprezentácie pokrývajúce dlhé úseky by sa mohli použiť na začiatku inferencie ako dočasná výplň pre hlavné časti odpovede. Tie by postupne doplnili reprezentácie kratších úsekov alebo jednotlivých tokenov.

Trénovanie a inferencia s pohyblivými reprezentáciami

Túto myšlienku grafu reprezentácií nosím v hlave už roky, ale detaily sú stále hmlisté. Optimalizácia v Gaussian splatting je trochu komplikovaná, pretože škvrny majú tendenciu uviaznuť v lokálnych optimách. Predpokladám, že inferencia s pohyblivými textovými reprezentáciami bude tiež zložitá. Na zložitosti však nezáleží. Záleží na tom, že všetky reprezentácie sa dajú optimalizovať paralelne.

Zatiaľ čo inferencia sa zdá byť len komplikovaná, trénovanie modelu pre pohyblivé reprezentácie sa zdá byť úplne nedosiahnuteľné. Základná myšlienka je, že musí existovať nejaká neurónová sieť, ktorá optimalizuje každú reprezentáciu, pričom berie do úvahy informácie zo všetkých ostatných reprezentácií. Niečo ako transformátor, ale bez pevných pozícií. Ale aký by bol cieľ trénovania? To fakt neviem.

Háčik je v tom, že pri Gaussian splatting sa posúvanie škvŕn do súladu s danou sadou obrázkov implementuje pomocou gradientového zostupu. Textovým ekvivalentom je posúvať výstupné reprezentácie tak, aby zodpovedali danému pevnému vstupu. To znamená, že by sme na inferenciu používali gradientový zostup. Ak sa gradientový zostup používa na inferenciu, ako by vyzeralo trénovanie? Možno tu analógia s Gaussian splatting zlyháva a pohyblivé textové reprezentácie si vyžadujú iný prístup k inferencii a trénovaniu.

Ak by sme najprv vedeli natrénovať samostatný model pre reprezentácie s pevnou pozíciou, potom by bolo ľahšie na ňom natrénovať model s pohyblivými reprezentáciami. Dobrá reprezentácia s pevnou pozíciou reprezentuje tak aktuálny token, ako aj okolitý kontext. Dobrá reprezentácia s pevnou pozíciou je tiež aditívna, takže sa dá skonštruovať ako vážený súčet niekoľkých pohyblivých reprezentácií. Potom by sme mohli nejakým spôsobom porovnať podobnosť medzi očakávanými reprezentáciami s pevnou pozíciou a vygenerovanými pohyblivými reprezentáciami, možno zrekonštruovaním vygenerovaných reprezentácií s pevnou pozíciou z tých pohyblivých. To by nám snáď dalo diferencovateľnú nákladovú funkciu. Stále by sme však museli vymyslieť, ako konštruovať trénovacie vzorky s rôznou mierou šumu.

Každopádne, toto je už nad moje sily. Týmto článkom som chcel povedať, že potrebujeme nejakú architektúru na úpravu sekvencií. Vyššie uvedený nápad s pohyblivými reprezentáciami je len príkladom toho, ako by takáto architektúra mohla vyzerať. Rád by som videl výskum, ktorý sa uberá týmto smerom.