Composing Linear Layers from Irreducibles

Résumé Technique : Composition de couches linéaires à partir d'irréductibles

Énoncé du Problème
Les modèles de langage de grande taille (LLM) contemporains présentent des comportements suggérant qu'ils s'appuient sur des primitives de bas niveau qui se composent en modules dotés d'une fonctionnalité plus riche. Cependant, les blocs fondamentaux de ces transformations linéaires restent mal compris. Alors que des travaux antérieurs ont exploré des approximations économes en paramètres (par exemple, l'adaptation de faible rang ou le Block-Hadamard) et des compositions modulaires (comme le Mixture-of-Experts), il manque une compréhension de la composition mécaniste des primitives de bas niveau qui génèrent des comportements linéaires complexes. Les auteurs posent un problème concret : la fonctionnalité d'une couche linéaire, qui contient typiquement des millions de paramètres ($O(d^2)$), peut-elle être synthétisée à partir d'un ensemble minimal d'objets géométriques irréductibles ?

Méthodologie
Les auteurs proposent un cadre algébrique utilisant l'algèbre de Clifford pour décomposer et reconstruire des transformations linéaires.

1. Formulation Algébrique :

   • Les applications linéaires standards sont traditionnellement représentées par des matrices denses. Les auteurs expriment plutôt les transformations linéaires comme des compositions de bivecteurs (objets géométriques codant des plans orientés) au sein de l'algèbre de Clifford $Cl(n)$.
   • Ils utilisent le groupe de Spin $Spin(n)$, un revêtement double du groupe spécial orthogonal $SO(n)$. Les éléments de ce groupe, appelés rotors ($r$), agissent sur les multivecteurs via le « produit sandwich » $x \mapsto rxr^\dagger$.
   • Bien que les multivecteurs arbitraires nécessitent $2^n$ paramètres, les auteurs restreignent la construction aux rotors générés par des bivecteurs. Un rotor $r$ est l'exponentielle d'un bivecteur $b$ : $r = \exp(b)$.

2. Efficacité des Paramètres :

   • Une matrice dense $d \times d$ nécessite $O(d^2)$ paramètres.
   • En représentant la transformation comme une composition de rotors agissant sur des sous-espaces, le nombre de paramètres scalaires requis passe à l'échelle $O(\log_2 d)$. Spécifiquement, pour une largeur $w$, le nombre de paramètres est approximativement $w \binom{\log_2 d}{2}$. Pour $d=2048$, cela réduit des millions de paramètres à quelques centaines.

3. Algorithme de Décomposition Différentiable :

   • L'application exponentielle pour des bivecteurs généraux implique une série infinie, ce qui est problématique sur le plan computationnel. Cependant, pour les bivecteurs simples (ceux représentant un plan de rotation unique), l'application exponentielle possède une solution en forme close : $\exp(b) = \cos(\|b\|) + \sin(\|b\|)\frac{b}{\|b\|}$.
   • Pour gérer les bivecteurs généraux, les auteurs emploient une décomposition invariante (basée sur Roelfs et Keninck [2021]), qui exprime tout bivecteur comme une somme d'au plus $k = \lfloor n/2 \rfloor$ bivecteurs simples mutuellement commutants et orthogonaux.
   • Ils introduisent un algorithme de décomposition invariante différentiable (Algorithmes 1 et 2) qui extrait de manière itérative ces bivecteurs simples en utilisant une méthode d'itération de puissance inspirée des sous-espaces de Krylov, adaptée à l'algèbre de Clifford. Cela évite la décomposition en valeurs propres explicite, garantissant la stabilité numérique et la différentiabilité pour l'optimisation basée sur le gradient.

4. Gadget de Rotor Généralisé :

   • Pour gérer des dimensions d'entrée/sortie arbitraires ($d_{in}, d_{out}$) qui ne sont pas des puissances de 2, les auteurs construisent une transformation généralisée $\psi$ en agrégeant plusieurs modules de produit sandwich de rotors opérant sur différents sous-espaces, suivis d'un opérateur de pooling.

Contributions Clés

1. Représentation Algébrique : Les auteurs expriment les transformations linéaires comme des compositions d'objets géométriques primitifs (bivecteurs) en utilisant des produits sandwich de rotors, réduisant la complexité des paramètres de $O(d^2)$ à $O(\log_2 d)$.
2. Décomposition Différentiable : Ils proposent un algorithme en forme close et différentiable pour mapper des bivecteurs arbitraires à leurs rotors correspondants via une décomposition invariante, permettant l'intégration avec les frameworks d'autograd.
3. Validation Empirique : La méthode est appliquée aux projections Key, Query et Value des couches d'attention des LLM (Key, Query et Value) et entraînée de bout en bout sur FMNIST.

Résultats

• Efficacité des Paramètres : Dans les expériences avec LLaMa-3.2 1B et Qwen-2.5 1.5B, les couches de rotor ont atteint des réductions massives de paramètres. Par exemple, la projection de requête dans LLaMa-3.2 1B est passée d'environ 4,19 M de paramètres (dense) à $\le 896$ paramètres (rotor), soit une réduction d'environ 4700x.
• Performance : Malgré cette compression extrême, les couches basées sur les rotors ont égalé ou surpassé de solides bases de comparaison, incluant les approximations de faible rang (LR1, LR4) et Block-Hadamard (BH1).
  • Sur Wikitext2 (LLaMa-3.2 1B), remplacer une seule couche d'attention a donné un log-PPL de 2,629 pour les rotors, contre 2,636 pour la meilleure base (BH1).
  • Sur Arc Challenge, les rotors ont atteint 55,31 % de précision (LLaMa-3.2 1B), surpassant BH1 (54,83 %) et LR4 (53,84 %).
• Robustesse : La performance est restée constante à travers diverses combinaisons de couches (remplacement de 1, 2 ou 3 couches) et divers jeux de données (Wikitext2, C4, PTB, Arc, HellaSwag).
• Entraînement de Bout en Bout : Lorsqu'ils sont entraînés à partir de zéro sur FMNIST, le MLP basé sur les rotors a convergé vers une précision finale (88,36 %) comparable à la base dense (89,67 %), bien que la convergence ait été légèrement plus lente.

Signification et Revendications
Le papier prétend démontrer la faisabilité de la synthèse de couches linéaires à partir d'un ensemble compact de blocs de construction géométriques. Les auteurs soulignent que leur objectif n'est pas simplement de fournir un remplacement direct des couches denses (ce qui souffre actuellement de temps d'inférence plus longs en raison de l'absence de noyaux optimisés pour le matériel), mais de fournir une perspective algébrique sur la manière dont les primitives de bas niveau se composent en fonctions de plus haut niveau.

Les résultats suggèrent que la fonctionnalité complexe des couches linéaires dans les modèles profonds peut être réalisée par une composition structurée de quelques centaines d'objets géométriques (rotors) plutôt que par des millions de poids indépendants. Cela offre de nouvelles avenues pour :

• L'Interprétabilité : Comprendre les couches linéaires comme des compositions de plans orientés.
• L'Efficacité des Paramètres : Réduire drastiquement le nombre de paramètres tout en maintenant la performance.
• Le Futur du Matériel : Potentiel de troquer la puissance de calcul contre la bande passante mémoire en synthétisant les poids sur puce à partir de paramètres géométriques, répondant ainsi aux goulots d'étranglement de la bande passée mémoire lors de l'inférence des grands modèles.

Les auteurs restent modestes, notant que les bénéfices pratiques nécessiteront une intégration au niveau du système et des implémentations adaptées au matériel, ce qui dépasse le cadre de ce travail.