On The Application of Linear Attention in Multimodal Transformers

Problemstellung

Multimodale Transformer-Architekturen bilden das Rückgrat moderner Vision-Language-Modelle (z. B. für Bildbeschreibung oder visuelle Fragebeantwortung). Ein zentrales Hindernis für deren Skalierbarkeit ist jedoch die quadratische Komplexität der Standard-Self-Attention-Mechanismen ($O(N^2D)$), die in Abhängigkeit von der Sequenzlänge $N$ steht.
Da multimodale Eingaben oft aus Text und Hunderten oder Tausenden von visuellen Tokens bestehen, führt diese quadratische Skalierung zu erheblichen Rechen- und Speicherkosten, insbesondere bei langen Sequenzen und hochauflösenden Eingaben. Obwohl Hardware-Optimierungen wie FlashAttention die Geschwindigkeit verbessern, bleibt die asymptotische Komplexität ein Engpass.

Methodik

Die Autoren untersuchen die Machbarkeit von Linearer Attention (LA) als effiziente Alternative innerhalb multimodaler Frameworks.

1. Grundlegender Ansatz:
   Statt des Softmax-Kernels wird ein dekomponierbarer Feature-Mapping-Funktion $\phi(\cdot)$ verwendet, um die Assoziativität der Matrixmultiplikation auszunutzen. Dies ändert die Berechnungsreihenfolge von $O(N^2D)$ auf $O(ND^2)$.
   Die Autoren verwenden einen affinen Kernel $f(x) = 1 + x$.

2. Identifizierte Herausforderungen & Anpassungen:

   • Problem der negativen Scores: Ein einfacher affiner Kernel kann negative Werte liefern, wenn $q_i \cdot k_n < -1$.
     • Lösung: Normalisierung der Query- und Key-Vektoren ($q_i, k_i$), sodass ihr Skalarprodukt im Bereich $[-1, 1]$ liegt.
   • Problem der Über-Glättung (Over-smoothing): Die ursprüngliche Formel mit Nenner (Normalisierung auf Summe 1) führt bei langen Sequenzen zu einer Abnahme der Aufmerksamkeitsscores und einer zu gleichmäßigen Verteilung (siehe Abbildung 1 im Paper). Dies erschwert es dem Modell, relevante Token-Paare zu unterscheiden, und behindert den Gradientenfluss.
     • Lösung: Die Autoren schlagen vor, den Nenner in der Attention-Formel zu entfernen und stattdessen eine feste Skalierung zu verwenden (Gleichung 7 im Paper). Dies erhält die Scores im Bereich $[0, 1]$ und verhindert die Über-Glättung, wodurch die Ausdruckskraft (Expressivity) des Modells erhalten bleibt.

3. Experimentelles Setup:

   • Modelle: Vision Transformer (ViT) in den Größen S/16, B/16 und L/16.
   • Datensatz: Pre-Training auf LAION-400M unter Verwendung des OpenCLIP-Frameworks.
   • Validierung: Zero-Shot-Accuracy auf ImageNet-21K.
   • Hardware: Tests auf NVIDIA A5500 und H200 GPUs.

Wichtige Beiträge

1. Erste systematische Untersuchung von LA in Multimodalität: Die Arbeit zeigt, dass Linear Attention in multimodalen Kontexten (Vision + Language) eingesetzt werden kann, ohne die Leistungsfähigkeit im Vergleich zu Standard-Attention signifikant zu beeinträchtigen.
2. Verbesserung der LA-Formulierung: Die Autoren identifizieren und lösen das Problem der Über-Glättung durch den Verzicht auf den Nenner in der Attention-Berechnung nach Vor-Normalisierung der Vektoren. Dies ist entscheidend für die Lernfähigkeit des Modells.
3. Analyse der Skalierungsgesetze: Es wird demonstriert, dass Modelle mit modifizierter Linear Attention denselben Skalierungsgesetzen (Power-Law zwischen Modellgröße und Loss) folgen wie Modelle mit Standard-Attention.

Ergebnisse

• Rechenzeit-Skalierung:
  • Die Messungen zeigen eine klare Skalierung von $O(N)$ für LA im Vergleich zu $O(N^2)$ für Standard-Attention.
  • Bei Sequenzlängen ab ca. $10^3$ Tokens wird ein deutlicher Geschwindigkeitsvorteil sichtbar. Bei sehr langen Sequenzen ($4 \times 10^6$ Tokens) erreicht LA eine bis zu $10^3$-fache Reduktion der Laufzeit.
• Trainingsleistung:
  • Mit den vorgeschlagenen Anpassungen (Gleichung 7) erreicht LA eine Endleistung, die der von Standard-Attention (Softmax) auf ViT-S, ViT-B und ViT-L Modellen entspricht.
  • Die konventionelle LA-Formulierung (ohne Anpassungen) zeigt eine sehr langsame Konvergenz und schlechte Leistung, was die Notwendigkeit der vorgeschlagenen Modifikation unterstreicht.
  • Ein Nachteil ist, dass LA etwas mehr Epochen benötigt, um das gleiche Leistungsniveau zu erreichen, was jedoch durch die schnellere Iterationszeit pro Epoche kompensiert wird.
• Skalierungsgesetze (Scaling Laws):
  • Die empirischen Fits zeigen, dass der Loss mit zunehmender Modellgröße für beide Mechanismen ähnlich abnimmt:
    • Standard-Attention: $L(N) \approx 593 N^{-0.362}$
    • Linear-Attention (angepasst): $L(N) \approx 376 N^{-0.332}$
  • Dies beweist, dass LA keine Leistungseinbußen bei der Skalierung auf große Modelle erfährt.

Bedeutung und Fazit

Die Studie positioniert Linear Attention als eine robuste und skalierbare Lösung für die nächste Generation multimodaler Transformer.

• Effizienz: Sie ermöglicht die Verarbeitung extrem langer Sequenzen und hochauflösender Bilder, die mit Standard-Attention aufgrund des Speicher- und Rechenaufwands oft unmöglich wären.
• Leistung: Durch die vorgeschlagenen Modifikationen (Entfernung des Nenners, Normalisierung) wird die Expressivität des Modells erhalten, sodass keine Kompromisse bei der Genauigkeit (Zero-Shot-Accuracy) eingegangen werden müssen.
• Zukunftsperspektive: Die Arbeit legt nahe, dass LA ein praktikabler Ersatz für Softmax-Attention in großen multimodalen Systemen ist, insbesondere in Szenarien, in denen die Sequenzlänge der limitierende Faktor ist.