MemoryLLM od Apple skúma, či sa vrstvy FFN dajú čítať ako pamäť transformera

Apple Machine Learning Research publikoval prácu MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers. Zameriava sa na časť architektúry transformerov, ktorá býva menej viditeľná než attention mechanizmus, no v skutočnosti tvorí veľký podiel parametrov aj výpočtu: feed-forward siete, často označované ako FFN. Autori navrhujú pohľad, v ktorom sa tieto vrstvy skúmajú ako tokenovo prístupná pamäť uložená v parametroch modelu.

Bežné vysvetlenia transformerov zdôrazňujú self-attention, teda mechanizmus, ktorým si tokeny navzájom venujú pozornosť podľa kontextu. FFN vrstva však po každom attention bloku spracúva jednotlivé tokeny nelineárnou transformáciou a významne ovplyvňuje, aké znalosti a schopnosti model prejaví. Problém je, že FFN sú silno previazané so zvyškom architektúry, takže je ťažké povedať, ktoré vstupy otvárajú ktoré „pamäťové“ oblasti a akú hodnotu tieto oblasti majú pre konkrétne úlohy.

MemoryLLM sa pokúša tento uzol rozviazať. Autori trénujú FFN izolovane od self-attention priamo s tokenovými embeddingmi, čím vytvárajú kontextovo nezávislý pohľad na ich správanie. V tomto nastavení sa dá skúmať, ako vstupné tokeny pristupujú k pamäťovým miestam v parametroch FFN a ako dôležité sú tieto miesta pri rôznych downstream úlohách. Nejde teda iba o ďalšiu optimalizáciu výkonu, ale aj o interpretablitný nástroj pre pochopenie toho, čo sa v transformeri deje mimo attention máp.

Z praktického hľadiska je zaujímavá aj výpočtová rovina práce. Keď sa FFN správa ako kontextovo nezávislé tokenové vyhľadávanie, jeho výstupy sa dajú predpočítať ako token-wise lookups. Autori tým otvárajú možnosť presúvať časť pamäťových výpočtov medzi VRAM a pomalším úložiskom podľa potreby. Pri veľkých modeloch, kde je pamäť grafickej karty jedným z hlavných limitov nasadenia, je takýto smer relevantný aj pre inferenčnú infraštruktúru.

Práca zároveň zavádza variant Flex-MemoryLLM, ktorý má stáť medzi konvenčným transformerom a plne oddeleným MemoryLLM. Tento kompromis má znižovať stratu výkonu spôsobenú tým, že FFN sa trénujú s kontextovo nezávislými tokenovými embeddingmi. Pre výskumníkov je to dôležité, pretože čisto interpretačný model by bol menej užitočný, keby sa príliš vzdialil od správania reálne nasadzovaných architektúr. Flex verzia naznačuje, že medzi čitateľnosťou a výkonom môže existovať plynulejšie spektrum.

Význam pre širšiu AI komunitu je dvojitý. Po prvé, práca prispieva k snahám rozumieť modelom na úrovni mechanizmov, nie iba na úrovni výstupných benchmarkov. Ak vieme lepšie zistiť, ktoré tokeny používajú ktoré časti parametrov, môžeme presnejšie skúmať špecializáciu, zabúdanie, prenos medzi úlohami alebo dôvody neočakávaného správania. Po druhé, rovnaký pohľad môže inšpirovať efektívnejšie spôsoby nasadenia modelov, najmä keď sa čoraz viac rieši cena a latencia inferencie.

Netreba však čakať, že MemoryLLM okamžite nahradí štandardné architektúry. Ide o výskumný rámec, ktorý oddeľuje jeden komponent modelu a testuje, čo sa z takého oddelenia dá naučiť. Pri produkčných modeloch zostáva kľúčové, ako sa podobný prístup správa pri veľkých kontextoch, rôznych jazykoch, nástrojovom používaní a modeloch trénovaných na veľmi heterogénnych dátach. Aj samotná predstava FFN ako pamäte je užitočná skratka, nie úplné vysvetlenie všetkej vnútornej dynamiky.

Pre tímy, ktoré sa venujú bezpečnosti, optimalizácii alebo auditovaniu modelov, je však táto línia práce dôležitá. Čím viac schopností presúvame do veľkých transformerov, tým menej stačí vedieť, že model dosiahol dobré skóre v teste. Potrebujeme nástroje na zistenie, ktoré časti modelu nesú akú informáciu, ako stabilne sa používajú a či ich možno kontrolovať bez úplného pretrénovania. MemoryLLM ukazuje jednu možnú cestu: pozerať sa na FFN nie ako na neprehľadnú čiernu skrinku, ale ako na štruktúrovanú pamäť, ktorú možno analyzovať aj technicky využiť.

V širšom kontexte to zapadá do trendu, v ktorom sa optimalizácia inferencie a interpretovateľnosť stretávajú. Ak sa niektoré výpočty dajú pomenovať ako opakovateľné prístupy k pamäťovým miestam, infraštruktúra môže lepšie rozhodovať, čo držať vo VRAM, čo predpočítať a čo načítať až pri potrebe. To je praktická otázka pre prevádzkovateľov veľkých modelov, nielen akademická otázka pre interpretovateľnosť.