Enhancing Layer Interaction Using Key-Correlated Layer Attention

Technisches Resümee: Verbesserung der Schichtinteraktion durch Key-Korrelierte Schicht-Attention

Problemstellung
Jüngste Fortschritte im Design tiefer neuronaler Netze (DNN) haben die Layer Attention genutzt, um Interaktionen zwischen den Schichten zu verbessern, indem die Ausgaben der Schichten als Sequenz-Token behandelt werden. Die Standard-Layer-Attention leidet jedoch unter einer quadratischen Komplexität ($O(D^2)$) in Bezug auf die Netzwerktiefe $D$ sowie einem signifikanten Speicherwachstum durch das notwendige Caching der Ausgaben aller vorangegangenen Schichten.

Um diese Probleme zu adressieren, führten frühere Arbeiten die Recurrent Layer Attention (RLA) und lineare Approximationen wie MRLA-L ein. Während RLA den Rechenaufwand durch das Caching von Key- ($K$) und Value-Repräsentationen ($V$) reduziert, fungiert es als statischer Mechanismus, bei dem die Merkmale früherer Schichten feststehen, sobald sie generiert wurden, was die dynamische Aktualisierung von Informationen einschränkt. Darüber hinaus versuchen lineare Approximationen wie MRLA-L, die Komplexität durch eine lineare Approximation der Attention-Ausgabe zu reduzieren. Die Autoren beobachten, dass dieser Ansatz dazu führt, dass die Beiträge früherer Schichten exponentiell abgeschwächt werden (wie in ihrer Analyse der Zerfallskoeffizienten $\beta_l$ gezeigt), was es schwierig macht, effektive langreichweitige Abhängigkeiten über mehrere Schichten hinweg zu etablieren. Zudem versagen bestehende Linearisierungsmethoden oft darin, den Speicherengpass beim Caching vorangegangener Schichtausgaben oder die spezifischen Herausforderungen bei der Anwendung linearer Attention auf hierarchische DNN-Strukturen zu adressieren.

Methodik: Key-Correlated Layer Attention (KCLA)
Die vorgeschlagene Key-Correlated Layer Attention (KCLA) ist eine neuartige linearisierte Layer-Attention-Methode, die darauf ausgelegt ist, eine lineare Rechenkomplexität zu erreichen und gleichzeitig dynamische Aktualisierungen sowie die Modellierung von Langreichweiten-Abhängigkeiten zu bewahren. Die Methodik basiert auf drei Kernbeobachtungen und Mechanismen:

1. Beobachtung der Key-Korrelation: Eine empirische Analyse zeigt, dass Key-Repräsentationen über verschiedene Schichten hinweg in der Standard-Layer-Attention eine hohe Kosinus-Ähnlichkeit aufweisen. Unter Ausnutzung dessen approximiert KCLA die Standard-Layer-Attention-Ausgabe, indem sie die lineare Transformation des Values ($V$) herausfaktorisiert und die Key-Repräsentationen als richtungsorientiert korreliert approximiert. Dies ermöglicht die Aggregation von Schichtausgaben in einen rekurrenten Zustand ($U^t$), ohne die vollständige Historie der $K$- und $V$-Matrizen speichern zu müssen.
2. Dynamische Softmax-Restaurierung: Eine kritische Herausforderung bei der Linearisierung von Attention besteht im Verlust der Softmax-Operation, die essenziell für die Normalisierung von Attention-Scores und die Ermöglichung dynamischer Verteilungsanpassungen ist. KCLA führt einen modifizierten Multi-Head-Attention-Mechanismus ein, um diese Fähigkeit wiederherzustellen:
   • Es verwendet mehrere Köpfe mit unterschiedlichen Skalierungsfaktoren (Temperaturen), die aus einer logarithmisch gleichmäßigen Sequenz abgeleitet sind.
   • Es nutzt eine lernbare lineare Transformation $\Phi$, um die Norm der Query ($\|Q^t\|$) und den Korrelationswinkel ($\cos \theta$) in kopfspezifische Gewichte abzubilden.
   • Dies erlaubt es dem Modell, die Attention-Verteilung in jeder Schicht dynamisch anzupassen, wodurch die Rolle des Temperaturparameters der Softmax-Funktion effektiv simuliert wird, ohne die vollständige $N \times N$-Matrix berechnen zu müssen.
3. Cross-Stage-Adaption: Um DNNs mit variierenden Feature-Map-Größen und Kanaldimensionen über verschiedene Stadien hinweg zu unterstützen, implementiert KCLA ein spezifisches Adaptionsverfahren (Algorithmus 1). Dies beinhaltet die Normalisierung akkumulierter Merkmale, die Anwendung einer $1 \times 1$-Konvolution mit Stride zur Downsampling-Operation und die Neuberechnung von Normalisierungsfaktoren, um die dimensionalen Disparitäten zwischen den Stadien zu überbrücken.

Zentrale Beiträge

• Beobachtung und Beweis: Die Autoren demonstrieren, dass Key-Repräsentationen über verschiedene Schichten in der Standard-Layer-Attention in ihrer Richtung hoch korreliert sind, was die theoretische Basis für die lineare Approximation liefert.
• Neuartige Linearisierung: KCLA wird als linearisierte Layer-Attention-Methode vorgeschlagen, die die statische Natur von RLA und die exponentielle Abschwächung von MRLA-L vermeidet. Sie erreicht eine lineare Rechenkomplexität und eine konstante räumliche Komplexität, die unabhängig von der Netzwerktiefe ist.
• Dynamische Softmax-Approximation: Eine neue Methode wird eingeführt, um die Softmax-Operation mittels Multi-Head-Skalierung und lernbarer Transformationen zu approximieren, was eine dynamische und adaptive Anpassung der Attention-Verteilungen in jeder Schicht ermöglicht.

Experimentelle Ergebnisse
Die Autoren evaluierten KCLA in drei grundlegenden Vision-Aufgaben: Bildklassifizierung, Objekterkennung und medizinische Bildsegmentierung.

• Bildklassifizierung: Auf CIFAR-100 erreichte KCLA die beste Leistung über ResNet-Tiefen von 20, 56 und 110 hinweg und übertraf dabei MRLA-L, DLA-L und die Standard-Layer-Attention. Auf ImageNet-1K verbesserte die Integration von KCLA in ResNet-50 und ResNet-101 die Top-1-Genauigkeit um jeweils +1,7 % bzw. +1,4 % gegenüber der Baseline. Es behielt einen leichten, aber konsistenten Vorteil gegenüber MRLA-L bei, während es deutlich weniger Parameter als die dynamische DLA-L-Methode verwendete.
• Objekterkennung: Auf dem COCO2017-Datensatz mit Faster R-CNN und Mask R-CNN zeigte KCLA eine kompetitive Leistung. Mit ResNet-101 übertraf es die komplexere DLA-L-Architektur in mehreren Metriken (AP$_{75}$, AP$_S$, AP$_M$), während es 2,2 Millionen Parameter weniger benötigte.
• Medizinische Bildsegmentierung: Auf den Datensätzen ISIC2017, ISIC2018 und Kvasir-SEG übertraf KCLA konsistent MRLA-L und die Baseline ResNet-UNet und erzielte überlegene mIoU- und DSC-Werte bei minimalem Parameter-Overhead.
• Komplexitätsanalyse: Im Gegensatz zu MRLA-B, dessen Speicherbedarf und FLOPs quadratisch mit der Tiefe wachsen, weist KCLA ein lineares Wachstum sowohl in der Rechenlast als auch in den Speicherkosten auf.

Bedeutung und Ansprüche
Das Paper behauptet, dass KCLA die Lücke zwischen der Effizienz linearer Attention und der Repräsentationskraft der Standard-Attention erfolgreich schließt. Durch die direkte Approximation des Standard-Layer-Attention-Mechanismus basierend auf Key-Korrelationen vermeidet KCLA die Probleme der statischen Merkmalsrepräsentation von RLA und die Degradierung der langreichweitigen Abhängigkeiten von MRLA-L.

Die Autoren betonen, dass ihre Methode die Fähigkeit beibehält, langreichweitige Cross-Layer-Abhängigkeiten effizient zu modellieren, und bietet somit ein wünschenswertes Gleichgewicht zwischen Effizienz und Leistung. Sie merken an, dass die vorgeschlagene Methode zur Wiederherstellung der Softmax-Funktion ein entscheidender Faktor für das Erreichen dieser Ergebnisse ist. Während die aktuelle Implementierung auf vordefinierten Temperaturparametern basiert, identifizieren die Autoren das Design einer vollständig selbstadaptiven, dynamischen Temperatur-Multi-Head-Linearen-Layer-Attention als Richtung für zukünftige Optimierungen. Der Code ist öffentlich zugänglich, um die Reproduzierbarkeit zu erleichtern.