PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training

Résumé Technique : PRISM

Énoncé du Problème

Les fonctions de matrice, telles que les racines carrées, les racines inverses et l'orthogonalisation (décomposition polaire), sont fondamentales pour les méthodes de gradient préconditionnées dans l'entraînement des réseaux de neurones (ex. : Shampoo, Muon). Bien que les algorithmes itératifs reposant sur des multiplications de matrices générales (GEMM) soient préférés pour l'accélération sur GPU par rapport aux méthodes nécessitant une décomposition en valeurs singulières (SVD) explicite ou une estimation des valeurs propres, les approches de pointe existantes souffrent de limitations significatives :

1. Manque d'adaptativité : Les méthodes récentes (ex. : PolarExpress, CANS) reposent souvent sur des mises à jour polynomiales à ordre fixe optimisées pour un intervalle spectral prédéfini (ex. : $[\ell, u]$). Si la distribution réelle des valeurs singulières de la matrice d'entrée dévie de cet intervalle, la convergence se dégrade considérablement.
2. Dépendance aux connaissances préalables : Ces méthodes nécessitent généralement une connaissance préalable de la distribution des valeurs singulières ou des bornes serrées sur les plus petites et plus grandes valeurs singulières, qui sont souvent aussi coûteuses à estimer que la fonction de matrice elle-même.
3. Portée étroite : De nombreuses techniques d'accélération récentes sont adaptées à des cas d'utilisation spécifiques (ex. : uniquement la décomposition polaire dans l'optimiseur Muon) plutôt que de fournir un cadre général pour diverses fonctions de matrice.

Méthodologie : Le Cadre PRISM

Les auteurs proposent PRISM (Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation), un méta-algorithme général conçu pour accélérer les calculs de fonctions de matrice itératives sans nécessiter de bornes spectrales explicites. PRISM opère en deux parties distinctes :

Partie I : Configuration de Base (Approximation Polynomiale Itérative)

PRISM cadre le calcul d'une fonction de matrice $T(A)$ comme un problème d'approximation polynomiale itérative.

• Pour un scalaire $x$ approximant $T(a)$, le problème est reformulé comme l'approximation d'une fonction résiduelle $f(\xi)$ où $\xi$ mesure la distance à la solution.
• Les méthodes itératives standards (comme Newton-Schulz) sont dérivées en utilisant un polynôme de Taylor d'ordre fixe $f_d(\xi)$ autour de $\xi=0$.
• L'itération de la matrice prend la forme $X_{k+1} = X_k f_d(R_k)$, où $R_k$ est la matrice résiduelle (ex. : $I - X_k^2$ pour les fonctions de signe).

Partie II : Accélération (Ajustement Adaptatif et Sketching)

Pour surmonter la lenteur de la convergence initiale des polynômes de Taylor fixes lorsque le spectre de $R_k$ n'est pas proche de 0, PRISM introduit deux mécanismes :

1. Ajustement Polynomial Adaptatif :
   Au lieu d'utiliser un polynôme de Taylor fixe, PRISM remplace le coefficient de plus haut degré par un paramètre apprenable $\alpha$. Le polynôme est mis à jour à chaque itération pour s'ajuster au spectre actuel de la matrice résiduelle $R_k$.

   • Le polynôme de mise à jour est défini par $g_d(\xi; \alpha) = f_{d-1}(\xi) + \alpha \xi^d$.
   • L' $\alpha_k^*$ optimal est trouvé en minimisant la norme résiduelle $\|I - X_{k+1}^2\|_F^2$ (ou équivalent pour d'autres fonctions) par rapport à $\alpha$. Cela revient à ajuster le polynôme aux valeurs propres de $R_k$ dans un sens de moindres carrés.
   • Pour garantir la stabilité, $\alpha_k$ est contraint à un intervalle fixe $[\ell, u]$ (ex. : $[1/2, 1]$ pour $d=1$), ce qui garantit des propriétés de convergence indépendantes du spectre d'entrée.

2. Sketching Randomisé :
   Le calcul de l' $\alpha_k^*$ exact nécessite l'évaluation des traces de puissances élevées de la matrice résiduelle, ce qui est coûteux en calcul ($O(n^3)$ ou plus).

   • PRISM emploie le sketching randomisé (utilisant des plongements de sous-espaces obligatoires, ex. : matrices gaussiennes) pour approximer le problème de moindres carrés.
   • En résolvant le problème de minimisation sur une version "sketchée" du résidu, le coût de la recherche de $\alpha_k$ est réduit à $O(n^2 \log n)$ (ou $O(n^2 p)$ où $p \ll n$), ce qui est négligeable par rapport au coût $O(n^3)$ des multiplications de matrices dans l'itération principale.
   • Les garanties théoriques (Théorème 2) montrent qu'avec une dimension de sketch $p = O(\log n)$, le taux de convergence reste essentiellement identique à celui de la minimisation exacte.

Contributions Clés

1. Cadre d'Accélération Général : PRISM fournit une approche unifiée applicable à une large gamme d'algorithmes itératifs, incluant Newton-Schulz (pour les racines carrées, racines inverses, le signe et la décomposition polaire), les méthodes de Chebyshev et les méthodes de Newton inverse.
2. Adaptativité Indépendante de la Distribution : Contrairement aux travaux antérieurs, PRISM s'adapte dynamiquement au spectre évolutif de l'itération actuelle sans nécessiter de connaissance préalable des bornes ou des distributions des valeurs singulières.
3. Implémentation Efficace : L'intégration du sketching randomisé garantit que le surcoût de l'ajustement polynomial adaptatif est minimal, préservant la nature de l'itération dominée par les GEMM, favorable aux GPU.
4. Garanties Théoriques : L'article fournit des preuves de convergence montrant que les itérations accélérées par PRISM maintiennent des taux de convergence quadratiques (ou d'ordre supérieur) même avec un sketching approximatif, à condition que la dimension du sketch soit suffisamment grande.

Résultats Expérimentaux

Les auteurs valident PRISM à travers des benchmarks synthétiques et son intégration dans des optimiseurs de réseaux de neurones :

• Vitesse de Convergence : Sur des matrices avec des rapports d'aspect et des distributions spectrales variés (incluant des matrices aléatoires gaussiennes et des distributions de Marchenko-Pastur à queue lourde imitant des modèles pré-entraînés), PRISM surpasse systématiquement les méthodes Newton-Schulz classiques et les méthodes à intervalle fixe comme PolarExpress.
  • PolarExpress ralentit souvent ou échoue à converger lorsque la plage réelle des valeurs singulières dévie de son intervalle optimisé.
  • PRISM maintient des accélérations stables sur toute la plage de $\sigma_{min}$, en adaptant automatiquement le coefficient $\alpha_k$.
• Entraînement de Réseaux de Neurones :
  • Optimiseur Shampoo : L'intégration de PRISM pour calculer les préconditionneurs de racine carrée inverse pour ResNet-20/32 sur les jeux de données CIFAR a entraîné des temps d'entraînement plus rapides par rapport aux références de décomposition en valeurs propres et PolarExpress.
  • Optimiseur Muon : L'application de PRISM à la décomposition polaire des matrices de gradient lors de l'entraînement de GPT-2 (jeu de données FineWeb) a démontré une amélioration des pertes de validation et de l'efficacité de l'entraînement par rapport à AdamW, PolarExpress et les variantes de Newton-Schulz à degré fixe.

Signification et Revendications

L'article affirme que PRISM offre une voie pratique et générale vers un calcul de fonction de matrice adaptatif à l'instance et compatible avec les GPU. Sa principale importance réside dans la transformation de l'« adaptativité spectrale » — auparavant une source de charge de réglage nécessitant des bornes spectrales manuelles — en un outil algorithmique fiable et automatisé. En découplant l'accélération de la connaissance spectrale préalable, PRISM permet des performances robustes à travers divers régimes spectraux communs en apprentissage automatique (ex. : distributions à queue lourde dans les modèles pré-entraînés), ce qui en fait un outil polyvalent pour accélérer l'entraînement de réseaux de neurones à grande échelle.