곡률 제어: 안정적인 Transformer 학습을 위한 아키텍처 웜업 (Architecture Warm-Up)

수십억 개의 파라미터를 가진 Transformer를 학습시키는 것은 일시적인 손실 스파이크(loss spikes)와 발산(divergence)으로 인해 컴퓨팅 자원을 낭비하는 등 종종 불안정한 모습을 보입니다. 최근 개발된 안정성의 가장자리 (Edge of Stability, EoS) 이론이 (전처리된 (preconditioned)) 곡률(curvature)을 통해 최적화 방법의 안정성을 이해하고 제어할 수 있는 강력한 도구를 제공함에도 불구하고, 이러한 곡률 제어 방법들은 곡률 추정의 복잡성으로 인해 대규모 Transformer 학습에서 널리 사용되지 못하고 있습니다. 이를 위해, 본 연구에서는 Hessian-vector product를 이용한 파워 반복법(power iteration)의 웜스타트(warm-started) 변형을 기반으로, 가장 큰 (전처리된) Hessian 고유값(즉, 곡률)을 빠르게 추정하는 온라인 추정기를 최초로 도입합니다. 우리는 제안된 방법이 수십억 파라미터 규모에서도 반복당 곡률 추적을 가능하게 하면서도 더 정확하다는 것을 이론적으로 보여주고 경험적으로 검증합니다. 이 도구를 사용하여, 우리는 학습 불안정성이 전처리된 곡률의 급증과 일치하며, 곡률이 깊이(depth)에 따라 증가한다는 것을 발견했습니다. 이러한 관찰에 착안하여, 우리는 아키텍처 웜업(architecture warm-up)을 제안합니다. 이는 전처리된 Hessian을 신중하게 제어하고 학습을 안정화하기 위해 네트워크의 깊이를 점진적으로 성장시키는 방식입니다. 대규모 Transformer에 대한 실험을 통해, 우리의 접근 방식이 수렴 속도를 늦추지 않으면서도 기존의 최첨단 안정화 기술과 비교하여 효율적인 곡률 추적을 가능하게 하고 불안정성을 줄인다는 것을 입증했습니다.