Résumé Technique : GEFEN : UN OPTIMISEUR STOCHASTIQUE OPTIMISÉ

Énoncé du Problème

L'apprentissage profond moderne repose largement sur l'optimiseur AdamW en raison de ses performances robustes sur diverses architectures. Cependant, AdamW maintient des moyennes mobiles du premier et du second moment pour chaque paramètre, typiquement en FP32. Cela se traduit par une empreinte mémoire de l'état de l'optimiseur d'environ deux fois la taille des paramètres du modèle (soit environ 8 octets par paramètre). Pour les modèles à grande échelle (ex. : 13 milliards de paramètres), cet état peut à lui seul consommer environ 97 own GiB de mémoire, ce qui limite sévèrement la taille des modèles entraînables, la taille maximale des lots (batch sizes) et le débit d'entraînement sur le matériel à mémoire contrainte.

Les alternatives existantes en matière d'efficacité mémoire présentent des compromis importants :

• Optimiseurs Quantifiés (ex. : Adam8bit, Adam4bit) : Réduisent la mémoire mais reposent souvent sur des codebooks manuellement conçus et sous-optimaux, ou des tailles de blocs fixes (introduisant des hyperparamètres), et peuvent ne pas être universellement applicables (ex. : éviter la quantification des embeddings).
• Optimiseurs Structurés (ex. : Adam-mini) : Partagent les statistiques du second moment basées sur des règles manuelles imitant les structures diagonales de blocs de la Hessienne. Cependant, ceux-ci nécessitent des métadonnées spécifiques à l'architecture (ex. : nombre de têtes d'attention) et manquent d'une justification théorique solide quant à l'intérêt d'un regroupement aligné sur la Hessienne.

L'objectif est de concevoir un optimiseur qui réduit significativement l'empreinte mémoire (au-delà d'Adam-mini) tout en maintenant des performances de niveau AdamW, sans configuration manuelle, sans hyperparamètres supplémentaires et sans réajustement du taux d'apprentissage.

Méthodologie

La méthode proposée, Gefen, atteint la réduction de mémoire via deux mécanismes principaux : le partage automatique du second moment et la quantification apprise du premier moment.

1. Motivation Théorique

L'approche est fondée sur un résultat théorique (Théorème 4.1) démontrant que pour les paramètres présentant de grands éléments mixtes de la Hessienne, le rapport de leurs carrés de gradients est contraint vers l'unité. Cela implique que les paramètres possédant une forte affinité avec la Hessienne possèdent naturellement des magnitudes de carrés de gradients similaires, ce qui en fait des candidats idéaux pour le partage des estimations du second moment. Bien que le calcul direct de la Hessienne soit intraitable à grande échelle, les auteurs observent que la structure des carrés de gradients lors de la première étape d'entraînement est corrélée à cette structure de la Hessienne.

2. Partitionnement Automatique par Blocs

Au lieu de s'appuyer sur des règles manuelles ou le calcul de la Hessienne, Gefen infère la structure de bloc optimale à partir des carrés de gradients de la première étape d'optimisation ($g_0$).

• Algorithme : Pour un tenseur de paramètres aplati, l'algorithme teste tous les diviseurs propres $p$ de la taille du tenseur $n$. Il reforme les carrés de gradients en blocs de taille $p$ et calcule l'hétérogénéité intra-bloc $E(p)$ (la variance des carrés de gradients au sein des blocs).
• Stratégie de Sélection : L'algorithme sélectionne la période $p$ qui produit la première chute prononcée de $E(p)$ à mesure que la taille du bloc augmente, privilégiant les blocs plus larges pour l'efficacité mémoire tout en garantissant la similitude des gradients. Si aucune chute significative n'est trouvée ou si le $p$ optimal $< 8$, il revient par défaut à $p=1$ (aucun partage).
• Utilisation : Ce schéma de partitionnement unique est utilisé à la fois pour le partage du second moment et pour la quantification du premier moment.

3. Quantification Exacte par Programmation Dynamique

Pour réduire la mémoire du premier moment (momentum), Gefen emploie un codebook apprenable et exact plutôt que des schémas fixes ou heuristiques.

• Apprentissage du Codebook : En utilisant les gradients initiaux, la méthode construit un histogramme des valeurs de gradient normalisées. Elle résout ensuite le problème de quantification scalaire optimale en utilisant la Programmation Dynamique (DP) (inspirée de Wang & Song, 2011) pour trouver le codebook exact pour la distribution et la largeur de bits données.
• Stabilité : Contra Contrairement à Lloyd-Max, qui peut souffrir de "contraction" (réduction de la plage de quantification au fil du temps) et de sensibilité à l'initialisation, Gefen force les extrémités du codebook à être fixées à $-1$ et $1$. Cela garantit que la plage de quantification est préservée tout au long de l'entraînement.
• Efficacité : Le codebook est appris une seule fois au début de l'entraînement (en utilisant CUDA pour l'histogramme et CPU/Numba pour la DP) et est réutilisé pour toute la durée de l'entraînement.

4. Flux de l'Algorithme

Gefen fonctionne comme un remplacement direct (drop-in replacement) pour AdamW :

1. Initialisation : Partitionner les paramètres basés sur $g_0$ et apprendre le codebook de quantification.
2. Étape de Mise à Jour :
   • Calculer les gradients.
   • Mettre à jour le premier moment (momentum) en déquantifiant les valeurs du codebook stockées, en appliquant le momentum, en re-quantifiant, et en mettant à jour le maximum absol 대로 par bloc.
   • Mettre à jour le second moment en calculant la moyenne des carrés de gradients au sein de chaque bloc et en appliquant une moyenne mobile exponentielle.
   • Appliquer la correction de biais et mettre à jour les paramètres en utilisant le second moment partagé et le premier moment quantifié.

Principales Contributions

1. Efficacité Mémoire : Gefen réduit l'empreinte mémoire de l'optimiseur d'environ 8× par rapport à AdamW (passant d'environ 2,0× la taille des paramètres à ~0,25×), ce qui correspond à une réduction d'environ 6,5 own GiB par milliard de paramètres.
2. Adaptation Automatique : La méthode ne nécessite aucune métadonnée spécifique à l'architecture, aucune règle de regroupement manuelle, ni hyperparamètres supplémentaires au-delà des valeurs par défaut d'AdamW. Elle infère automatiquement les structures de blocs à partir des données.
3. Fondement Théorique : Fournit un lien théorique entre les entrées de la Hessienne et les rapports de carrés de gradients, justifiant le partage des statistiques du second moment pour les paramètres alignés sur la Hessienne.
4. Quantification Optimale : Introduit une approche exacte par programmation dynamique basée sur l'histogramme pour quantifier le momentum, évitant la sous-optimalité et l'instabilité des méthodes itératives comme Lloyd-Max.
5. Parité de Performance : Démontre que ces réductions ne se font pas au détriment de la qualité du modèle, égalant les performances d'AdamW sur diverses tâches.

Résultats Expérimentaux

Les auteurs ont évalué Gefen contre AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor et SM3 sur des modèles CNN, ResNet, GPT-2 et Llama 3.

• Empreinte Mémoire de Pointe : Gefen a systématiquement obtenu l'utilisation de mémoire de pointe la plus faible parmi les méthodes comparées. Pour un modèle de 1,5 milliard de paramètres, Gefen a réduit l'état de l'optimiseur de 11,2 own GiB (AdamW) à 1,5 own GiB.
• Performance (Qualité) : Sur l'ensemble des configurations (OpenWebText, C4, ImageNet, MNIST), Gefen a obtenu des pertes de validation et des précocités comparables à AdamW. Notamment, il a surpassé Adam4bit et Adam-mini dans plusieurs métriques tout en conservant le même calendrier de taux d'apprentissage.
• Débit (Throughput) : La réduction de l'empreinte mémoire a permis des micro-lots (microbatches) plus grands dans l'entraînement distribué :
  • FSDP (Llama 3-1.5B) : Gefen a permis un microbatch de taille 2 (contre 1 pour les autres), entraînant une augmentation du débit de 56 % (tokens/sec).
  • DDP (GPT-2-1B) : Gefen a permis l'entraînement avec un microbatch de taille 2, alors qu'AdamW ne pouvait même pas supporter un microbatch de taille 1. Cela a entraîné une amélioration du débit de 21 % par rapport à Adam-mini.
• Stabilité : Les études d'ablation ont confirmé que le codebook appris reste stable tout au long de l'entraînement, validant la décision de le calculer uniquement à l'initialisation.

Signification et Revendications

Le papier positionne Gefen comme un remplacement pratique et direct ("drop-in") pour AdamW, qui traite le goulot d'étranglement de la mémoire dans l'entraînement à grande échelle sans sacrifier la performance ni nécessiter de réglage manuel.

• Impact Pratique : En réduant l'état de l'optimiseur, Gefen permet d'entraîner des modèles plus grands ou d'utiliser des lots plus importants sur un matériel fixe, améliorant directement le débit d'entraînement et l'efficacité des échantillons.
• Généralité : Contrairement aux méthodes précédentes qui reposent sur des heuristiques manuelles ou des noms de tenseurs spécifiques, Gefen est agnostique à l'architecture et s'adapte automatiquement aux données.
• Complémentarité : Les auteurs notent que la compression de l'état de l'optimiseur de Gefen est orthogonale à la compression de communication des gradients (ex. : 1-bit Adam) et aux stratégies d'entraînement distribué (FSDP/DDP), ce qui signifie qu'elle peut être combinée avec ces techniques pour des gains supplémentaires.

Les auteurs fournissent une implémentation Python complète avec des noyaux CUDA fusionnés, mettant l'accent sur la reproductibilité et la facilité d'intégration dans les pipelines d'entraînement existants.