[논문 리뷰] Small-scale proxies for large-scale Transformer training instabilities

Mitchell Wortsman, Peter J. Liu외 14명|arXiv (Cornell University)|2023. 09. 25.

Explainable Artificial Intelligence (XAI)인용 수 4

한 줄 요약

이 논문은 대규모 트랜스포머에서의 훈련 불안정성 문제를 연구하기 위해 소규모 프록시 실험을 도입한다. 여기서 핵심적인 불안정성—attention 레이어에서의 로짓 증가와 로그 확률에서의 발산—은 qk-layernorm 및 z-loss와 같은 기법을 사용해 소규모 모델에서도 재현되고 완화됨을 보여준다. 저자들은 학습률 민감도가 안정성의 신뢰할 수 있는 진단 도구임을 입증하며, 활성화 및 기울기 노름의 스케일링 행동이 불안정성이 발생하기 전에 이를 예측할 수 있음을 보여준다.

ABSTRACT

Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the $μ$Param (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.

연구 동기 및 목표

• 동일한 초모수를 사용함에도 불구하고 소규모에서는 나타나지 않는 대규모 트랜스포머에서의 훈련 불안정성을 조사한다.
• 대규모 컴퓨팅 자원이 필요 없이 이러한 불안정성을 연구하기 위해 소규모 프록시 모델을 개발한다.
• 학습률 변화에 걸쳐 안정성을 높이는 효과적인 조치(예: qk-layernorm, z-loss)를 식별하고 검증한다.
• 최적화기 및 모델 수정(예: 웜업, 가중치 감소, μ-Param)이 소규모 모델에서의 학습률 민감도에 미치는 영향을 분석한다.
• 활성화 및 기울기 노름의 스케일링 추세가 훈련 중 불안정성이 발생하기 전에 이를 예측할 수 있는지 탐색한다.

제안 방법

• Flax와 AdamW 최적화기를 사용하여 GPT-2 스타일 아키텍처의 소규모 디코더 전용 트랜스포머 모델을 훈련하고, 기본 초모수를 사용한다.
• 세 개의 학습률 주기 범위에서 학습률을 체계적으로 변화시키며 최종 평가 손실을 측정하여 학습률(LR) 민감도를 요약 통계량으로 계산한다.
• LR 민감도를 감소시키고 넓은 LR 범위에서의 훈련 안정성을 높이기 위해 qk-layernorm 및 z-loss 정규화를 도입하고 평가한다.
• 웜업 기간, 가중치 감소, μ-Param의 영향을 소규모 모델에서의 LR 민감도와 훈련 안정성에 대해 연구한다.
• 모델 스케일과 학습률에 따라 활성화 노름과 기울기 노름의 스케일링 행동을 분석하여 불안정성의 조기 징후를 탐지한다.
• 대규모 훈련에서 기본값인 AdamW epsilon이 너무 크기 때문에 업데이트가 너무 작아지는 점을 관찰하여 불안정성 메커니즘과 연결한다.

실험 결과

연구 질문

• RQ1 대규모 트랜스포머에서 관찰된 훈련 불안정성—예를 들어 attention 로짓 증가와 출력 로짓 발산—은 소규모 모델에서도 재현 가능한가?
• RQ2 기존의 대규모 훈련에서의 완화 기법(예: qk-layernorm, z-loss)이 소규모 프록시 환경에서도 여전히 효과적인가?
• RQ3 표준 최적화기 및 모델 수정(예: 웜업, 가중치 감소, μ-Param)은 소규모 모델에서의 학습률 민감도에 어떤 영향을 미치는가?
• RQ4 활성화 및 기울기 노름의 스케일링 추세는 불안정성이 나타나기 전에 이를 예측할 수 있는가?
• RQ5 AdamW epsilon 초모수는 불안정성, 특히 기울기 노름 스케일링과 관련하여 어떤 역할을 하는가?

주요 결과

• 고학습률로 훈련할 경우 소규모 모델에서도 attention 로짓 증가 및 출력 로짓 발산 등의 훈련 불안정성이 재현된다.
• qk-layernorm 및 z-loss 정규화는 학습률 민감도를 감소시켜 소규모 모델에서 학습률 변화 범위가 세 배수에 이르는 동안에도 안정적인 훈련을 가능하게 한다.
• 모델 깊이의 스케일링은 너비의 스케일링보다 학습률 민감도를 더 빠르게 증가시키며, 더 깊은 모델일수록 학습률 선택에 더 민감함을 시사한다.
• 기울기 노름은 모델 스케일과 학습률 증가에 따라 감소하며, 이는 기본값인 AdamW epsilon(1e-8)가 너무 크기 때문에 파라미터 업데이트가 너무 작아지는 결과를 초래함을 시사한다.
• 점진적 날카움과 적응형 최적화기 다이내믹스(예: β₂를 가진 Adam) 간의 상호작용이 주기적인 손실 폭발을 설명하며, β₂를 줄이면 조기 조절자가 더 빠르게 적응할 수 있어 이를 방지할 수 있다.
• 활성화 및 기울기 노름의 스케일링 행동은 불안정성이 발생하기 전에 이를 예측할 수 있으며, 이는 대규모 실행 없이도 훈련 안정성의 진단 도구로 활용될 수 있다.