Natalie Abreu Nikhil Vyas Sham Kakade Depen Morwani

Résumé

Les récents efforts visant à accélérer l'entraînement préalable des modèles de langage à grande échelle (LLM) se sont concentrés sur des approximations à faible coût computationnel, exploitant la structure du second ordre. Cela soulève une question fondamentale pour l'entraînement à grande échelle : quelle perte de performance est-elle induite par ces approximations ? Pour examiner cette question, nous établissons une borne supérieure pratique sur la complexité itérative en appliquant une préconditionnement complet du type Gauss-Newton (GN) à des modèles Transformer de taille allant jusqu'à 150 millions de paramètres. Nos expériences montrent que les mises à jour complètes GN permettent des gains substantiels par rapport aux optimiseurs existants, réduisant le nombre d'itérations d'entraînement de 5,4 fois par rapport à des références solides telles que SOAP ou Muon. En outre, nous constatons qu’un préconditionneur GN précis par couche, qui ignore les informations croisées entre couches, atteint presque la performance du méthode GN complète. Collectivement, nos résultats suggèrent que : (1) l’approximation GN est très efficace pour la précondition, ce qui implique que les termes du second ordre de la fonction de perte pourraient ne pas être critiques pour la vitesse de convergence ; (2) la structure de la hessienne par couche contient suffisamment d’information pour atteindre la majeure partie des gains potentiels ; et (3) un écart significatif existe entre les méthodes approximatives actuelles et un oracle idéalisé par couche.