Perdu dans la rétropropagation : La tête de modèle de langage constitue un goulot d’étranglement du gradient

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

March 10, 2026

Auteurs: Nathan Godey, Yoav Artzi

cs.AI

Résumé

La dernière couche des modèles de langage neuronaux projette les caractéristiques de sortie de dimension D vers des logits de dimension V, la taille du vocabulaire, où généralement D ≪ V. On sait que cette inadéquance accroît les risques de limitation de l'expressivité dans les modèles de langage neuronaux, créant ce qu'on appelle un goulot d'étranglement du softmax. Nous montrons que ce goulot d'étranglement n'est pas seulement une limitation de l'expressivité mais aussi une limitation de l'optimisation. La rétropropagation de gradients de dimension V à travers une couche linéaire de rang D induit une compression inévitable, ce qui altère la rétroaction d'apprentissage fournie à la grande majorité des paramètres. Nous présentons une analyse théorique de ce phénomène et mesurons empiriquement que 95 à 99 % de la norme du gradient est supprimée par la couche de sortie, résultant en des directions de mise à jour largement sous-optimales. Nous menons des expériences d'apprentissage préalable contrôlé montrant que ce goulot d'étranglement du gradient rend des motifs triviaux inapprenables et affecte considérablement la dynamique d'apprentissage des grands modèles de langage. Nous soutenons que ce défaut inhérent contribue aux inefficacités d'entraînement à grande échelle, indépendamment de l'architecture du modèle, et soulève la nécessité de nouvelles conceptions de tête de modèle de langage.

English

The last layer of neural language models (LMs) projects output features of dimension D to logits in dimension V, the size of the vocabulary, where usually D ll V. This mismatch is known to raise risks of limited expressivity in neural LMs, creating a so-called softmax bottleneck. We show the softmax bottleneck is not only an expressivity bottleneck but also an optimization bottleneck. Backpropagating V-dimensional gradients through a rank-D linear layer induces unavoidable compression, which alters the training feedback provided to the vast majority of the parameters. We present a theoretical analysis of this phenomenon and measure empirically that 95-99% of the gradient norm is suppressed by the output layer, resulting in vastly suboptimal update directions. We conduct controlled pretraining experiments showing that the gradient bottleneck makes trivial patterns unlearnable, and drastically affects the training dynamics of LLMs. We argue that this inherent flaw contributes to training inefficiencies at scale independently of the model architecture, and raises the need for new LM head designs.