Perdidos en la Retropropagación: La Capa de Salida del LM es un Cuello de Botella del Gradiente

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

March 10, 2026

Autores: Nathan Godey, Yoav Artzi

cs.AI

Resumen

La última capa de los modelos de lenguaje neuronal (LM) proyecta las características de salida de dimensión D a logits en dimensión V, el tamaño del vocabulario, donde generalmente D ≪ V. Se sabe que esta discrepancia aumenta los riesgos de expresividad limitada en los LM neuronales, creando un llamado cuello de botella del softmax. Demostramos que el cuello de botella del softmax no es solo un cuello de botella de expresividad, sino también un cuello de botella de optimización. La retropropagación de gradientes V-dimensionales a través de una capa lineal de rango D induce una compresión inevitable, lo que altera la retroalimentación del entrenamiento proporcionada a la gran mayoría de los parámetros. Presentamos un análisis teórico de este fenómeno y medimos empíricamente que entre el 95% y el 99% de la norma del gradiente es suprimida por la capa de salida, resultando en direcciones de actualización enormemente subóptimas. Realizamos experimentos controlados de preentrenamiento que muestran que el cuello de botella del gradiente hace que patrones triviales sean inaprendibles y afecta drásticamente la dinámica del entrenamiento de los LLM. Argumentamos que este defecto inherente contribuye a ineficiencias en el entrenamiento a escala, independientemente de la arquitectura del modelo, y plantea la necesidad de nuevos diseños para la cabeza del LM.

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.