처음부터 만드는 LLM 32e부: 학습률 스케줄링과 최적화 개입

섹션별 상세

고정 학습률의 한계와 스케줄링의 필요성을 분석했다. 학습 초기에는 큰 보폭으로 손실 지형을 탐색해야 하지만, 최적점 근처에서는 보폭을 줄여야 정밀한 수렴이 가능하다. 특히 V자형 손실 곡선에서는 고정 학습률이 최소값 주변을 무한히 진동하게 만들 수 있음을 확인했다.

코사인 감쇠(Cosine Decay) 스케줄링을 도입했다. Chinchilla 논문의 권장 사항에 따라 학습률을 피크값의 1/10 수준까지 코사인 곡선을 그리며 낮추는 방식을 채택했다. 이는 초기 탐색과 후기 미세 조정 단계를 자연스럽게 연결하여 최적의 성능을 이끌어낸다.

학습 초기 안정성을 위한 선형 웜업(Linear Warmup)을 적용했다. 모델 초기화 직후의 무작위 에러로 인한 급격한 파라미터 업데이트를 방지하기 위해, 전체 학습 스텝의 약 5% 구간 동안 학습률을 0에서 피크값까지 서서히 높이는 과정을 추가했다.

DeepSeek 논문의 공식을 활용해 최적 피크 학습률을 산출했다. 연산 예산(FLOPs)과 학습률 사이의 상관관계를 정의한 스케일링 법칙 공식을 통해, 현재 프로젝트의 예산인 5e18 FLOPs에 적합한 학습률이 약 1.4e-3임을 도출하고 이를 실험에 적용했다.

python

C = 5e18 # Compute budget in FLOPs
eta_opt = 0.3118 * C**-0.1250
# Result: 0.00143385376262387

DeepSeek 논문의 공식을 활용하여 연산 예산에 따른 최적 학습률을 산출하는 코드

PyTorch의 SequentialLR을 이용한 구현 및 체크포인트 관리 방법을 정립했다. LinearLR과 CosineAnnealingLR을 순차적으로 연결하여 구현했으며, 학습 중단 후 재개 시 스케줄러의 state_dict를 함께 로드해야 학습률이 초기화되는 문제를 방지할 수 있음을 확인했다.

python

peak_lr = 0.0014
warmup_period = 1600
decay_period = 32000

optimizer = torch.optim.AdamW(
    model.parameters(), lr=peak_lr, weight_decay=0.1
)

warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
    optimizer, start_factor=0.00001, end_factor=1.0, total_iters=warmup_period
)

cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=decay_period, eta_min=peak_lr / 10
)

scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer, 
    schedulers=[warmup_scheduler, cosine_scheduler], 
    milestones=[warmup_period]
)

PyTorch의 SequentialLR을 사용하여 선형 웜업과 코사인 감쇠를 결합한 스케줄러 구현 예시

용어 해설

학습률 스케줄링(Learning Rate Scheduling): — 학습 과정 중에 학습률을 동적으로 변경하는 기법이다. 초기에는 큰 학습률로 손실 함수 공간을 빠르게 탐색하고, 후반부에는 학습률을 낮추어 최적점에 정밀하게 수렴하도록 돕는다.
코사인 어닐링(Cosine Annealing): — 학습률을 코사인 함수 곡선에 따라 감소시키는 스케줄링 방식이다. 초기 높은 학습률에서 시작해 부드럽게 감소하며, Chinchilla 논문 등에서 LLM 학습의 표준적인 기법으로 권장된다.
웜업(Warmup): — 학습 초기 단계에서 학습률을 0에 가까운 값에서 목표 피크값까지 점진적으로 높이는 과정이다. 모델 초기화 직후 발생하는 불안정한 그래디언트가 최적화 상태를 망치는 것을 방지한다.
친칠라 최적성(Chinchilla Optimality): — DeepMind의 연구에서 제시된 개념으로, 주어진 연산 예산(FLOPs) 내에서 모델 파라미터 수와 학습 데이터 토큰 수 사이의 최적의 균형을 의미한다.
AdamW: — Adam 최적화 알고리즘에서 가중치 감쇠(Weight Decay)를 수식적으로 분리하여 일반화 성능을 개선한 알고리즘이다. 현대적인 대형 언어 모델 학습에서 가장 널리 사용되는 옵티마이저이다.

기술

PyTorch
AdamW
GPT-2
SequentialLR

활용 사례

LLM Pre-training
Hyperparameter Optimization
Training Stability Improvement

언급된 리소스

문서Build a Large Language Model (from Scratch)

논문DeepSeek-V3 Technical Report