게이트형 언어모델의 0 근처 임베딩이 학습 업데이트를 조용히 지운다
Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models
무엇인가
이 논문은 공개 가중치 언어모델의 입력 임베딩에서 발견되는 근접 영벡터 토큰이 게이트형 시퀀스 모델의 전체 파라미터 학습을 조용히 방해하는 문제를 다룬다. 저자들은 291개 HuggingFace 저장소와 280개 임베딩 테이블을 조사해 101개 weight-tied 모델에서는 도달 가능한 근접 영벡터 토큰이 0개였지만, untied 모델 190개 중 67개에서 발견했다고 보고한다. 특히 untied 게이트·하이브리드 체크포인트 38개 중 20개인 52.6%가 10^-6 미만 노름의 도달 가능한 입력 임베딩을 가졌다. 원인은 사전학습에서 본 적 없는 희귀 토큰이나 템플릿 토큰의 입력 행이 0 그래디언트를 받는 동안 AdamW weight decay가 untied 입력 임베딩을 계속 줄이기 때문이다. 테스트한 Transformer 구성은 비초기 위치나 BOS 토큰이 있으면 이 실패 모드를 피했지만, Mamba·Gated DeltaNet·xLSTM 같은 게이트형 모델은 취약했다.
어떻게 동작하나
메커니즘은 게이트형 블록의 현재 토큰 판독 계수에서 시작한다. SiLU 게이트나 선형 쿼리는 입력이 0일 때 0이 되므로, 근접 영벡터 입력은 과거 상태 판독을 거의 0으로 만든다. 그러면 연속 레이어의 잔차 RMS가 RMSNorm 바닥값의 제곱근보다 작아지고, Pre-RMSNorm이 역전파를 레이어당 최대 10^3배씩 증폭해 깊이 방향으로 10^30을 넘길 수 있다. 이때 FP32 전역 노름 클리핑에서 어떤 그래디언트 원소가 1.84×10^19를 넘으면 제곱이 무한대로 넘치고, 유한한 그래디언트를 무한대로 나눈 결과 정확히 0.0이 되어 NaN 없이 현재 배치의 모든 그래디언트가 지워진다. PyTorch clip_grad_norm_, HuggingFace Trainer, FSDP2에서 이런 일이 발생하며, Nemotron-3-Nano-4B의 텔루구 계속 사전학습에서는 약 20% 스텝의 업데이트가 버려지고 트리거된 임베딩 행은 얼어붙은 채 남았다.
무엇과 다른가
출시된 체크포인트를 위한 해법이 GroupClip이다. 백본 가중치는 기존처럼 max_norm 1.0으로 클리핑하되, 입력 임베딩 그래디언트는 64비트 노름으로 따로 계산해 min(1.0, 1.0/(gn+1e-12))를 곱하는 세 줄짜리 분리 클리핑이다. 깨끗한 배치에서는 입력 임베딩 그래디언트가 백본보다 훨씬 작을 때 표준 전역 클리핑과 사실상 같고, 트리거 배치에서는 백본 업데이트를 보존한다. FP64 노름은 1.84×10^19 초과 오버플로를 막아 한 번의 단위 노름 AdamW 스텝으로 트리거 행을 오버플로 영역 밖으로 끌어올린다. FP64 노름은 3.1GB 임시 메모리와 클리핑 호출 18.8ms 대 21.2ms를 썼지만, 복사 없는 재스케일 변형은 이를 제거한다고 저자들은 밝힌다.
어떻게 쓰나
사전학습부터 하는 경우에는 GateFloor를 제안한다. 먼저 입력 임베딩의 weight_decay를 0으로 두는 것을 기본 권장한다. 그래도 SiLU 게이트는 작은 노름 초기화나 어휘 확장에서 1/r 증폭을 유지하고, Sigmoid 게이트는 0에서 0.5로 양수지만 양수 출력을 1로 제한해 퍼플렉시티를 해친다. GateFloor는 gδ(z)=SiLU(z)+δ, δ∈[0.01,0.1]로 게이트에 작은 양수 오프셋을 더한다. z=0에서 gδ(0)=δ>0이므로 gated RMSNorm 앞에 δ·y를 통과시키고, t>0에서는 과거 상태가 0이 아님으로써, t=0에서는 Mamba-2의 conv1d bias로 인해 Layer 0이 sub-ε 영역에서 벗어난다. 추가 파라미터는 없다.
전제와 한계
GroupClip 실험은 FineWeb-2 계속 사전학습을 3개 시드, 1,000스텝, 32×512=16,384 토큰/스텝으로 Nemotron-3-Nano-4B 텔루구와 Falcon-H1-1.5B 몽골어에 적용하고 FLORES-200 devtest PPL로 평가했다. Nemotron 텔루구에서 FP32 전역 클리핑은 19.8% 스텝을 멈추고 13.44±0.17 PPL을 냈지만, GroupClip FP64는 멈춘 스텝 0.0, 13.04±0.19 PPL로 Oracle Mask와 같았다. 30.7±0.6개의 근접 영벡터 행이 첫 접촉에서 3.5×10^-4 스텝으로 치유됐고, 10^10 초과 스파이크는 21.0±1.0스텝에서만 발생했다. Falcon-H1 몽골어는 트리거 배치가 96.6%였고 GroupClip이 100스텝 PPL을 40.40에서 37.47로 낮췄으며, 1,000스텝에서 15.58±0.10으로 FP32 전역 클리핑 15.68±0.11과 Oracle Mask 15.63±0.10을 앞섰다. 300스텝 비교에서 Skip-on-Inf 16.63±0.36, ZClip 16.57±0.39, GradientStabilizer 16.23±0.30, Freeze 16.16±0.24, FP64 전역 노름 16.09±0.35, Mean Re-init 15.96±0.24, GroupClip 15.95±0.25, Oracle Mask 15.94±0.26였다. Mean Re-init은 Falcon3-7B에서 희귀 토큰 꼬리 NLL을 3.8에서 9.8로 악화시켰다.
GateFloor 실험은 FineWeb-Edu에서 어휘 50,304의 untied 모델 14개를 사전학습했다. 24층 Mamba-2 165M 모델은 1.05B 토큰과 2개 시드, δ 스윕을 포함했고, 16층 3:1 하이브리드 203M과 20층 3:1 하이브리드 370M은 각각 1.05B와 8×H100에서 1.57B 토큰을 학습했다. δ=0.1, 최대 370M에서 r=10^-30 중간 시퀀스 그래디언트를 1.6×10^28에서 0.13으로 28자리 이상 줄였고, 처리량은 542.8k 대 527.2k tok/s로 비슷했으며 검증 PPL은 기준선 대비 ±0.17 이내였다. 165M에서 r=10^-30 그래디언트는 δ에 반비례해 4.14→1.51→0.60→0.29로 줄었고, 기준선은 5.6×10^28이었다. Sigmoid 게이트는 규모별로 PPL을 0.18~0.56 악화시켰다.
개발자 관점에서는 untied 입력 임베딩을 쓰는 게이트형 체크포인트를 전체 파라미터로 계속 사전학습하거나 다국어·템플릿 데이터로 미세조정할 때 증상이 없다고 안심하면 안 된다. 예외나 NaN 없이 grad_norm이 inf로 기록되고 현재 배치 업데이트가 0이 될 수 있으므로, 입력 임베딩 weight decay 여부와 근접 영벡터 행을 점검하고, 해당 작업에는 GroupClip을 적용하는 것이 직접적이다. LoRA처럼 입력 임베딩을 고정하는 학습과 추론은 이 실패 모드를 피했고, on-policy RL 롤아웃 샘플링에서는 근접 영벡터 토큰이 0개 나왔지만, 전체 파라미터 RL 업데이트는 근접 영벡터 채팅 태그가 프롬프트에 들어오면 여전히 트리거될 수 있다. 새로 사전학습한다면 입력 임베딩 weight decay를 끄고 GateFloor를 검토해야 한다.
저자들이 밝힌 한계는 범위다. 감사와 그래디언트 탐침은 최대 48B 모델까지지만 계속 사전학습 실험은 1.5B~4B에 머물고, 처음부터 학습하는 GateFloor 실험은 370M, 시드 1~2개에 그친다. 더 큰 규모에서 GateFloor의 퍼플렉시티와 안정성, GroupClip의 재스케일 변형 비용을 추가 검증해야 한다. 또한 GroupClip의 FP64 노름은 일시적으로 3.1GB 메모리를 쓰므로 메모리 여유가 작은 환경에서는 복사 없는 재스케일 구현을 확인할 필요가 있다.