Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings

Technische Zusammenfassung: Entkopplung von „Was“ und „Wo“ durch Polar Coordinate Positional Embedding

Problemstellung
In Transformer-Architekturen gleicht der Attention-Mechanismus Keys mit Queries basierend auf zwei Faktoren ab: dem Inhalt („Was“) und der Sequenzposition („Wo“). Die Autoren argumentieren, dass das weit verbreitete Rotary Position Embedding (RoPE), das in Frontier-Modellen wie Llama 3 und DeepSeek verwendet wird, diese beiden Faktoren miteinander verwebt. In RoPE hängt die Rotation, die auf Key- und Query-Vektoren angewendet wird, sowohl von den Amplituden des Inhalts als auch von deren relativen Positionen ab. Diese Verwebung erzeugt ein „Was-Wo“-Konfund (Vermischung), was die Leistung des Modells beeinträchtigt, wenn Aufgaben eine unabhängige Übereinstimmung von nur Inhalt oder nur Position erfordern. Speziell führt der Interaktionsterm in RoPE dazu, dass die Phase der Rotation durch die Inhaltsmerkmale beeinflusst wird, was es für das Modell schwierig macht, Regeln zu erlernen, die strikt auf einer Positionsindexierung oder einem Inhaltsabgleich beruhen.

Methodik: Polar Coordinate Positional Embedding (PoPE)
Die Autoren schlagen PoPE vor, eine Modifikation von RoPE, die Inhalt und Position durch eine Neuformulierung der Key- und Query-Repräsentationen in Polarkoordinaten entkoppelt.

• Neuformulierung: In RoPE werden Komponenten basierend auf der Position rotiert, aber die initiale Phase wird durch den Inhalt bestimmt. In PoPE transformieren die Autoren die reellen Key- und Query-Vektoren in komplexe Vektoren, bei denen:
  • Die Amplitude ($\mu$) ausschließlich aus den Inhaltsmerkmalen mittels einer Softplus-Aktivierungsfunktion ($\sigma(x) = \ln(1 + e^x)$) abgeleitet wird, um Nicht-Negativität zu gewährleisten.
  • Die Phase ($\phi$) strikt durch die Sequenzposition ($s$ oder $t$) und eine komponenten-spezifische Frequenz ($\theta_c$) bestimmt wird.
• Attention-Score: Der Attention-Score wird als der Realteil des Skalarprodukts zwischen der konjugiert transponierten Query und dem Key berechnet. Mathematisch ergibt dies:
  $$a^{PoPE}_{ts} = \sum_{c=1}^{d} \mu_{\tilde{q}_{tc}} \mu_{\tilde{k}_{sc}} \cos((s - t)\theta_c + \delta_c)$$
  wobei $\delta_c$ ein lernbarer Bias-Term ist.
• Wesentliche Unterschiede zu RoPE:
  1. Entkopplung: Der Interaktionsterm ($\phi_{k_{sc}} - \phi_{q_{tc}}$), der in RoPE Inhalt und Position vermischt, wird eliminiert. Die Amplitude (Inhalt) und die Phase (Position) sind strikt getrennt.
  2. Frequenzauflösung: PoPE operiert auf $d$ Komponenten (einzelne Elemente) statt auf $d/2$ Paaren, was die Anzahl der Frequenzkanäle effektiv verdoppelt.
  3. Lernbarer Bias: Ein fester, aber lernbarer Bias $\delta_c$ kann eingeführt werden, um den optimalen relativen Versatz für jede Frequenz abzustimmen, ohne das „Was-Wo“-Konfund erneut einzuführen.

Wesentliche Beiträge

1. Diagnostische Aufgabe (Indirekte Indizierung): Die Autoren führen eine Aufgabe ein, die Modelle dazu zwingt, ein Zielzeichen basierend auf einem relativen Versatz von einem Quellzeichen zu identifizieren (z. B. „Finde das Zeichen 3 Positionen links von 'z'“). Diese Aufgabe erfordert die unabhängige Manipulation von Inhalt und Position.
2. Theoretische Erkenntnis: Die Arbeit liefert eine algebraische Herleitung, die zeigt, wie RoPE Inhalt und Position konfundiert, während PoPE es dem Attention-Mechanismus ermöglicht, Regeln als Konjunktion unabhängiger „Was“- und „Wo“-Abgleiche zu implementieren.
3. Effiziente Implementierung: Eine effiziente Implementierung unter Verwendung von Triton und Flash Attention 2 wird bereitgestellt, welche den Realteil des komplexen Skalarprodukts berechnet, ohne die vollständige komplexe Matrix zu materialisieren, wodurch lediglich ein einzelner Multiplikations-Overhead im Vergleich zur Standard-Attention entsteht.

Ergebnisse

• Indirekte Indizierung: Bei der diagnostischen Aufgabe erreichte RoPE eine durchschnittliche Genauigkeit von nur 11,16 % und hatte Schwierigkeiten, den erforderlichen unabhängigen Abgleich zu erlernen. Im Gegensatz dazu erreichte PoPE eine Genauigkeit von 94,82 % und demonstrierte damit seine Fähigkeit zur Entkopplung der Faktoren.
• Sequenzmodellierung (Musik und Genomik):
  • Musik (JSB und MAESTRO): PoPE übertraf RoPE im Hinblick auf die negative Log-Likelihood (NLL) in beiden Datensätzen und erzielte eine geringere Perplexität.
  • Genomik (Humanes Referenzgenom): PoPE zeigte einen signifikanten Rückgang der NLL im Vergleich zur RoPE-Baseline.
• Sprachmodellierung (OpenWebText): Über verschiedene Modellskalen hinweg (124M, 253M und 774M Parameter) erreichte PoPE konsistent eine geringere Perplexität als RoPE. Der Leistungsunterschied blieb stabil oder nahm mit der Modellgröße sogar zu.
• Downstream-Aufgaben: In Zero-Shot-Evaluierungen auf sechs Downstream-Aufgaben (LAMBADA, BLiMP, CBT, HellaSwag, PIQA, ARC-E) zeigten PoPE-basierte Modelle eine höhere mittlere Genauigkeit über alle Modellgrößen hinweg im Vergleich zu RoPE-basierten Modellen.
• Längenextrapolation:
  • Zero-Shot: PoPE demonstrierte starke Zero-Shot-Längenextrapolationsfähigkeiten auf Sequenzen, die bis zu 10-mal länger sind als der Trainingskontext (1024 vs. 10420 Token), ohne jegliches Fine-Tuning oder Frequenzinterpolation.
  • Vergleich mit YaRN: Während YaRN (eine Methode zur Extrapolation) die Degradation von RoPE bis zu seinem Fine-Tuning-Limit (4096 Token) abmildert, verschlechtert es sich jenseits dessen signifikant. PoPE behielt die Leistung auf längeren Sequenzen bei, ohne die Interpolationstechniken oder das umfangreiche Fine-Tuning zu benötigen, das YaRN erforderlich macht.
  • Analyse der Frequenznutzung: Eine Analyse der Feature-Normen ergab, dass RoPE-Modelle dazu neigen, hochfrequente Kanäle zu unterdrücken, um Rauschen zu vermeiden, während PoPE-Modelle einen breiter verteilten Bereich von Frequenzen, einschließlich hoher Frequenzen, über die Layer hinweg nutzen.

Bedeutung und Ansprüche
Das Paper behauptet, dass PoPE eine prinzipielle Lösung für die in RoPE inhärente „Was-Wo“-Verwebung bietet. Durch die Entkopplung von Inhalt und Position verbessert PoPE die Dateneffizienz, die asymptotische Genauigkeit und vor allem die Generalisierung der Kontextlänge. Die Autoren betonen, dass diese Gewinne durch eine geringfügige Modifikation der bestehenden RoPE-Architektur erzielt werden, die keine komplexen architektonischen Änderungen erfordert.

Die Arbeit wird als fundamentale Forschung zum Thema Sequenzmodellierung präsentiert. Die Autoren merken an, dass die Methode zwar die Fähigkeiten in Domänen wie Musik, Genomik und natürlicher Sprache verbessert (was potenziell Anwendungen in der Proteinfaltung oder Wirkstoffforschung unterstützen kann), sie jedoch Dual-Use-Risiken birgt, die für Foundation Models üblich sind, wie etwa die Generierung schädlicher Inhalte oder ausgefeigter Deepfakes. Der primäre Beitrag ist der Nachweis, dass die Trennung von Inhalts- und Positionsinformationen zu robusteren und generalisierbareren Transformer-Modellen führt.