Perdidos na Retropropagação: A Cabeça do Modelo de Linguagem é um Gargalo de Gradiente

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

March 10, 2026

Autores: Nathan Godey, Yoav Artzi

cs.AI

Resumo

A última camada dos modelos de linguagem neural (LMs) projeta as características de saída de dimensão D para logits em dimensão V, o tamanho do vocabulário, onde geralmente D << V. Sabe-se que esta discrepância aumenta os riscos de expressividade limitada nos LMs neurais, criando um chamado gargalo da softmax. Demonstramos que o gargalo da softmax não é apenas um gargalo de expressividade, mas também um gargalo de otimização. A retropropagação de gradientes V-dimensionais através de uma camada linear de posto D induz uma compressão inevitável, que altera o *feedback* de treinamento fornecido à vasta maioria dos parâmetros. Apresentamos uma análise teórica deste fenômeno e medimos empiricamente que 95-99% da norma do gradiente é suprimida pela camada de saída, resultando em direções de atualização vastamente subótimas. Realizamos experiências controladas de pré-treinamento que mostram que o gargalo do gradiente torna padrões triviais inaprendíveis e afeta drasticamente a dinâmica de treinamento dos LLMs. Argumentamos que esta falha inerente contribui para ineficiências de treinamento em escala, independentemente da arquitetura do modelo, e levanta a necessidade de novos projetos para a camada final (head) dos LMs.

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.