AdaMerge zrýchľuje Vision Transformery bez tréningu cez pozornejšie spájanie tokenov

Na arXive pribudol preprint AdaMerge, ktorý sa venuje praktickému problému Vision Transformerov: ich pozornosť nad obrazovými tokenmi je výpočtovo drahá a pri veľkých vstupoch rýchlo naráža na kvadratické náklady. Výskumníci preto navrhujú metódu salience-aware adaptive token merging, teda spájanie tokenov, ktoré berie do úvahy ich významnosť, a zároveň mení intenzitu kompresie podľa vrstvy a konkrétneho vstupu.

Vision Transformer rozdelí obraz na tokeny a spracúva ich pomocou mechanizmu pozornosti. Čím viac tokenov, tým viac párových vzťahov musí model porovnávať. Jednou z populárnych ciest k zrýchleniu je token merging: podobné tokeny sa spoja, počet prvkov klesne a model potrebuje menej operácií. Problém je, že jednoduché spájanie často predpokladá, že tokeny majú podobnú hodnotu. Pri obrazoch to neplatí. Niektoré časti nesú kľúčový objekt alebo hranu, iné sú redundantné pozadie.

AdaMerge sa pokúša túto nerovnosť zachytiť dvoma mechanizmami. Prvý používa salience-weighted similarity, teda podobnosť váženú významnosťou. Autori počítajú proxy mieru dôležitosti cez centrálne postavenie tokenu vo vzťahoch medzi príznakmi a túto informáciu vkladajú do skóre pri párovaní tokenov. Zjednodušene: ak je token pre reprezentáciu obrazu dôležitý, nemal by byť agresívne pohltený menej dôležitým susedom len preto, že vektorovo vyzerajú podobne.

Druhý mechanizmus je adaptive merging intensity. Namiesto pevného počtu spojených tokenov v každej vrstve metóda využíva predpočítané štatistiky podobnosti po vrstvách a podľa konkrétneho vstupu upravuje, kde sa dá komprimovať viac a kde menej. To je dôležité, pretože redundantnosť sa nemusí v modeli rozkladať rovnomerne. Niektoré vrstvy môžu mať veľa opakovania, iné nesú jemnejšie rozlíšenie medzi triedami.

Autori tvrdia, že AdaMerge je tréningovo nenáročná metóda: netreba znovu učiť celý model ani vytvárať špeciálnu menšiu architektúru. To je praktická výhoda pre tímy, ktoré už majú nasadené alebo odladené ViT modely a chcú znížiť výpočtové náklady bez rizika rozsiahleho retréningu. V produkcii býva práve takáto kompatibilita rozhodujúca, pretože aj malá zmena modelu môže vyžadovať nové overovanie presnosti, latencie a správania na okrajových dátach.

V experimentoch na ImageNet-1k s ViT-B/16 má AdaMerge podľa preprintu prekonávať ToMe, PiToMe aj DSM pri rovnakých úrovniach FLOPs. Pri prevádzkovom bode 13,4G FLOPs uvádzajú autori pokles Top-1 presnosti len o 1,06 percentuálneho bodu, zatiaľ čo PiToMe klesá o 1,45 bodu a DSM o 4,62 bodu. Takéto čísla treba brať ako výsledky z konkrétneho experimentálneho nastavenia, no naznačujú, že výber tokenov podľa významnosti môže byť výhodnejší než čisto rovnostárske spájanie.

Pre širší trh AI akcelerácie je práca zaujímavá tým, že nejde o ďalší väčší model ani o hardvérový trik. Skôr ukazuje, že veľká časť efektivity môže prísť z lepšieho hospodárenia s reprezentáciou počas inferencie. Podobný motív vidno aj pri jazykových modeloch v kompresii kontextu, selektívnom načítaní pamäte či špekulatívnom dekódovaní. V počítačovom videní sa táto otázka prejavuje najmä v tom, ktoré časti obrazu treba niesť ďalej a ktoré už sú pre rozhodnutie nadbytočné.

Praktické nasadenie však bude závisieť od robustnosti mimo ImageNetu. Tokeny dôležité pre klasifikáciu prirodzených obrázkov nemusia stačiť pri medicínskom zobrazení, pri detekcii malých objektov alebo pri priemyselnej kontrole kvality, kde rozhoduje drobný detail. Ak kompresia zle odhadne salienciu, môže odstrániť práve signál, ktorý používateľ potrebuje. Preto bude dôležité testovať podobné metódy nielen na priemernej presnosti, ale aj na citlivosti voči malým objektom, posunom domény a bezpečnostne kritickým triedam.

AdaMerge je napriek tomu dobrým príkladom smeru, ktorý bude pre nasadzovanie multimodálnych modelov čoraz dôležitejší. Ako rastie počet vizuálnych asistentov, robotických systémov a modelov spracúvajúcich video, nestačí zvyšovať výkon hrubou silou. Potrebné budú metódy, ktoré vedia dynamicky rozhodnúť, čo si model má ponechať. Ak sa salienčne riadené spájanie tokenov potvrdí aj v ďalších úlohách, môže byť jednou z ciest, ako dostať Vision Transformery bližšie k lacnejšej produkčnej inferencii.