Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer
텐서 프로그램 V: 제로샷 하이퍼파라미터 전이를 통한 대규모 신경망 튜닝
Greg Yang, Edward J. Hu, Igor Babuschkin, et al. (2022)
Tensor Programs 이론에 기반한 최대 업데이트 파라미터화(μP, Maximal Update Parameterization)를 제안하여, 작은 모델에서 찾은 최적 하이퍼파라미터를 대규모 모델에 제로샷으로 전이(zero-shot transfer)할 수 있게 함으로써 대규모 모델 학습의 비용을 획기적으로 절감했다.
배경
대규모 신경망의 학습에서 학습률, 초기화 스케일, 가중치 감쇠 등 하이퍼파라미터 튜닝은 엄청난 계산 비용을 요구한다. GPT-3(175B)의 단일 학습 실행에 수백만 달러가 소요되므로, 여러 하이퍼파라미터 설정을 시도하는 것은 비현실적이다. 표준 파라미터화(SP, Standard Parameterization)에서는 모델 폭(width)이 변하면 최적 하이퍼파라미터도 변하여, 작은 프록시 모델에서 튜닝한 값을 큰 모델에 적용할 수 없었다. 이는 대규모 모델 학습이 항상 '추측'에 의존해야 한다는 근본적 문제를 야기했다.
핵심 아이디어
μP의 핵심 통찰은 신경망의 초기화 스케일과 레이어별 학습률을 폭에 대한 적절한 함수로 설정하면, 무한 폭 극한에서 모든 은닉 레이어의 활성화와 업데이트가 동일한 스케일을 유지하며, 이에 따라 최적 하이퍼파라미터가 폭에 무관(width-invariant)해진다는 수학적 결과이다. 구체적으로, 은닉 가중치의 초기화를 1/√(fan_in) 대신 1/fan_in으로 조정하고, 출력 레이어의 학습률을 폭에 반비례하게 스케일링하며, 출력 가중치를 0으로 초기화하는 등의 변경이 필요하다. 이 파라미터화 하에서는 width=128인 모델에서 찾은 학습률, 초기화 등이 width=8192인 모델에서도 최적에 가깝다.
방법론
Tensor Programs 프레임워크를 통해 신경망의 무한 폭 극한 행동을 분석한다. μP는 다음 규칙을 따른다: (1) 입력/출력 가중치는 O(1)로 초기화하되, 은닉 가중치는 O(1/√n)으로 초기화(n은 폭), (2) 은닉 레이어의 학습률은 기준과 동일하게 유지하되, 출력 레이어의 학습률은 1/n으로 스케일링, (3) 출력 가중치의 곱셈 상수를 1/n으로 설정. 이를 기존 학습 코드에 적용하려면 coord_check 유틸리티로 활성화 스케일을 확인하고 mup 라이브러리를 통해 자동 변환할 수 있다.
주요 결과
GPT-3 계열 모델에서 폭 128부터 8192까지 μP를 적용한 결과, 작은 모델에서 찾은 최적 학습률이 큰 모델에서도 최적에 가까웠다. 반면 표준 파라미터화(SP)에서는 최적 학습률이 폭에 따라 크게 변했다. 6.7B 파라미터 모델에서 40M 프록시에서 전이한 하이퍼파라미터가 직접 튜닝과 동등한 검증 손실을 달성하여, 수천 GPU-시간의 하이퍼파라미터 탐색 비용을 절약했다. 다양한 아키텍처(Transformer, ResNet)와 태스크(언어 모델, 이미지 분류)에서 전이 가능성이 확인되었다.
임팩트
μP는 대규모 모델 학습의 경제학을 근본적으로 변화시킬 잠재력을 가진 연구이다. Microsoft의 대규모 모델 학습에 실제 적용되었으며, Cerebras 등의 AI 칩 기업도 μP를 기본 파라미터화로 채택했다. 이론적으로도 신경망의 무한 폭 극한과 실질적 유한 네트워크 사이의 관계를 정립한 중요한 기여이며, 하이퍼파라미터 전이라는 새로운 연구 방향을 열었다. 이후 μTransfer(깊이 방향 전이), Depth-μP 등의 후속 연구로 확장되고 있다.