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

Résumé Technique : Découpler le « Quoi » et le « Où » avec l'Embedding de Position en Coordonnées Polaires

Énoncé du Problème
Dans les architectures Transformer, le mécanisme d'attention associe les clés aux requêtes en se basant sur deux facteurs : le contenu (« quoi ») et la position dans la séquence (« où »). Les auteurs soutiennent que l'Embedding de Position Rotatif (RoPE), largement adopté dans des modèles de pointe comme Llama 3 et DeepSeek, enchevêtre ces deux facteurs. Dans RoPE, la rotation appliquée aux vecteurs de clé et de requête dépend à la fois de l'amplitude de leur contenu et de leurs positions relatives. Cet enchevêtrement crée un amalgame « quoi-où », nuisant aux performances du modèle lorsque les tâches nécessitent une correspondance indépendante sur le contenu ou la position seuls. Plus précisément, le terme d'interaction dans le score d'attention de RoPE fait que la phase de la rotation est influencée par les caractéristiques de contenu, ce qui rend difficile pour le modèle l'apprentissage de règles dépendant strictement de l'indexation positionnelle ou de la correspondance de contenu.

Méthodologie : Embedding de Position en Coordonnées Polaires (PoPE)
Les auteurs proposent PoPE, une modification de RoPE qui découple le contenu et la position en reformulant les représentations de clé et de requête en coordonnées polaires.

• Reformulation : Dans RoPE, les composantes sont rotées en fonction de la position, mais la phase initiale est déterminée par le contenu. Dans PoPE, les auteurs transforment les vecteurs de clé et de requête à valeurs réelles en vecteurs complexes où :
  • L'amplitude ($\mu$) est dérivée uniquement des caractéristiques de contenu en utilisant une fonction d'activation softplus ($\sigma(x) = \ln(1 + e^x)$) pour garantir la non-négativité.
  • La phase ($\phi$) est déterminée strictement par la position dans la séquence ($s$ ou $t$) et une fréquence spécifique à la composante ($\theta_c$).
• Score d'Attention : Le score d'attention est calculé comme la partie réelle du produit scalaire entre la transposée conjuguée de la requête et la clé. Mathématiquement, cela donne :
  $$a^{PoPE}_{ts} = \sum_{c=1}^{d} \mu_{\tilde{q}_{tc}} \mu_{\tilde{k}_{sc}} \cos((s - t)\theta_c + \delta_c)$$
  où $\delta_c$ est un biais apprenable.
• Différences Clés avec RoPE :
  1. Découplage : Le terme d'interaction ($\phi_{k_{sc}} - \phi_{q_{tc}}$) qui mélange le contenu et la position dans RoPE est éliminé. L'amplitude (contenu) et la phase (position) sont strictement séparées.
  2. Résolution de Fréquence : PoPE opère sur $d$ composantes (éléments individuels) plutôt que sur $d/2$ paires, doublant ainsi efficacement le nombre de canaux de fréquence.
  3. Biais Apprenable : Un biais fixe mais apprenable $\delta_c$ peut être introduit pour ajuster l'offset relatif optimal pour chaque fréquence sans réintroduire le amalgame contenu-position.

Principales Contributions

1. Tâche de Diagnostic (Indexation Indirecte) : Les auteurs introduisent une tâche exigeant que les modèles identifient un caractère cible basé sur un décalage relatif par rapport à un caractère source (ex: « trouvez le caractère 3 positions à gauche de 'z' »). Cette tâche nécessite une manipulation indépendante du contenu et de la position.
2. Aperçu Théorique : L'article fournit une dérivation algébrique montrant comment RoPE amalgame le contenu et la position, tandis que PoPE permet au mécanisme d'attention d'implémenter des règles comme une conjonction de correspondances indépendantes de « quoi » et de « où ».
3. Implémentation Efficace : Une implémentation efficace utilisant Triton et Flash Attention 2 est fournie, laquelle calcule la partie réelle du produit scalaire complexe sans matérialiser la matrice complexe complète, n'ajoutant qu'un seul surcoût de multiplication par rapport à l'attention standard.

Résultats

• Indexation Indirecte : Sur la tâche de diagnostic, RoPE a atteint une précision moyenne de seulement 11,16 %, peinant à apprendre la correspondance indépendante requise. En revanche, PoPE a atteint 94,82 % de précision, démontrant sa capacité à découpler les facteurs.
• Modélisation de Séquences (Musique et Génomique) :
  • Musique (JSB et MAESTRO) : PoPE surpasse RoPE en termes de log-vraisemblance négative (NLL) sur les deux jeux de données, atteignant une perplexité plus faible.
  • Génomique (Génome de Référence Humain) : PoPE montre une baisse significative de la NLL par rapport à la ligne de base Ro_PE.
• Modélisation de Langage (OpenWebText) : À travers différentes échelles de modèles (124M, 253M et 774M de paramètres), PoPE obtient systématiquement une perplexité plus faible que RoPE. L'écart de performance reste stable ou augmente avec la taille du modèle.
• Tâches Aval : Les évaluations zero-shot sur six tâches aval (LAMBADA, BLiMP, CBT, HellaSwag, PIQA, ARC-E) ont montré que les modèles basés sur PoPE atteignent une précision moyenne plus élevée sur toutes les tailles de modèles par rapport aux modèles basés sur RoPE.
• Extrapolation de Longueur :
  • Zero-Shot : PoPE démontre de fortes capacités d'extrapolation de longueur zero-shot sur des séquences jusqu'à 10 fois plus longues que le contexte d'entraînement (1024 vs 10240 jetons) sans aucune interpolation de fréquence ou de fine-tuning.
  • Comparaison avec YaRN : Bien que YaRN (une méthode conçue pour l'extrapolation) atténue la dégradation de RoPE jusqu'à sa limite de fine-tuning (4064 jetons), il se dégrade considérablement au-delà de celle-ci. PoPE maintient ses performances sur des séquences plus longues sans nécessiter les techniques d'interpolation ou le fine-tuning extensif que YaRN impose.
  • Analyse de l'Utilisation des Fréquences : L'analyse des normes de caractéristiques a révélé que les modèles RoPE ont tendance à supprimer les canaux de haute fréquence pour éviter le bruit, tandis que les modèles PoPE utilisent une plage de fréquences plus distribuée, incluant les hautes fréquences, à travers les couches.

Signification et Revendications
L'article affirme que PoPE offre une solution fondée sur des principes au problème de l'enchevêtrement « quoi-où » inhérent à RoPE. En découplant le contenu et la position, PoPE améliore l'efficacité des données, la précision asymptotique et, surtout, la généralisation de la longueur du contexte. Les auteurs soulignent que ces gains sont obtenus avec une modification mineure de l'architecture RoPE existante, ne nécessitant aucun changement architectural complexe.

Le travail est présenté comme une recherche fondamentale sur la modélisation de séquences. Les auteurs notent que si la méthode améliore les capacités dans des domaines comme la musique, la génomique et le langage naturel (pouvant aider des applications telles que le repliement des protéines ou la découverte de médicaments), elle comporte des risques de double usage communs aux modèles de fondation, tels que la génération de contenus malveillants ou de deepfakes sophistiqués. La contribution principale est la démonstration que la séparation des informations de contenu et de position conduit à des modèles Transformer plus robustes et plus généralisables.