역전파에서 길을 잃다: LM 헤더의 그래디언트 병목 현상

초록

신경망 언어 모델(LM)의 마지막 층은 차원 D의 출력 특징을 어휘 집합의 크기인 차원 V의 로짓으로 투사하며, 일반적으로 D는 V에 비해 매우 작습니다. 이러한 불일치는 신경망 LM의 표현력 제한 위험을 초래하는 소위 소프트맥스 병목 현상으로 알려져 있습니다. 우리는 소프트맥스 병목 현상이 단순히 표현력의 병목이 아닌 최적화의 병목이기도 함을 보입니다. V차원 그래디언트를 rank-D 선형 층을 통해 역전파하면 불가피한 압축이 발생하며, 이는 대다수 매개변수에 제공되는 훈련 피드백을 왜곡합니다. 우리는 이 현상에 대한 이론적 분석을 제시하고, 출력 층에 의해 그래디언트 노름의 95-99%가 억제되어 심각하게 비최적의 업데이트 방향이 초래됨을 실증적으로 측정합니다. 통제된 사전 훈련 실험을 통해 그래디언트 병목 현상이 사소한 패턴의 학습을 불가능하게 하며, 대규모 언어 모델(LLM)의 훈련 역학에 극적인 영향을 미친다는 것을 확인합니다. 우리는 이 본질적 결함이 모델 아키텍처와 무관하게 대규모 훈련 비효율성에 기여하며, 새로운 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.

역전파에서 길을 잃다: LM 헤더의 그래디언트 병목 현상

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

초록

Support