On The Application of Linear Attention in Multimodal Transformers

Probleemstelling

Multimodale Transformers vormen de ruggengraat van de meest geavanceerde vision-language modellen (zoals CLIP). Hun succes is grotendeels te danken aan zelf- en kruis-attentie mechanismen die relaties binnen en tussen modaliteiten (beeld en tekst) kunnen vastleggen. Echter, de standaard Softmax-attentie heeft een kwadratische complexiteit ($O(N^2D)$) ten opzichte van de sequentielengte $N$.

In multimodale contexten, waar invoer vaak honderden of duizenden visuele tokens bevat naast teksttokens, wordt deze kwadratische schaling een kritieke bottleneck. Dit beperkt de schaalbaarheid van modellen voor lange sequenties en hoge resolutie-invoer, wat de verwerking van steeds grotere en complexere datasets (zoals LAION-400M) bemoeilijkt.

Methodologie

De auteurs onderzoeken de haalbaarheid van Lineaire Attentie (LA) als een efficiënter alternatief binnen multimodale frameworks.

1. Kernidee: In plaats van de niet-lineaire Softmax-kern te gebruiken, vervangt LA deze door een ontbindbare feature-map $\phi(\cdot)$, waardoor de berekeningsvolgorde kan worden herschreven om associativiteit van matrixvermenigvuldiging te benutten. Dit verlaagt de complexiteit van $O(N^2D)$ naar $O(ND^2)$.
2. Specifieke Implementatie:
   • De auteurs gebruiken een affiene kern gedefinieerd als $f(x) = 1 + x$.
   • Normalisatie: Om te voorkomen dat de affiene kern negatieve scores oplevert (wat zou gebeuren als $q_i \cdot k_n < -1$), normaliseren de auteurs de query- en key-vectoren zodat hun dot-product tussen -1 en 1 blijft.
3. Probleem met de standaard LA-formulering:
   • De standaard LA-formulering (met deling door de som van de termen) leidt bij toenemende sequentielengte tot over-smoothing van de attentie-matrices. De scores worden te uniform, waardoor het model moeite heeft om relevante token-paren te onderscheiden en de gradiëntstroom tijdens training wordt gehinderd.
4. De voorgestelde oplossing (Equation 7):
   • De auteurs stellen voor om de noemer in de standaard formule weg te laten. Omdat $Q$ en $K$ al genormaliseerd zijn, blijven de individuele termen $(1 + q_i \cdot k_n)$ natuurlijk begrensd.
   • De nieuwe formule wordt: $o_{ij} = \frac{\sum_{n=1}^i (1 + q_i \cdot k_n) v_{nj}}{2}$.
   • Dit behoudt de attentiescores tussen 0 en 1, voorkomt over-smoothing en verbetert het expressievermogen van het model.

Belangrijkste Bijdragen

• Onderzoek naar LA in Multimodale Modellen: De auteurs trainen ViT-S-16, ViT-B-16 en ViT-L-16 modellen op het LAION-400M dataset en tonen aan dat LA prestaties kan bereiken die vergelijkbaar zijn met standaard Softmax-attentie.
• Efficiëntie-analyse: Ze analyseren de computationele efficiëntie van LA in multimodale settings, waarbij ze aantonen dat de lineaire schaling aanzienlijke tijdwinst oplevert voor lange sequenties.
• Schalingswetten: Ze bewijzen dat LA dezelfde schalingswetten (accuracy vs. modelgrootte) volgt als standaard attentie, wat aangeeft dat LA schaalbaar is zonder prestatieverlies bij grotere modellen.
• Verbeterde Formulering: Ze introduceren een specifieke aanpassing (weglaten van de noemer na normalisatie) die het probleem van over-smoothing oplost en de trainbaarheid van LA in multimodale contexten mogelijk maakt.

Resultaten

De evaluatie is uitgevoerd op ViT-S, ViT-B en ViT-L architecturen getraind op LAION-400M, met validatie via ImageNet-21K zero-shot nauwkeurigheid.

1. Tijdschaal:
   • Standaard attentie toont $O(N^2)$ schaling, terwijl LA $O(N)$ schaling toont.
   • Bij sequentielengtes vanaf $10^3$ tokens is het verschil merkbaar; bij $4 \times 10^6$ tokens is LA ongeveer $10^3$ keer sneller.
2. Trainingsprestaties:
   • Met de voorgestelde aanpassingen convergeert LA naar een vergelijkbare eindwaarde (verlies en nauwkeurigheid) als standaard attentie.
   • De standaard LA-formulering (zonder aanpassing) convergeert extreem traag of faalt vanwege het gebrek aan contrast in de attentiewaarden.
   • Hoewel LA meer epochs nodig heeft om te convergeren, is de totale rekentijd lager door de efficiëntere per-epoch verwerking.
3. Schalingswetten:
   • De training loss neemt af volgens een machtswet ($L(N) \approx aN^{-\alpha}$) voor zowel standaard als lineaire attentie.
   • De exponenten zijn vergelijkbaar ($\approx -0.36$ voor standaard vs. $\approx -0.33$ voor LA), wat aantoont dat LA geen prestatiedegradatie ondergaat bij het schalen naar grotere modellen.

Betekenis en Conclusie

Dit paper positioneert Lineaire Attentie als een robuust en schaalbaar alternatief voor Softmax-attentie in de volgende generatie multimodale Transformers. De belangrijkste conclusies zijn:

• Schaalbaarheid: LA lost het kwadratische bottleneck-probleem op, waardoor het mogelijk wordt om modellen te trainen op zeer lange sequenties en hoge resoluties zonder onbeperkte rekenkosten.
• Prestatiebehoud: Door de specifieke aanpassingen (normalisatie en weglating van de noemer) behoudt LA het expressievermogen dat nodig is voor complexe multimodale taken.
• Toekomstperspectief: De bevindingen suggereren dat LA een ideale oplossing is voor toekomstige grote multimodale modellen die worden getraind op steeds grotere datasets, waarbij efficiëntie en schaalbaarheid cruciaal zijn.