ByteDance Seed propose AIR, un calendrier de pas d'apprentissage : une erreur finale inférieure de 0,01 à celle de cosine

Des chercheurs de ByteDance Seed (Qifan Yu, Xinyu Ma, Xingyan Bin, Di He) proposent AIR, un calendrier de vitesse d'apprentissage. Lors du pré-entraînement d'un modèle Mixture-of-Experts (MoE) sur 128 milliards de tokens, il réduit l'erreur finale sur les données d'entraînement (training loss) d'environ 0,01 par rapport à la décroissance cosinus du pas (cosine). La comparaison a été menée à somme de pas identique sur l'ensemble de l'entraînement. Ici, la vitesse d'apprentissage (learning rate, LR) désigne la taille du pas avec lequel l'optimiseur modifie les poids du modèle. On la diminue généralement au fil de l'entraînement, le plus souvent selon une courbe cosinus.

Pourquoi la normalisation des poids fait parler d'elle

Pour contrôler la taille des poids, on utilise d'ordinaire le weight decay (WD) : à chaque pas, il tire légèrement les poids vers zéro. Récemment, la weight normalization (WN), c'est-à-dire l'entraînement avec des poids normalisés, a été envisagée comme alternative. Plusieurs travaux ont rapporté un gain notable de loss en pré-entraînement.

Les auteurs ont vérifié l'origine de ce gain et arrivent à une autre explication. Avec WN, le LR nominal se traduit par une taille de mise à jour réelle différente de celle obtenue avec WD. Il en résulte une trajectoire propre du pas d'apprentissage effectif (effective learning rate, ELR). Les chercheurs ont ajusté l'ELR des deux méthodes pour que les trajectoires coïncident. Les courbes de loss de WN et de WD se sont alors presque confondues. Le gain tenait donc à la forme de l'ELR, et non à la normalisation en soi.

D'où une conclusion pratique : il faut construire le calendrier pour l'ELR, et pas seulement pour le LR nominal.

Comment fonctionne AIR

Le nom signifie Affine Inverse-square-Root decay, soit une décroissance affine en racine carrée inverse. Le calendrier comprend deux parties.

• La base. Le pas décroît proportionnellement à 1/√t, où t est le numéro du pas d'entraînement. Les auteurs ont établi cette forme par la théorie. • La correction. Une transformation linéaire (affine). L'entraînement a une durée finie, et une décroissance pure en 1/√t ne donne pas les valeurs voulues au début et à la fin. La correction permet de fixer à l'avance le LR de pic et le LR final.

Comme avec cosine, l'utilisateur indique deux points extrêmes : le pic et la valeur finale.

Les auteurs font aussi une observation sur la phase tardive. Si plusieurs entraînements partagent la même fin de calendrier d'ELR, celle-ci atténue l'influence des choix de LR antérieurs. Les auteurs voient dans la planification tardive de l'ELR l'une des principales sources du gain.

Ce que montre l'expérience

La vérification a porté sur le pré-entraînement d'un modèle MoE sur 128 milliards de tokens. Dans un MoE, seule une partie des blocs d'experts travaille sur chaque token. AIR a été comparé à cosine et à d'autres calendriers candidats. Les résultats :

• La training loss est plus basse avec AIR. La loss finale est inférieure d'environ 0,01 à celle de cosine, à somme de LR identique sur l'ensemble de l'entraînement. • AIR est aussi meilleur sur la qualité en tâches aval (downstream). • Le calendrier est plus robuste au choix des hyperparamètres.

Dans une expérience distincte, la même trajectoire d'ELR a été appliquée au weight decay classique. La dynamique d'entraînement s'est reproduite. Les auteurs en concluent que l'intérêt ne se limite pas à l'entraînement avec une norme de poids fixe, et désignent l'ELR comme l'objet qu'il convient de planifier.

Ce que cela implique en pratique

Un écart de 0,01 de loss est modeste, mais il s'obtient en changeant un seul calendrier de pas. Le travail donne trois repères pratiques :

• vérifier comment votre schéma de poids et votre optimiseur transforment le LR nominal en LR effectif ; • essayer le calendrier AIR, où seuls le LR de pic et le LR final sont à définir ; • accorder plus d'attention à la fin de l'entraînement, car une fin de calendrier commune atténue l'influence des réglages antérieurs.

Dans l'abstract, les expériences sont décrites sur un pré-entraînement MoE de 128 milliards de tokens. Il ne donne aucune information sur d'autres tailles de modèles ni d'autres volumes de données.

L'article est publié sur arXiv sous le numéro 2610.40362 et figure dans la sélection Hugging Face Daily Papers.