Unlocking Feature Learning in Gated Delta Networks at Scale

Résumé Technique : Débloquer l'apprentissage de caractéristiques dans les réseaux Delta à porte (Gated Delta Networks) à grande échelle

Énoncé du Problème
La mise à l'échelle des grands modèles de langage (LLM) fait face à deux goulots d'étranglement principaux : le coût computationnel élevé de l'entraînement de modèles plus larges et la complexité quadratique de l'auto-attention standard des Transformers. Les architectures récurrentes linéaires, telles que les réseaux Delta à porte (Gated Delta Networks - GDN), sont apparues comme des alternatives efficaces qui maintiennent une inférence en temps linéaire tout en capturant des dépendances à longue portée. Cependant, un fossé critique subsiste dans la capacité à mettre ces modèles à l'échelle efficacement. Bien que la paramétrisation par mise à jour maximale (Maximal Update Parametrization - µP) ait permis avec succès le transfert de zéro-shot des hyperparamètres (spécifiquement les taux d'apprentissage) pour les Transformers standards et certains modèles linéaires, son application aux réseaux Delta à porte reste inexplorée. Le défi central réside dans la dynamique de l'état récurrent de l'architecture, qui est mise à jour via des mécanismes de porte dépendants des données et des règles de delta à matrice complète, plutôt que par les récurrences scalaires ou les structures de feedforward standards supposées par les dérivations existantes de µP. Des travaux antérieurs sur les modèles d'espace d'état structurés (SSM) comme Mamba ont montré que le µP classique échoue pour les SSM diagonaux, mais l'utilisation par les GDN de mises à jour d'état par matrice complète rend ces corrections spécifiques inapplicables.

Méthodologie
Les auteurs comblent ce fossé en dérivant rigoureusement la formulation µP pour les réseaux Delta à porte. La méthodologie consiste à propager les estimations de taille de coordonnées à travers l'ensemble de la passe avant, incluant les mécanismes de porte uniques et les transitions d'état récurrentes, afin de déterminer les règles de mise à l'échelle nécessaires pour l'initialisation, les multiplicateurs de passe avant et les taux d'apprentissage.

Étapes analytiques clés :

1. Analyse de la passe avant : Les auteurs dérivent les tailles de coordonnées des caractéristiques projetées, de l'état latent $S_t$, et de la lecture (readout). Ils identifient que la normalisation L2 des requêtes (queries) et des clés (keys), combinée aux mises à jour d'écriture de rang un, entraîne des mises à l'échelle de variance spécifiques. Crucialement, ils proposent d'insérer un multiplicateur $\sqrt{d}$ avant la couche RMSNorm pour garantir que l'entrée de la couche de normalisation maintienne une taille de coordonnée $\Theta(1)$, empêchant ainsi la perturbation de la mise à l'échelle du gradient requise pour l'apprentissage de caractéristiques.
2. Analyse du mécanisme de porte : Le papier analyse les paramètres de porte non linéaires $\alpha_t$ (décroissance) et $\beta_t$ (force d'écriture). Contrairement aux poids standards, ceux-ci sont dérivés de pré-activations scalaires impliquant des vecteurs de poids entraînables ($W_\alpha, W_\beta$) et des paramètres scalaires ($a_{log}, b$). L'analyse révèle que ces composants de porte ne suivent pas les lois de mise à l'échelle µP standard en raison de leurs rôles spécifiques dans l'équation de mise à jour de l'état.
3. Passe arrière et dérivation du taux d'apprentissage :
   • AdamW : En raison de la normalisation des gradients par le second moment effectuée par AdamW, l'amplitude effective de la mise à jour est cohérente entre les différentes classes de poids, simplifiant les exigences de taux d'apprentissage.
   • SGD : Sous un SGD pur, les amplitudes de gradient brutes varient considérablement selon les classes de poids. Les auteurs dérivent que les matrices de poids de porte ($W_\alpha, W_\beta$) nécessitent une mise à l'échelle du taux d'apprentissage de $\Theta(1/\sqrt{d})$, et que les paramètres de porte scalaires ($a_{log}, b$) nécessitent une mise à l'échelle de $\Theta(\sqrt{d})$. Ces mises à l'échelle divergent du cadre µP standard où les taux d'apprentissage sont généralement uniformes ou suivent une règle simple en $1/n$.

Contributions Clés

1. Dérivation Théorique : Le papier fournit la première formulation µP complète pour les réseaux Delta à porte. Il dérive les estimations de taille de coordonnées à travers toute la passe avant et établit des variances d'initialisation et des multiplicateurs de passe avant fondés sur des principes pour toutes les classes de poids.
2. Nouvelles Règles de Mise à l'Échelle : Les auteurs découvrent que les composants de porte nécessitent des mises à l'échelle de taux d'apprentissage non standard sous SGD :
   • Les matrices de poids de porte ($W_\alpha, W_\beta$) nécessitent une mise à l'échelle de $\Theta(1/\sqrt{d})$.
   • Les paramètres de porte scalaires ($a_{log}, b$) nécessitent une mise à l'échelle de $\Theta(\sqrt{d})$.
3. Validation Empirique : Les auteurs pré-entraînent des modèles de langage GDN à travers plusieurs largeurs (allant de 256 à 1563 dimensions cachées) en utilisant les optimiseurs AdamW et SGD.
   • Résultats : Sous la configuration µP proposée, le taux d'apprentissage optimal est transféré en zéro-shot à travers toutes les largeurs de modèle pour les deux optimiseurs. En revanche, la paramétrisation standard (SP) et la configuration µP originale échouent au transfert, le taux d'apprentissage optimal changeant significativement à mesure que la largeur du modèle augmente.
   • Stabilité : Les configurations proposées permettent un transfert de taux d'apprentissage stable, confirmant les dérivations théoriques.

Signification et Revendications
Le papier affirme que son analyse résout le problème ouvert de la paramétrisation des réseaux Delta à porte pour l'apprentissage de caractéristiques à grande échelle. En démontrant que la formulation µP dérivée permet un transfert de hyperparamètres en zéro-shot, ce travail valide la correction de la dérivation théorique et son utilité pratique. Les auteurs soulignent qu'il s'agit du premier travail dérivant et validant une paramétrisation cohérente avec µP spécifiquement pour les réseaux Delta à porte, se distinguant ainsi des travaux antérieurs sur les SSM diagonaux. Les conclusions suggèrent qu'une paramétrisation appropriée est essentielle pour la mise à l'échelle des modèles récurrents linéaires et que les mécanismes de porte spécifiques dans les GDN nécessitent des règles de mise à l'échelle uniques qui diffèrent à la fois des Transformers standards et d'autres architectures linéaires. Les auteurs espèrent que ces dérivations guideront les recherches futures sur les lois d'échelle d'autres architectures linéaires ou hybrides.