[PYTORCH] Warmup Step이 학습 안정성에 미치는 5가지 영향과 해결 방법

Warmup Step

 

현업 딥러닝 엔지니어가 전하는 초대형 모델 학습의 필수 테크닉: 왜 초기 학습률 제어가 모델의 운명을 결정하는가?

---

1. Warmup Step이란 무엇이며 왜 중요한가?

딥러닝 모델, 특히 Transformer나 ResNet과 같이 층이 깊은 네트워크를 학습시킬 때, 초기 가중치는 무작위(Random)로 설정되어 있습니다. 이 상태에서 매우 높은 학습률(Learning Rate)을 적용하면 그래디언트가 폭주(Exploding)하거나, 가중치가 최적 해(Global Minimum)에서 너무 멀어져 학습이 불가능한 상태에 빠지기 쉽습니다. Warmup Step은 학습 초기에 매우 낮은 학습률에서 시작하여 설정한 목표 학습률까지 점진적으로 높여가는 과정을 말합니다. 이는 엔진을 예열하는 과정과 유사하며, 네트워크의 각 층이 초기 그래디언트에 적응할 시간을 부여하여 전체적인 학습 안정성을 획기적으로 개선합니다.

2. Warmup 유무에 따른 학습 특성 차이 및 비교

학습 초기 단계에서 Warmup이 적용되었을 때와 그렇지 않았을 때의 실무적 차이를 분석한 결과입니다.

비교 항목	No Warmup (급격한 시작)	With Warmup (점진적 시작)	비고
초기 Loss 변동성	매우 높음 (NaN 발생 위험)	매우 낮고 안정적임	학습 초기 1~2 epoch 기준
최종 성능(Accuracy)	상대적으로 낮거나 수렴 실패	안정적으로 최고 성능 도달	하이퍼파라미터 민감도 차이
그래디언트 분포	특정 레이어에 편중되거나 폭주	전체 레이어에 고르게 분포	Vanishing/Exploding 방지
Batch Size 의존성	Large Batch 시 학습 파괴 위험	Large Batch에서도 안정적	분산 학습 필수 요소
학습 시간	초기 수렴은 빠르나 불안정	수렴까지의 총 시간은 비슷	안정성 확보가 우선

3. 실무자를 위한 PyTorch 기반 Warmup 구현 Example (7가지)

단순한 개념 이해를 넘어, 실제 프로젝트에서 바로 복사하여 사용할 수 있는 7가지 핵심 예제 코드입니다.

Example 1: LambdaLR을 이용한 선형 Warmup (Linear Warmup)

가장 기본적이면서 강력한 방식으로, 지정된 스텝 동안 선형적으로 학습률을 증가시킵니다.

import torch
from torch.optim.lr_scheduler import LambdaLR

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
warmup_steps = 1000

def lr_lambda(current_step: int):
    if current_step < warmup_steps:
        return float(current_step) / float(max(1, warmup_steps))
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)
        

Example 2: 지수적 Warmup (Exponential Warmup)

초기 변화량을 더 미세하게 조정하고 싶을 때 사용합니다.

def lr_lambda_exp(current_step: int):
    if current_step < warmup_steps:
        return (current_step / warmup_steps) ** 2
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda_exp)
        

Example 3: Cosine Annealing과 Warmup의 결합

Warmup 이후에 자연스럽게 감쇄하는 스케줄러를 구성하는 실무적인 방법입니다.

from torch.optim.lr_scheduler import CosineAnnealingLR, SequentialLR

warmup_scheduler = LambdaLR(optimizer, lr_lambda=lambda step: step / warmup_steps)
cosine_scheduler = CosineAnnealingLR(optimizer, T_max=9000)

# SequentialLR을 통해 두 스케줄러 연결 (PyTorch 1.10+)
scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_steps])
        

Example 4: 훈련 루프 내 조건부 Warmup (Manual Implementation)

스케줄러 클래스를 사용하지 않고 직접 루프 내에서 처리하는 유연한 방식입니다.

base_lr = 1e-4
for step in range(total_steps):
    if step < warmup_steps:
        curr_lr = base_lr * (step / warmup_steps)
        for param_group in optimizer.param_groups:
            param_group['lr'] = curr_lr
    # ... train step ...
        

Example 5: HuggingFace Style Warmup (Transformers 라이브러리 활용)

NLP 모델 학습 시 가장 널리 쓰이는 표준적인 Warmup 방식입니다.

from transformers import get_linear_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=500, 
    num_training_steps=10000
)
        

Example 6: 단계별(Stepwise) Warmup

특정 구간마다 계단식으로 학습률을 올리는 독특한 시나리오를 위한 구현입니다.

def stepwise_warmup(step):
    if step < 200: return 0.2
    if step < 500: return 0.5
    if step < 1000: return 0.8
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=stepwise_warmup)
        

Example 7: Warmup 중 Gradient Clipping 적용 해결

Warmup 구간에서도 발생할 수 있는 이상치를 방어하기 위한 필수 안전 장치입니다.

loss.backward()
# Warmup 중에는 더 엄격하게 클리핑을 적용하는 전략
clip_value = 0.5 if current_step < warmup_steps else 1.0
torch.nn.utils.clip_grad_norm_(model.parameters(), clip_value)
optimizer.step()
scheduler.step()
        

4. Warmup Step 설정 시 발생하는 흔한 문제와 해결 전략

단순히 Warmup을 적용한다고 모든 문제가 해결되지는 않습니다. 상황에 따른 최적화 전략을 공유합니다.

• 문제: Warmup 종료 후 Loss가 갑자기 튈 때
• 해결: Warmup 종료 지점의 학습률과 이후 스케줄러의 시작 학습률이 일치하는지 확인하십시오. 또한 Warmup 기간을 총 학습 스텝의 5~10% 정도로 더 길게 설정하는 것이 도움이 됩니다.
• 문제: Batch Size를 키웠는데 학습이 안 될 때
• 해결: "Linear Scaling Rule"에 따라 Batch Size가 커지면 학습률도 높여야 합니다. 이때 Warmup Step도 비례해서 늘려야 초기 불안정성을 극복할 수 있습니다.
• 문제: Warmup 기간 중 Loss가 줄지 않을 때
• 해결: 초기 학습률이 너무 낮게 설정되었을 수 있습니다. Warmup 시작점을 $0$이 아닌 목표 학습률의 $1/100$ 정도로 설정해 보세요.

5. 결론 및 요약

Warmup Step은 단순한 기법이 아니라, 현대적인 고성능 딥러닝 모델 학습의 기반입니다. 특히 거대 언어 모델(LLM)이나 고해상도 이미지 생성 모델에서 Warmup 없이 학습을 시작하는 것은 사실상 불가능에 가깝습니다. 위에서 제시한 7가지 예제 코드를 바탕으로 본인의 데이터셋과 모델 구조에 최적화된 Warmup 전략을 수립하시기 바랍니다.

---

내용 출처 및 참고 문헌

• Goyal et al., "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour", Facebook AI Research.
• Vaswani et al., "Attention Is All You Need", NeurIPS 2017.
• PyTorch 공식 문서: torch.optim.lr_scheduler API Reference.
• HuggingFace Documentation: Optimization methods for Transformers.