Softmax 병목 현상: LLM을 크게 만든다고 항상 더 똑똑해지는 것은 아닌 이유

연구자들이 언어 모델의 크기를 확장할 때—더 많은 파라미터, 더 많은 레이어, 더 넓은 은닉 차원—암묵적인 가정이 있습니다. 즉, 더 큰 모델이 더 많은 것을 표현할 수 있다는 것입니다. 더 높은 표현력, 더 많은 지식, 더 나은 예측입니다. 대부분의 경우 이것은 사실입니다. 하지만 스케일링만으로는 높일 수 없는 구조적 한계가 존재하며, 이는 네트워크의 최종 레이어에 위치합니다. 이를 Softmax 병목 현상(softmax bottleneck)이라고 부릅니다. 이 현상을 이해하면 왜 일부 모델이 순수한 컴퓨팅 파워로는 해결할 수 없는 성능 벽에 부딪히는지, 그리고 단순히 모델 크기를 늘리는 것 외에 특정 아키텍처 선택(Mixture of Experts, 출력 분해(output factorisation), Softmax 혼합(mixture of softmaxes))이 존재하는 이유를 알 수 있습니다.

Softmax 병목 현상이 실제로 무엇인지
언어 모델의 최종 단계에서, 여러분은 어휘집(vocabulary)에 있는 모든 토큰에 대한 확률 분포를 생성해야 합니다. 일반적으로 이 크기는 30,000개에서 200,000개의 토큰입니다. 모델은 은닉 상태 벡터 $h$ (차원 $d$)를 출력 임베딩 행렬 $W$ (크기 $d imes V$, 여기서 $V$는 어휘집 크기)와 곱하고 softmax를 적용하여 이를 수행합니다. 문제는 이 행렬 곱의 결과가 $ ext{rank-}d$ 행렬이라는 것입니다. 만약 여러분이 근사하려고 하는

모델은 낮은 랭크 투영(low-rank projection)을 통해 높은 랭크 함수를 표현하도록 요청받고 있는 것입니다. 왜 이것이 실제에서 나타나는가? 10만 개의 토큰으로 구성된 어휘집은 엄청난 수의 문맥적 구별 능력을 가지고 있습니다. 예를 들어, 단어 'bank'가 'river', 'financial', 'blood', 'memory'와 같은 다른 단어들로 앞설 때 다음 토큰에 대한 확률을 얼마나 다르게 분포해야 하는지 생각해 보세요. 가능한 모든 선행 문맥에 걸쳐 전체 분포 행렬은 잠재적으로 매우 높은 랭크를 가집니다. 즉, 각 문맥이 어휘집에 대한 고유한 확률 분포를 생성하며, 이 분포들은 서로 거의 선형 독립일 수 있습니다. 은닉 차원 $d = 4096$을 가진 모델은 모델 본체의 매개변수(parameter) 수와 관계없이 최대 4096개의 선형 독립적인 출력 분포만을 생성할 수 있습니다. 트랜스포머 블록(transformer blocks)은 임의로 깊고 강력해질 수 있지만, 결국 $d$-차원 벡터를 생성하며, 이 벡터는 다음 토큰 분포의 제한된 다양성만을 표현할 수 있습니다. 이는

이는 효과적인 출력 랭크(effective output rank)가 단순히 $d$가 아니라 $K imes d$로 확장될 수 있게 합니다. Yang 등이 제안한 자체 해결책은 작동하지만, 추론 비용을 $K$에 비례하여 증가시킵니다. 입력/출력 임베딩 연결 (Tied input/output embeddings): 흥미로운 부작용으로, 입력 임베딩 행렬을 출력 행렬과 연결하는 것(매개변수 수를 줄이는 널리 사용되는 트릭)은 실제로 일부 구성에서 병목 현상 완화에 도움이 됩니다. 왜냐하면 입력 임베딩이 더 풍부한 토큰-토큰 관계를 인코딩하고, 출력 투영(output projection)이 이를 상속받기 때문입니다. 전문가 혼합 (Mixture of Experts, MoE): 서로 다른 전문가가 다른 입력에 대해 활성화될 때, 출력 단계의 효과적인 표현력은 활성화되는 전문가 수에 따라 확장되며, 이는 랭크 제약을 부분적으로 완화합니다. 이것이 바로 MoE 모델이 활성화된 매개변수 가중치보다 더 나은 성능을 낼 수 있는 간과된 이유 중 하나입니다. 마지막 레이어의 더 큰 은닉 차원 (Larger hidden dimensions in the final layers): 일부 아키텍처는 병목 현상이 출력 단계에서 가장 두드러지다는 것을 인식하고, 의도적으로 마지막 몇 개의 트랜스포머 블록을 넓히거나(widen) 다른 (더 넓은) 투영 헤드(projection head)를 사용합니다. 실무자를 위한 의미: 기본 모델을 미세 조정(fine-tuning)하면서 검증 손실(validation loss)이 작업에 비해 비정상적으로 높은 값에서 평탄화되는 것을 발견한다면, 병목 현상은 데이터나 훈련 문제가 아니라 아키텍처적인 문제일 수 있습니다.

이것은 다음과 같은 경우에 더 문제가 될 수 있습니다: 작업이 큰 어휘 집합(코드 생성, 다국어 작업)에 걸쳐 세밀한 토큰 수준의 구별을 요구하는 경우 모델의 은닉 차원(hidden dimension)이 어휘 크기에 비해 작은 경우 어휘 토큰(도메인별 용어)을 추가했지만 출력 아키텍처를 조정하지 않은 경우 해결책은 거의 '더 오래 훈련시키는 것'이 아닙니다. 이는 $d$를 늘리거나, 출력 인수분해(output factorisation)를 적용하거나, 모델의 토큰 분포 표현력에 구조적 한계가 있다는 점을 받아들이는 것입니다. 더 큰 그림 소프트맥스 병목 현상은 매개변수 수나 FLOP 추정치에서는 나타나지 않지만, 근본적으로 모델이 표현할 수 있는 것을 제한하는 일련의 아키텍처 제약 조건의 명확한 예입니다. 이 분야는 스케일링 법칙(scaling laws) — 더 많은 데이터, 더 많은 컴퓨팅, 더 나은 성능 — 에 초점을 맞추는 경향이 있으며, 이러한 법칙들은 실제로 존재합니다. 그러나 그것들은 아키텍처적 범위 내에서 작동합니다. 그러한 범위의 천장에 가까워지면, 더 많은 컴퓨팅은 도움이 되지 않습니다. 어디에 천장이 있는지 이해하는 것이 아키텍처 직관과 벤치마크 추구 능력을 구분하는 것입니다.