거의 0인 토큰 임베딩이 게이트형 모델 학습을 멈추게 하는 이유

들어가며

거의 0에 가까운 토큰 임베딩은 단순한 작은 값처럼 보인다. 하지만 게이트형 시퀀스 모델에서는 학습을 멈추게 하는 숨은 결함의 출발점이 될 수 있다. 논문 ‘Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models’에 따르면, 매우 드문 토큰이나 템플릿 토큰은 장기간 입력 방향의 그래디언트를 받지 못할 수 있다. 입력 임베딩과 출력 임베딩이 묶여 있지 않다면 AdamW의 가중치 감쇠는 계속 적용되고, 해당 행의 값은 사실상 0에 가까워진다.

핵심 내용

• 연구진은 291개 오픈 웨이트 언어 모델 저장소와 280개의 서로 다른 임베딩 테이블을 조사했다. 임베딩을 공유하지 않는 게이트형·하이브리드 체크포인트 38개 중 20개에서 노름이 10^-6보다 작은 접근 가능한 입력 행이 확인됐다. 가중치 공유 모델에서는 해당 사례가 없었다.
• 아키텍처에 따라 영향이 달랐다. 실험한 Transformer 설정에서는 시퀀스의 첫 위치가 아닌 위치나 BOS 토큰이 문제를 막았다. 반면 Mamba, Gated DeltaNet, xLSTM에서는 거의 0인 토큰이 입력 임베딩 그래디언트를 30자릿수 이상, 최대 35자릿수까지 키울 수 있었다.
• 거의 0인 입력에서는 SiLU 게이트나 선형 쿼리 같은 현재 토큰 읽기 요소가 사라지기 쉽다. 그러면 과거 상태를 읽는 경로가 약해진다. 여러 층에서 잔차 RMS가 RMSNorm 바닥값의 제곱근보다 낮아지면 Pre-RMSNorm이 역전파 그래디언트를 크게 증폭시킬 수 있다.
• 그래디언트의 한 요소가 약 1.84×10^19을 넘으면 PyTorch, Hugging Face Trainer, FSDP2의 FP32 전역 노름 클리핑이 무한대로 오버플로할 수 있다. 그 결과 현재 배치의 그래디언트가 0이 된다. Nemotron-3-Nano-4B의 텔루구어 지속 사전학습 실험에서는 약 20%의 단계에서 업데이트가 손실됐고, 문제가 발생한 임베딩 행은 계속 고정됐다.

의미와 영향

이 연구는 손실 곡선이나 전체 그래디언트 노름만 보는 것만으로는 충분하지 않다고 지적한다. 특히 임베딩을 공유하지 않는 모델과 게이트형 아키텍처에서는 접근 가능한 입력 임베딩 행을 개별적으로 검사해야 한다. 드물게 등장하는 토큰 하나가 게이트, 정규화, 클리핑을 거치며 전체 학습 안정성에 영향을 줄 수 있기 때문이다.

공개된 체크포인트를 위한 GroupClip은 입력 임베딩의 클리핑을 다른 파라미터와 분리하고, 64비트 또는 재스케일된 노름을 사용한다. 보고된 실험에서 재스케일 방식은 업데이트 손실을 없애고 문제가 발생한 행을 한 번의 업데이트로 회복했으며, 측정된 전체 오버헤드는 매우 작았다. 사전학습용 GateFloor는 SiLU(z)+δ 형태를 사용하며 δ는 0.01~0.1 범위다. 이를 통해 입력이 0이어도 게이트가 양수가 되고, 최대 3억 7천만 파라미터 규모의 실험에서 거의 0인 그래디언트를 28자릿수 줄이면서 기준 모델과 비슷한 퍼플렉서티를 보였다.

희귀 토큰의 임베딩은 사용 빈도가 낮다고 해서 무해한 것은 아니다. 게이트형 모델을 개발하거나 배포할 때 임베딩의 최소 노름과 비정상적인 클리핑 패턴을 점검하는 것이 실용적인 안정성 검사로 자리 잡을 수 있다.

출처: Hugging Face Daily Papers