Technische Samenvatting: Rotary Position Encodings voor Grafen (WIRE)

Probleemstelling

Transformers vertrouwen zwaar op positie-encodings (PE's) om de ruimtelijke of topologische relaties tussen tokens vast te leggen. Hoewel Rotary Position Encodings (RoPE) de standaard zijn geworden voor Large Language Models (LLM's) en Vision Transformers (ViT's) vanwege hun vermogen om relatieve positionele informatie via token-rotatie te injecteren, blijft hun toepassing op graafgestructureerde gegevens onderbelicht.

Grafen missen een canoniek coördinatenstelsel, waardoor de definitie van "positie" en "afstand" niet triviaal is. Bestaande grafische PE's vallen over het algemeen uiteen in twee categorieën:

1. Absolute Position Encodings (APE's): Vaak gebaseerd op de eigenvectoren van de grafische Laplaciaan. Hoewel deze structurele informatie vastleggen, presteren ze over het algemeen slechter dan relatieve methoden en missen ze specifieke invariantie-eigenschappen.
2. Relative Position Encodings (RPE's): Typisch geïmplementeerd als biases die worden toegevoegd aan attention-logits (bijv. gebaseerd op de kortste padafstand of effectieve weerstand). Hoewel effectief, vereisen deze methoden de instantiatie van de volledige $N \times N$ attention-matrix, wat ze incompatibel maakt met lineaire attention-mechanismen (die schalen als $\mathcal{O}(N)$) en efficiënte KV-caching verhindert.

De paper stelt de vraag: In welke mate kunnen rotary position encodings worden toegepast op graafgestructureerde gegevens?

Methodologie: Wave-Induced Rotary Encodings (WIRE)

De auteurs stellen WIRE (Wave-Induced Rotary Encodings) voor, een methode die RoPE aanpast voor grafen door gebruik te maken van het grafische spectrum.

Kernalgoritme

1. Spectrale Feature Extractie: Bereken de laagste $m$ eigenvectoren ($\mathbf{u}_k$) en eigenwaarden ($\lambda_k$) van de grafische Laplaciaan $\mathbf{L}$.
2. Coördinaatdefinitie: Definieer spectrale features voor elke node $i$ als een vector $\mathbf{r}_i$ bestaande uit de waarden van de node in deze eigenvectoren (bijv. $\mathbf{r}_i = [\mathbf{u}_k[i]]_{k=0}^{m-1}$).
3. Rotary Transformatie: Pas de standaard RoPE-rotatie toe op de query ($\mathbf{q}_i$) en key ($\mathbf{k}_i$) vectoren met behulp van deze spectrale coördinaten:
   $$\mathbf{z}_i \to \text{RoPE}(\mathbf{r}_i)\mathbf{z}_i$$
   waarbij $\text{RoPE}(\mathbf{r}_i)$ 2D-blokken van de token-vector roteert door hoeken die worden bepaald door de spectrale coördinaten.

Theoretische Eigenschappen

De paper stelt theoretische garanties vast voor WIRE:

• Permutatie-equivariantie: WIRE is equivariant onder node-herordening (tot tekenomklappingen en rotaties in degeneratieve subruimten), een noodzakelijke eigenschap voor graafleren.
• Herstel van Standaard RoPE: Op roostergrafen (Cartesiaanse producten van padgrafen) herstelt WIRE het gedrag van standaard RoPE gebruikt in LLM's en ViT's. De spectrale coördinaten op roosters komen monotoon overeen met ruimtelijke coördinaten, waardoor ze effectief worden genormaliseerd naar het bereik $(-1, 1)$.
• Afhankelijkheid van Effectieve Weerstand: Theoretisch, wanneer WIRE-frequenties willekeurig worden gesampled, dempt de verwachte transformatie de query-key attention-scores proportioneel aan de effectieve weerstand tussen nodes. Effectieve weerstand is een metriek die dient als een ondergrens voor de kortste padafstand.
• Compatibiliteit met Lineaire Attention: In tegenstelling tot bias-gebaseerde RPE's, wijzigt WIRE de tokens (queries en keys) direct in plaats van de attention-logits. Dit maakt het mogelijk om WIRE te gebruiken met lineaire attention-mechanismen (bijv. Performers) zonder de volledige attention-matrix te instantiëren, waardoor $\mathcal{O}(N)$ schaling behouden blijft.

Belangrijkste Bijdragen

1. Introductie van WIRE: Een nieuwe position encoding voor grafen die RoPE generaliseert door grafische Laplaciaan-eigenvectoren als positionele coördinaten te gebruiken.
2. Theoretische Analyse:
   • Bewijs dat WIRE standaard RoPE herstelt op roostergrafen.
   • Demonstratie dat WIRE impliciet de attention moduleert op basis van de effectieve weerstand van de graaf, wat een geprincipeerde topologische inductieve bias biedt.
   • Vaststelling van compatibiliteit met lineaire attention, een kenmerk dat voorheen niet beschikbaar was voor weerstands-gebaseerde relatieve encodings.
3. Empirische Validatie: Uitgebreide experimenten over synthetische taken, puntenwolken (point clouds) en standaard grafische benchmarks.

Experimentele Resultaten

De auteurs trainden meer dan 200 transformer-modellen om WIRE te evalueren:

• Synthetische Taken:
  • Monochrome Subgrafen: WIRE verminderde de test-RMSE significant vergeleken met baselines zonder WIRE, vooral naarmate de grafische topologie complexer werd (door randen uit roosters verwijderd).
  • Kortste Pad Voorspelling: WIRE halveerde bijna de test-RMSE op Watts-Strogatz grafen vergeleken met APE-alleen baselines, door gebruik te maken van de afhankelijkheid van effectieve weerstand.
• Point Cloud Transformers: Toegepast op ModelNet40 (classificatie) en ShapeNet (segmentatie). WIRE leverde consistente winsten op over standaard PCT-baselines en kwam overeen met of overtrof Cartesiaanse RoPE, terwijl het SE(3)-invariantie behield.
• Grafische Benchmarks (GNN Taken):
  • Geëvalueerd op datasets inclusief MNIST, CIFAR-10, PATTERN, CLUSTER, en OGB datasets (bijv. ogbg-molhiv).
  • Met behulp van Performers (lineaire attention), verbeterde WIRE consequent de prestaties ten opzichte van baselines.
  • Opvallend genoeg verkleinde WIRE vaak de prestatiekloof tussen efficiënte lineaire attention-modellen en dure full-rank softmax transformers. Bijvoorbeeld, op MalNet-Tiny bereikten WIRE-enabled Performers een nauwkeurigheid die vergelijkbaar is met full-rank transformers, maar konden ze getraind worden op een enkele T4 GPU.
• Efficiëntie: WIRE-parameters vormen minder dan 1% van de totale modelparameters. De methode is compatibel met diverse attention-mechanismen, waaronder softmax, lineair (Performer), SGFormer, en BigBird.

Betekenis en Claims

De paper claimt dat WIRE erin slaagt de kloof te overbruggen tussen het succes van RoPE in sequentie- en beelddata en de behoeften van graafgestructureerde gegevens. De primaire betekenis ligt in:

1. Schaalbaarheid: Het maakt het mogelijk om sterke structurele inductieve biases (zoals effectieve weerstand) te gebruiken in lineaire attention modellen, wat voorheen onmogelijk was met bias-gebaseerde RPE's.
2. Prestaties: Het verhoogt de nauwkeurigheid van transformers op grafische taken, waardoor efficiënte lineaire modellen vaak de prestaties kunnen evenaren van computationeel duurdere softmax-gebaseerde modellen.
3. Theoretische Onderbouwing: Het biedt een wiskundig geprincipeerde manier om grafische topologie in attention-mechanismen te injecteren, waarbij bekend gedrag op roosters wordt hersteld terwijl nieuwe inzichten in effectieve weerstand op algemene grafen worden geboden.

De auteurs concluderen dat WIRE een "drop-in" toevoeging is aan bestaande grafische transformer-architecturen die een lichtgewicht, effectieve structurele inductieve bias biedt, wat verdere research stimuleert naar het generaliseren van rotary encodings naar andere datamodaliteiten.