Résumé technique : L'attention RoPE peut être entraînée en temps presque linéaire

Définition du problème

Le mécanisme de l'encodage de position rotatif (RoPE - Rotary Position Embedding) est devenu un composant standard des modèles de langage de grande taille (LLM) de pointe tels que Llama, Claude et les modèles d'Apple, offrant une expressivité supérieure pour capturer les relations entre les jetons par rapport aux encodages positionnels traditionnels. Cependant, les rotations dépendantes de la position inhérentes à RoPE compliquent le calcul du mécanisme d'attention.

Alors que des travaux récents ([AS24a]) ont établi un algorithme en temps presque linéaire ($n^{1+o(1)}$) pour le calcul de la passe avant (forward) de l'attention RoPE sous le régime des « entrées bornées » (où les entrées de la matrice sont limitées par un paramètre $B$), le calcul de la passe arrière (backward, calcul du gradient pour l'entraînement) est resté non traité. Le calcul de la passe arrière est intrinsèquement plus complexe car il implique des transformations non linéaires de la matrice d'attention et des encodages positionnels. La question centrale abordée par ce travail est de savoir si le calcul du gradient de la passe arrière pour l'attention RoPE peut atteindre la même efficacité temporelle presque linéaire que le calcul de la passe avant sous des conditions d'entrées bornées.

Méthodologie

Les auteurs développent le premier algorithme de calcul de la passe arrière de l'attention RoPE qui s'exécute en temps presque linéaire. L'approche repose sur une combinaison de dérivation de gradient sous forme fermée, d'approximation de bas rang, de méthodes polynomiales et de la transformée de Fourier rapide (FFT).

1. Reformulation du gradient en forme fermée

Le papier dérive d'abord une expression en forme fermée pour le gradient de la fonction de perte de l'attention RoPE par rapport aux matrices de poids. En utilisant l'astuce tensorielle (produits de Kronecker) et en reformulant la matrice d'attention $A(X)$, le gradient est exprimé comme :
$$\frac{d\text{Loss}(x)}{dx} = \tilde{A}^\top \text{vec}(\gamma(x))$$
où $\gamma(x)$ est une fonction matricielle complexe impliquant :

• $s(x)$ : Le vecteur Softmax normalisé.
• $\ell(x)$ : Un terme d'erreur dérivé de la différence entre la sortie de l'attention et la cible.
• $\beta(x)$ : Un terme combinant l'erreur et la matrice de valeur.
• $\gamma(x)$ : Un terme impliquant la diagonale de $s(x)$ et le produit extérieur $s(x)s(x)^\top$ agissant sur $\beta(x)$.

2. Stratégie d'approximation de bas rang

Pour atteindre une complexité temporelle presque linéaire, les auteurs approximent les composantes de $\gamma(x)$ à l'aide de matrices de bas rang. La stratégie consiste à décomposer $\gamma(x)$ en deux parties, $\gamma_1(x)$ et $\gamma_2(x)$, et à approximer chacune séparément :

• Approximation de $s(x)$ et $\ell(x)$ : En s'appuyant sur l'algorithme de la passe avant de [AS24a], les auteurs montrent que le Softmax normalisé $s(x)$ peut être approximé par des matrices de bas rang $U_1 V_1^\top$ en temps $n^{1+o(1)}$. Le terme d'erreur $\ell(x)$ est ensuite approximé en utilisant ce résultat.
• Approximation de $\beta(x)$ : Puisque $\beta(x)$ est un produit impliquant la matrice de valeur et le terme d'erreur, il est approximé en construisant des facteurs de bas rang basés sur les approximations de ses composantes.
• Approximation de $\gamma(x)$ :
  • $\gamma_1(x) = \text{diag}(s(x))\beta(x)$ est approximé en combinant les facteurs de bas rang de $s(x)$ et $\beta(x)$ en utilisant des produits de Kronecker par ligne.
  • $\gamma_2(x) = s(x)s(x)^\top\beta(x)$ est approximé en précalculant des termes intermédiaires et en utilisant la structure de bas rang de $s(x)$ et $\beta(x)$.

3. Analyse de dureté

Pour établir la nécessité de la condition d'entrée bornée, les auteurs dérivent des bornes inférieures basées sur l'hypothèse de l'exponentielle forte du temps (SETH). Ils prouvent que si la borne d'entrée $B$ dépasse un certain seuil (spécifiquement $B = \omega(\sqrt{\log n})$), aucun algorithme ne peut calculer le gradient en temps sous-quadratique ($O(n^{2-q})$) en supposant SETH. Cela confirme que l'hypothèse d'entrée bornée n'est pas seulement une commodité technique, mais un prérequis fondamental pour une performance sous-quadratique.

Contributions clés

1. Gradient en forme fermée : Le papier fournit la première formulation en forme fermée du gradient de l'attention RoPE (Lemme 4.1) et analyse sa complexité temporelle exacte, identifiant le goulot d'étranglement quadratique dans le calcul naïf.
2. Algorithme en temps presque linéaire : Les auteurs présentent le premier algorithme pour approximer le gradient de la passe arrière de l'attention RoPE en $n^{1+o(1)}$ sous des conditions d'entrées bornées (Théorème 5.7). Cela correspond à l'efficacité de la passe avant.
3. Bornes inférieures théoriques : Le travail établit que la condition d'entrée bornée est nécessaire pour une performance sous-quadratique, fournissant un résultat de dureté dérivé de SETH (Théorème 6.1).
4. Techniques algorithmiques : L'approche intègre des méthodes d'approximation polynomiale et la FFT avec des techniques d'approximation de bas rang spécifiquement adaptées aux contraintes structurelles de RoPE.

Résultats

Le résultat principal (Théorème 5.7) démontre que pour des paramètres $d = O(\log n)$ et $B = o(\sqrt{\log n})$, il existe un algorithme pour résoudre le problème de calcul du gradient de l'attention RoPE avec une erreur additive bornée par $1/\text{poly}(n)$ en temps $n^{1+o(1)}$.

Inversement, le résultat de dureté (Théorème 6.1) montre que si $B = \omega(\sqrt{\log n})$, le calcul du gradient en un temps $O(n^{2-q})$ est impossible sous l'hypothèse SETH.

Signification

Ce travail comble une lacune critique dans la compréhension théorique des Transformers basés sur RoPE. En prouvant que le calcul de la passe arrière peut être aussi efficace que celui de la passe avant sous des entrées bornées, l'article lève une barrière de calcul importante pour l'entraînement de modèles à grande échelle utilisant RoPE. Les conclusions suggèrent que l'efficacité de l'entraînement des modèles basés sur RoPE est théoriquement comparable à celle des modèles utilisant l'attention standard, à condition que le régime d'entrées bornées soit respecté.

L'article caractérise la complexité fine des calculs de la passe arrière de RoPE, étendant les résultats antérieurs sur les calculs de la passe avant. Il souligne l'interaction entre la conception d'algorithmes et la théorie de la complexité computationnelle, offrant une base pour de futures recherches sur les calculs de sous-gradients pour d'autres variantes avancées de l'attention et des mécanismes d'encodage de position. Les auteurs notent que des travaux futurs pourraient explorer les cas d'entrées non bornées et les implications pratiques de ces bornes théoriques pour l'entraînement des LLM réels.