[논문리뷰] Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models

링크: 논문 PDF로 바로 열기

메타데이터

저자: Xin Liu, Boyuan Sun, Yuxin Wang, Haoran Song, Mingjian Zhao, Mengyu Liu

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• Near-Zero Token: 사전 학습 중 사용되지 않은 토큰으로, Untied embedding 구조에서 AdamW weight decay로 인해 norm이 $10^{-6}$ 이하로 수렴한 토큰을 의미합니다.
• Gated Sequence Models: Mamba, Gated DeltaNet, xLSTM과 같이 어텐션 대신 gated recurrence를 사용하는 모델군을 지칭합니다.
• Gradient Amplification: Near-zero input이 네트워크 내부의 SiLU 게이트나 linear readout을 통과할 때, Pre-RMSNorm에 의해 역전파 그래디언트가 기하급수적으로 증폭되는 현상입니다.
• FP32 Overflow: 그래디언트 norm이 FP32의 최대 표현 범위를 초과하여 inf로 계산되고, 이로 인해 배치의 전체 그래디언트가 0으로 소실되는 수치적 불안정성입니다.
• GroupClip: 본 논문에서 제안하는 기법으로, 임베딩 레이어와 백본의 그래디언트 클리핑을 분리하여 FP32 overflow를 방지하고 학습을 안정화하는 방식입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 많은 오픈 웨이트 모델에서 발견되는 Near-Zero Token Embeddings가 gated sequence model의 학습을 예기치 않게 저해하는 메커니즘을 규명합니다. 기존 연구들은 Transformer 구조에서의 임베딩 문제를 다루었으나, gated 구조에서 발생하는 극단적인 그래디언트 증폭 및 클리핑 장애는 명확히 분석되지 않았습니다. 분석 결과, 291개 오픈 모델 중 52.6%의 untied gated 모델이 이 문제에 취약하며, 특정 토큰 등장 시 배치의 그래디언트가 완전히 0으로 처리되는 'silent update zeroing' 현상이 발생합니다 [Figure 1]. 이러한 현상은 모델이 새로운 언어나 템플릿에 적응할 때 학습 업데이트를 강제로 중단시켜 성능 저하를 야기합니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

저자들은 gated block 내의 SiLU 게이트가 0입력 근처에서 historical state readout을 0으로 수렴시켜 발생하는 Pre-RMSNorm 역전파 증폭 효과를 이론적으로 증명했습니다 [Figure 2]. 이를 해결하기 위해 사후 학습용 GroupClip과 사전 학습용 GateFloor를 제안합니다. GroupClip은 임베딩 레이어의 그래디언트를 별도의 64-bit norm으로 클리핑하여 전체 백본 업데이트를 보호하며, 이를 통해 Nemotron-3-Nano-4B 모델의 continual pretraining에서 발생하던 20% 수준의 업데이트 손실을 완전히 제거했습니다 [Table 2]. 또한 사전 학습 단계에서 적용 가능한 GateFloor ($\text{SiLU}(z) + \delta$)는 게이트 활성화를 강제로 유지하여 Near-zero 입력에 의한 그래디언트 증폭을 28 orders of magnitude 이상 감소시켰습니다 [Figure 6].

4. Conclusion & Impact (결론 및 시사점)

본 연구는 gated sequence model에서 임베딩 초기화와 AdamW decay 전략이 초래하는 수치적 취약점을 규명하고 이를 해결할 실질적인 기법을 제시했습니다. 제안된 GroupClip과 GateFloor는 추가적인 연산 오버헤드 없이 학습 안정성을 비약적으로 향상시킵니다. 이 연구는 대규모 언어 모델의 continual pretraining이나 도메인 적응 과정에서 발생하는 학습 정체 현상을 해결할 핵심적인 가이드를 제공하며, 향후 다양한 gated architecture의 견고한 설계에 기여할 것으로 기대됩니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.