Résumé Technique : Encodages de Position Rotatifs pour les Graphes (WIRE)

Énoncé du Problème

Les Transformers dépendent fortement des encodages de position (PE) pour capturer les relations spatiales ou topologiques entre les jetons. Bien que les encodages de position rotatifs (RoPE) soient devenus la norme pour les grands modèles de langage (LLM) et les Vision Transformers (ViT) grâce à leur capacité à injecter une information positionnelle relative via la rotation des jetons, leur application aux données structurées en graphes reste sous-explorée.

Les graphes manquent d'un système de coordonnées canonique, ce qui rend la définition de la « position » et de la « distance » non triviale. Les encodages de position pour graphes existants se répartissent généralement en deux catégories :

1. Encodages de Position Absolus (APE) : Souvent basés sur les vecteurs propres du laplacien du graphe. Bien qu'ils capturent l'information structurelle, ils sont généralement moins performants que les méthodes relatives et manquent de propriétés d'invariance spécifiques.
2. Encodages de Position Relatifs (RPE) : Typiquement implémentés comme des biais ajoutés aux logits d'attention (par exemple, basés sur la distance du plus court chemin ou la résistance effective). Bien qu'efficaces, ces méthodes nécessitent l'instanciation de la matrice d'attention complète $N \times N$, ce qui les rend incompatibles avec les mécanismes d'attention linéaire (qui passent à l'échelle en $\mathcal{O}(N)$) et empêche un mise en cache efficace des KV (KV-caching).

La question posée est : Dans quelle mesure les encodages de position rotatifs peuvent-ils être appliqués aux données structurées en graphes ?

Méthodologie : Wave-Induced Rotary Encodings (WIRE)

Les auteurs proposent WIRE (Wave-Induced Rotary Encodings), une méthode qui adapte le RoPE pour les graphes en utilisant le spectre du graphe.

Algorithme Central

1. Extraction de Caractéristiques Spectrales : Calculer les $m$ plus petits vecteurs propres ($\mathbf{u}_k$) et valeurs propres ($\lambda_k$) du laplacien du graphe $\mathbf{L}$.
2. Définition des Coordonnées : Définir les caractéristiques spectrales pour chaque nœud $i$ comme un vecteur $\mathbf{r}_i$ composé des valeurs du nœud dans ces vecteurs propres (par exemple, $\mathbf{r}_i = [\mathbf{u}_k[i]]_{k=0}^{m-1}$).
3. Transformation Rotative : Appliquer la rotation standard du RoPE au vecteur de requête ($\mathbf{q}_i$) et de clé ($\mathbf{k}_i$) en utilisant ces coordonnées spectrales :
   $$\mathbf{z}_i \to \text{RoPE}(\mathbf{r}_i)\mathbf{z}_i$$
   où $\text{RoPE}(\mathbf{r}_i)$ fait pivoter les blocs 2D du vecteur de jeton par des angles déterminés par les coordonnées spectrales.

Propriétés Théoriques

Le document établit des garanties théoriques pour WIRE :

• Équivariance par Permutation : WIRE est équivariant sous le réordonnancement des nœuds (à l'exception des inversions de signe et des rotations dans les sous-espaces dégénérés), une propriété nécessaire pour l'apprentissage sur graphes.
• Récupération du RoPE Standard : Sur les graphes de grille (produits cartésiens de graphes de chemins), WIRE récupère le comportement du RoPE standard utilisé dans les LLM et les ViT. Les coordonnées spectrales sur les grilles correspondent de manière monotone aux coordonnées spatiales, les normalisant efficacement dans l'intervalle $(-1, 1)$.
• Dépendance à la Résistance Effective : Théoriquement, lorsque les fréquences de WIRE sont échantillonnées de manière aléatoire, la transformation attendue réduit les scores d'attention requête-clé proportionnellement à la résistance effective entre les nœuds. La résistance effective est une métrique qui sert de borne inférieure pour la distance du plus court chemin.
• Compatibilité avec l'Attention Linéaire : Contrairement aux RPE basés sur des biais, WIRE modifie directement les jetons (requêtes et clés) plutôt que les logits d'attention. Cela permet d'utiliser WIRE avec des mécanismes d'attention linéaire (par exemple, les Performers) sans instancier la matrice d'attention complète, préservant ainsi une complexité de $\mathcal{O}(N)$.

Contributions Clés

1. Introduction de WIRE : Un nouvel encodage de position pour les graphes qui généralise le RoPE en utilisant les vecteurs propres du laplacien du graphe comme coordonnées positionnelles.
2. Analyse Théorique :
   • Preuve que WIRE récupère le RoPE standard sur les graphes de grille.
   • Démonstration que WIRE module implicitement l'attention en fonction de la résistance effective du graphe, offrant ainsi un biais inductif topologique fondé.
   • Établissement de la compatibilité avec l'attention linéaire, une caractéristique auparavant indisponible pour les encodages de position relatifs basés sur la résistance.
3. Validation Empirique : Expérimentations approfondies sur des tâches synthétiques, des nuages de points et des benchmarks de graphes standards.

Résultats Expérimentaux

Les auteurs ont entraîné plus de 200 modèles de transformers pour évaluer WIRE :

• Tâches Synthétiques :
  • Sous-graphes monochromatiques : WIRE a considérablement réduit l'erreur quadratique moyenne (RMSE) de test par rapport aux bases sans WIRE, particulièrement lorsque la topologie du graphe devient complexe (suppression d'arêtes dans les grilles).
  • Prédiction du plus court chemin : WIRE a presque divisé par deux la RMSE de test sur les graphes de Watts-Strogatz par rapport aux bases utilisant uniquement des APE, en tirant parti de sa dépendance à la résistance effective.
• Transformers de Nuages de Points : Appliqué à ModelNet40 (classification) et ShapeNet (segmentation). WIRE a apporté des gains constants par rapport aux bases PCT standard et a égalé ou surpassé le RoPE cartésien, tout en maintenant l'invariance SE(3).
• Benchmarks de Graphes (Tâches GNN) :
  • Évalué sur des jeux de données incluant MNIST, CIFAR-10, PATTERN, CLUSTER et des jeux de données OGB (ex: ogbg-molhiv).
  • En utilisant des Performers (attention linéaire), WIRE a systématiquement amélioré les performances par rapport aux modèles de base.
  • Notamment, WIRE a souvent réduit l'écart de performance entre les modèles d'attention linéaire efficaces et les modèles de softmax à rang complet coûteux. Par exemple, sur MalNet-Tiny, les modèles dotés de WIRE (Performers) ont atteint une précision comparable aux transformers à rang complet, mais pouvaient être entraînés sur un seul GPU T4.
• Efficacité : Les paramètres de WIRE constituent moins de 1 % du total des paramètres du modèle. La méthode est compatible avec divers mécanismes d'attention, notamment softmax, linéaire (Performer), SGFormer et BigBird.

Signification et Revendications

Le papier affirme que WIRE parvient à combler le fossé entre le succès de RoPE dans les données de séquences et d'images et les besoins des données structurées en graphes. Sa principale importance réside dans :

1. Scalabilité : Elle permet l'utilisation de biais inductifs structurels forts (comme la résistance effective) dans les modèles d'attention linéaire, ce qui était auparavant impossible avec les RPE basés sur des biais.
2. Performance : Elle augmente la précision des transformers sur les tâches de graphes, permettant souvent aux modèles linéaires efficaces de s'aligner sur les performances des modèles basés sur le softmax beaucoup plus coûteux.
3. Fondement Théorique : Elle fournit une manière mathématiquement fondée d'injecter la topologie du graphe dans les mécanismes d'attention, récupérant les comportements connus sur les grilles tout en offrant de nouvelles perspectives sur la résistance effective sur les graphes généraux.

Les auteurs concluent que WIRE est un ajout "prêt à l'emploi" (drop-in) aux architectures de graph transformers existantes, offrant un biais inductif structurel léger et efficace, et ouvrant la voie à la recherche sur la généralisation des encodages rotatifs à d'autres modalités de données.