LLM 사전 학습에서 가중치 재조정을 통한 분산 제어

Variance Control via Weight Rescaling in LLM Pre-training

March 21, 2025

저자: Louis Owen, Abhay Kumar, Nilabhra Roy Chowdhury, Fabian Güra

cs.AI

초록

대규모 언어 모델(LLM) 사전 학습의 결과는 가중치 초기화 및 분산 제어 전략에 크게 의존합니다. 일반적으로 신경망에서 초기 분산 제어의 중요성은 잘 문서화되어 있지만, 특히 LLM 사전 학습 동안의 초기화 및 그 성장 관리를 다룬 문헌은 다소 부족한 실정입니다. 본 논문에서는 레이어 인덱스 재조정(Layer Index Rescaling, LIR) 가중치 초기화 기법과 목표 분산 재조정(Target Variance Rescaling, TVR) 분산 제어 전략을 소개합니다. 10억 파라미터 규모의 LLaMA 모델을 대상으로 한 실험에서, 이러한 기법을 통해 더 나은 분산 관리를 수행함으로써 다운스트림 작업 성능이 상당히 향상되었으며(일반적인 사전 학습 벤치마크에서 최대 4.6% 향상), 극단적인 활성화 값이 감소하여 양자화 및 저정밀도 학습과 관련된 문제를 완화할 수 있음을 입증했습니다. 우리의 코드는 https://github.com/bluorion-com/weight_rescaling에서 확인할 수 있습니다.

English

The outcome of Large Language Model (LLM) pre-training strongly depends on weight initialization and variance control strategies. Although the importance of initial variance control has been well documented in neural networks in general, the literature on initialization and management of its growth during LLM pre-training, specifically, is somewhat sparse. In this paper, we introduce the Layer Index Rescaling (LIR) weight initialization scheme, and the Target Variance Rescaling (TVR) variance control strategy. Experiments on a 1B parameter LLaMA model demonstrate that better variance management using these techniques yields substantial improvements in downstream task performance (up to 4.6% on common pre-training benchmarks) and reduces extreme activation values, thus mitigating challenges associated with quantization and low-precision training. Our code is available at: https://github.com/bluorion-com/weight_rescaling.