학습률 스케줄과 워밍업

학습률(learning rate)은 단일 하이퍼파라미터(hyperparameter) 중 가장 중요합니다. 아키텍처(architecture)도, 데이터셋 크기(dataset size)도, 활성화 함수(activation function)도 아닙니다. 학습률입니다. 단 하나만 조정(tune)한다면 바로 이것을 조정합니다.

유형: Build

언어: Python

선수 지식: Lesson 03.06 옵티마이저(Optimizers), Lesson 03.08 가중치 초기화(Weight Initialization)

소요 시간: 약 90분

학습 목표

• 상수(constant), 계단형 감쇠(step decay), 코사인 어닐링(cosine annealing), 워밍업 + 코사인(warmup + cosine), 1cycle 다섯 가지 학습률 스케줄(learning rate schedule)을 처음부터 구현합니다.
• 학습률 선택에서 나타나는 세 가지 실패 모드(failure mode)인 발산(divergence, 너무 큼), 정체(stalling, 너무 작음), 진동(oscillation, 감쇠 없음)을 보여 줍니다.
• Adam 계열 옵티마이저(optimizer)에 워밍업(warmup)이 필요한 이유와 초기 학습(early training)을 안정화하는 방식을 설명합니다.
• 같은 과제(task)에서 다섯 스케줄의 수렴 속도(convergence speed)를 비교하고 학습 예산(training budget)에 맞는 스케줄을 선택합니다.

문제

학습률을 0.1로 설정해 보겠습니다. 학습(training)이 발산하고 손실(loss)은 단 3 스텝(step) 만에 무한대로 튑니다. 0.0001로 설정하면 100 에폭(epoch)이 지나도 모델(model)은 무작위(random) 초기 상태에서 거의 움직이지 않습니다. 0.01로 설정하면 50 에폭 동안 잘 작동하다가, 최솟값(minimum) 근처에서 스텝이 너무 커서 손실이 끊임없이 진동(oscillation)합니다.

최적(optimal) 학습률은 고정된 값이 아닙니다. 학습 도중 계속 바뀝니다. 초반에는 큰 스텝으로 빠르게 이동하고 싶고, 후반에는 작은 스텝으로 최솟값에 정착하고 싶습니다. 정확도(accuracy) 90% 모델과 95% 모델의 차이는 종종 스케줄 하나에서 결정됩니다.

최근 3년 동안 발표된 주요 모델은 모두 학습률 스케줄을 사용합니다. Llama 3는 정점(peak) lr=3e-4, 워밍업 2000 스텝, 코사인 감쇠(cosine decay) 3e-5까지를 사용했습니다. GPT-3는 lr=6e-4와 3억 7500만 토큰(375 million tokens)에 걸친 워밍업을 사용했습니다. 이는 임의의 선택이 아니라 수많은 하이퍼파라미터 탐색(hyperparameter sweep)의 결과로, 비용도 수백만 달러에 달합니다.

스케줄을 이해해야 하는 이유는 기본값(default)이 여러분의 문제에 맞지 않을 수 있기 때문입니다. 사전 학습된 모델(pretrained model)을 미세 조정(fine-tune)할 때의 올바른 스케줄은 처음부터 학습(scratch training)할 때와 다릅니다. 배치 크기(batch size)를 키우면 워밍업 기간도 바뀌어야 합니다. 1만 번째 스텝에서 학습이 깨졌을 때, 그것이 스케줄 문제인지 다른 문제인지 판단할 수 있어야 합니다.

개념

상수 학습률(Constant Learning Rate)

가장 단순한 방식입니다. 숫자 하나를 골라 모든 스텝에 같은 값을 사용합니다.

lr(t) = lr_0

대개 최적이 아닙니다. 학습 후반에는 값이 너무 높아 최솟값(minimum) 주변에서 진동하고, 학습 초반에는 너무 낮아 연산(compute)을 낭비하게 됩니다. 작은 모델을 디버깅(debugging)할 때에는 충분하지만, 한 시간 이상 학습하는 작업에는 적합하지 않습니다.

계단형 감쇠(Step Decay)

ResNet 시대에 널리 쓰였던 전통적인 방식입니다. 정해진 에폭마다 학습률을 일정 비율(보통 10배)만큼 떨어뜨립니다.

lr(t) = lr_0 * gamma^(floor(epoch / step_size))

gamma = 0.1, step_size = 30이면 30 에폭마다 10배씩 줄어듭니다. ResNet-50은 이 방식을 채택했고, lr=0.1에서 시작해 30, 60, 90 에폭에 10배씩 떨어뜨렸습니다.

문제는 최적 감쇠 시점(decay point)이 데이터셋(dataset)과 아키텍처(architecture)에 따라 달라, 새로운 문제로 옮기면 다시 조정해야 한다는 점입니다. 전환(transition)이 갑작스러워서 비율이 갑자기 바뀌는 순간 손실이 급등(spike)할 수도 있습니다.

코사인 어닐링(Cosine Annealing)

최대 학습률(maximum learning rate)에서 최소 학습률(minimum learning rate)까지, 코사인 곡선(cosine curve)을 따라 부드럽게 감소시킵니다.

lr(t) = lr_min + 0.5 * (lr_max - lr_min) * (1 + cos(pi * t / T))

여기서 t는 현재 스텝, T는 전체 스텝 수입니다.

t=0이면 코사인 항이 1이라 lr = lr_max, t=T이면 코사인 항이 -1이라 lr = lr_min이 됩니다. 초반에는 천천히 감쇠하다가 중간에 가속되고, 끝부분에서 다시 완만해집니다.

최근 학습에서 사실상의 기본값(default)으로 자리잡은 방식입니다. lr_max와 lr_min 이외에 추가로 조정할 하이퍼파라미터가 없습니다. 코사인 모양은 "학습의 대부분은 중반에 일어난다"는 경험적 관찰과 잘 맞으며, 가장 중요한 시기에 적절한 스텝 크기를 유지해 줍니다.

워밍업(Warmup): 작게 시작하는 이유

Adam을 비롯한 적응형 옵티마이저(adaptive optimizer)는 기울기(gradient)의 평균과 분산에 대한 이동 추정값(running estimate)을 유지합니다. 스텝 0에서는 이 추정값이 0으로 초기화되어 있고, 초반의 몇 개 기울기 갱신(update)은 아직 신뢰할 수 없는 통계량(statistics)에 기반합니다. 이때 학습률이 크면 모델은 크지만 방향이 어긋난 스텝을 밟게 됩니다.

워밍업은 이 문제를 해결합니다. 매우 작은 학습률(주로 lr_max / warmup_steps 또는 0)에서 시작해, 처음 N 스텝 동안 lr_max까지 선형(linear)으로 끌어올립니다. 학습률이 완전히 도달할 즈음이면 Adam의 통계량도 안정됩니다.

lr(t) = lr_max * (t / warmup_steps)     for t < warmup_steps

일반적인 워밍업 비율은 전체 학습 스텝의 1-5%입니다. Llama 3는 약 1.8조 토큰(1.8 trillion tokens) 동안 학습하면서 2000 스텝의 워밍업을 사용했고, GPT-3는 3억 7500만 토큰에 걸쳐 워밍업을 진행했습니다.

선형 워밍업 + 코사인 감쇠(Linear Warmup + Cosine Decay)

요즘 가장 널리 쓰이는 기본 조합입니다. 선형으로 끌어올린 뒤 코사인으로 감쇠합니다.

if t < warmup_steps:
    lr(t) = lr_max * (t / warmup_steps)
else:
    progress = (t - warmup_steps) / (total_steps - warmup_steps)
    lr(t) = lr_min + 0.5 * (lr_max - lr_min) * (1 + cos(pi * progress))

Llama, GPT, PaLM을 비롯한 대부분의 최신 트랜스포머(transformer)가 이 패턴을 사용합니다. 워밍업이 초기 불안정성(early instability)을 막아 주고, 코사인 감쇠가 모델을 좋은 최솟값으로 정착시킵니다.

1cycle 정책(1cycle Policy)

Leslie Smith가 2018년에 발견한 방식입니다. 학습 전반부에는 낮은 값에서 높은 값으로 학습률을 끌어올리고, 후반부에는 다시 매우 낮은 값으로 내립니다. 직관에 반하는 선택입니다 — 왜 학습 도중에 일부러 학습률을 올리는 것일까요?

이론적 설명은 이렇습니다. 높은 학습률은 최적화 궤적(optimization trajectory)에 잡음(noise)을 더해 일종의 정규화(regularization)처럼 작동합니다. 학습률이 올라가는 동안 모델은 손실 지형(loss landscape)을 더 넓게 탐색해 더 좋은 분지(basin)를 찾고, 이후 내려가는 구간에서 그 분지 안을 세밀히 다듬습니다.

Phase 1 (0 to T/2):    lr ramps from lr_max/25 to lr_max
Phase 2 (T/2 to T):    lr ramps from lr_max to lr_max/10000

1cycle은 고정된 연산 예산(fixed compute budget)에서 코사인 어닐링보다 빠르게 학습될 때가 많습니다. 대신 전체 스텝 수를 미리 알고 있어야 한다는 단점이 있습니다.

스케줄 모양 비교(Schedule Shapes)

graph LR
    subgraph "Constant"
        C1["lr"] --- C2["lr"] --- C3["lr"]
    end

    subgraph "Step Decay"
        S1["0.1"] --- S2["0.1"] --- S3["0.01"] --- S4["0.001"]
    end

    subgraph "Cosine Annealing"
        CS1["lr_max"] --> CS2["gradual"] --> CS3["steep"] --> CS4["lr_min"]
    end

    subgraph "Warmup + Cosine"
        WC1["0"] --> WC2["lr_max"] --> WC3["cosine"] --> WC4["lr_min"]
    end

스케줄 선택 흐름도(Decision Flowchart)

flowchart TD
    Start["LR schedule 선택"] --> Know{"total training steps를 아나요?"}

    Know -->|"Yes"| Budget{"Compute budget?"}
    Know -->|"No"| Constant["constant LR<br/>with manual decay"]

    Budget -->|"Large (days/weeks)"| WarmCos["Warmup + Cosine Decay<br/>(Llama/GPT default)"]
    Budget -->|"Small (hours)"| OneCycle["1cycle Policy<br/>(fast convergence)"]
    Budget -->|"Moderate"| Cosine["Cosine Annealing<br/>(safe default)"]

    WarmCos --> Warmup["Warmup = 1-5% of steps"]
    OneCycle --> FindLR["LR range test로 lr_max 찾기"]
    Cosine --> MinLR["lr_min = lr_max / 10"]

공개 모델의 실제 숫자(Real Numbers from Published Models)

graph TD
    subgraph "Published LR Configs"
        L3["Llama 3 (405B)<br/>Peak: 3e-4<br/>Warmup: 2000 steps<br/>Schedule: Cosine to 3e-5"]
        G3["GPT-3 (175B)<br/>Peak: 6e-4<br/>Warmup: 375M tokens<br/>Schedule: Cosine to 0"]
        R50["ResNet-50<br/>Peak: 0.1<br/>Warmup: none<br/>Schedule: Step decay x0.1 at 30,60,90"]
        B["BERT (340M)<br/>Peak: 1e-4<br/>Warmup: 10K steps<br/>Schedule: Linear decay"]
    end

만들어 보기

Step 1: 스케줄 함수(Schedule Functions)

각 함수는 현재 스텝을 입력으로 받아 해당 스텝의 학습률을 반환합니다.

import math


def constant_schedule(step, lr=0.01, **kwargs):
    return lr

def step_decay_schedule(step, lr=0.1, step_size=100, gamma=0.1, **kwargs):
    return lr * (gamma ** (step // step_size))

def cosine_schedule(step, lr=0.01, total_steps=1000, lr_min=1e-5, **kwargs):
    if step >= total_steps:
        return lr_min
    return lr_min + 0.5 * (lr - lr_min) * (1 + math.cos(math.pi * step / total_steps))

def warmup_cosine_schedule(step, lr=0.01, total_steps=1000, warmup_steps=100, lr_min=1e-5, **kwargs):
    if total_steps <= warmup_steps:
        return lr * (step / max(warmup_steps, 1))
    if step < warmup_steps:
        return lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return lr_min + 0.5 * (lr - lr_min) * (1 + math.cos(math.pi * progress))

def one_cycle_schedule(step, lr=0.01, total_steps=1000, **kwargs):
    mid = max(total_steps // 2, 1)
    if step < mid:
        return (lr / 25) + (lr - lr / 25) * step / mid
    else:
        progress = (step - mid) / max(total_steps - mid, 1)
        return lr * (1 - progress) + (lr / 10000) * progress

워밍업 + 코사인은 워밍업 구간에서 선형 증가(linear ramp)를 사용하고, 이후 구간에 코사인 감쇠를 적용합니다. 1cycle은 학습 전반부에 학습률을 끌어올리고 후반부에 끌어내립니다.

Step 2: 스케줄 시각화(Visualize All Schedules)

텍스트 플롯(text plot)으로 각 스케줄이 학습 동안 어떻게 변화하는지 출력합니다. lr_max, 워밍업, lr_min이 실제로 어떤 곡선을 만드는지 눈으로 확인합니다.

def visualize_schedule(name, schedule_fn, total_steps=500, **kwargs):
    steps = list(range(0, total_steps, total_steps // 20))
    if total_steps - 1 not in steps:
        steps.append(total_steps - 1)

    lrs = [schedule_fn(s, total_steps=total_steps, **kwargs) for s in steps]
    max_lr = max(lrs) if max(lrs) > 0 else 1.0

    print(f"\n{name}:")
    for s, lr_val in zip(steps, lrs):
        bar_len = int(lr_val / max_lr * 40)
        bar = "#" * bar_len
        print(f"  Step {s:4d}: lr={lr_val:.6f} {bar}")

Step 3: 학습 네트워크(Training Network)

원형 데이터셋(circle dataset)을 사용한 단순한 2-레이어 신경망(two-layer network)을 준비합니다. 이전 강의와 같은 네트워크 구조이지만, 이번에는 스케줄을 바꿔 가며 학습할 수 있도록 학습률을 매 스텝마다 갱신합니다.

import random


def sigmoid(x):
    x = max(-500, min(500, x))
    return 1.0 / (1.0 + math.exp(-x))


def relu(x):
    return max(0.0, x)


def relu_deriv(x):
    return 1.0 if x > 0 else 0.0


def make_circle_data(n=200, seed=42):
    random.seed(seed)
    data = []
    for _ in range(n):
        x = random.uniform(-2, 2)
        y = random.uniform(-2, 2)
        label = 1.0 if x * x + y * y < 1.5 else 0.0
        data.append(([x, y], label))
    return data


def train_with_schedule(schedule_fn, schedule_name, data, epochs=300, base_lr=0.05, **kwargs):
    random.seed(0)
    hidden_size = 8
    total_steps = epochs * len(data)

    std = math.sqrt(2.0 / 2)
    w1 = [[random.gauss(0, std) for _ in range(2)] for _ in range(hidden_size)]
    b1 = [0.0] * hidden_size
    w2 = [random.gauss(0, std) for _ in range(hidden_size)]
    b2 = 0.0

    step = 0
    epoch_losses = []

    for epoch in range(epochs):
        total_loss = 0
        correct = 0

        for x, target in data:
            lr = schedule_fn(step, lr=base_lr, total_steps=total_steps, **kwargs)

            z1 = []
            h = []
            for i in range(hidden_size):
                z = w1[i][0] * x[0] + w1[i][1] * x[1] + b1[i]
                z1.append(z)
                h.append(relu(z))

            z2 = sum(w2[i] * h[i] for i in range(hidden_size)) + b2
            out = sigmoid(z2)

            error = out - target
            d_out = error * out * (1 - out)

            for i in range(hidden_size):
                d_h = d_out * w2[i] * relu_deriv(z1[i])
                w2[i] -= lr * d_out * h[i]
                for j in range(2):
                    w1[i][j] -= lr * d_h * x[j]
                b1[i] -= lr * d_h
            b2 -= lr * d_out

            total_loss += (out - target) ** 2
            if (out >= 0.5) == (target >= 0.5):
                correct += 1
            step += 1

        avg_loss = total_loss / len(data)
        epoch_losses.append(avg_loss)

    return epoch_losses

Step 4: 스케줄 비교와 학습률 민감도(Compare Schedules and LR Sensitivity)

먼저 다섯 가지 스케줄(상수, 계단형 감쇠, 코사인, 워밍업 + 코사인, 1cycle)을 같은 네트워크에 적용해 시작 손실(start loss), 중간 손실(mid loss), 마지막 손실(end loss), 최저 손실(best loss)을 비교합니다. 이어서 상수 스케줄에서 학습률 1.0, 0.1, 0.01, 0.001, 0.0001을 비교해 너무 높음, 너무 낮음, 적당함의 세 가지 상태를 직접 확인합니다. 손실이 발산하거나 거의 움직이지 않는 경우를 한눈에 구분할 수 있습니다.

def compare_schedules(data):
    configs = [
        ("Constant", constant_schedule, {}),
        ("Step Decay", step_decay_schedule, {"step_size": 15000, "gamma": 0.1}),
        ("Cosine", cosine_schedule, {"lr_min": 1e-5}),
        ("Warmup+Cosine", warmup_cosine_schedule, {"warmup_steps": 3000, "lr_min": 1e-5}),
        ("1cycle", one_cycle_schedule, {}),
    ]

    print(f"\n{'Schedule':<20} {'Start Loss':>12} {'Mid Loss':>12} {'End Loss':>12} {'Best Loss':>12}")
    print("-" * 70)

    for name, schedule_fn, extra_kwargs in configs:
        losses = train_with_schedule(schedule_fn, name, data, epochs=300, base_lr=0.05, **extra_kwargs)
        mid_idx = len(losses) // 2
        best = min(losses)
        print(f"{name:<20} {losses[0]:>12.6f} {losses[mid_idx]:>12.6f} {losses[-1]:>12.6f} {best:>12.6f}")

def lr_sensitivity(data):
    learning_rates = [1.0, 0.1, 0.01, 0.001, 0.0001]

    print("\nLR Sensitivity (constant schedule, 100 epochs):")
    print(f"  {'LR':>10} {'Start Loss':>12} {'End Loss':>12} {'Status':>15}")
    print("  " + "-" * 52)

    for lr in learning_rates:
        losses = train_with_schedule(constant_schedule, f"lr={lr}", data, epochs=100, base_lr=lr)
        start = losses[0]
        end = losses[-1]

        if end > start or math.isnan(end) or end > 1.0:
            status = "DIVERGED"
        elif end > start * 0.9:
            status = "BARELY MOVED"
        elif end < 0.15:
            status = "CONVERGED"
        else:
            status = "LEARNING"

        end_str = f"{end:.6f}" if not math.isnan(end) else "NaN"
        print(f"  {lr:>10.4f} {start:>12.6f} {end_str:>12} {status:>15}")

사용하기

PyTorch는 torch.optim.lr_scheduler에서 다양한 스케줄러(scheduler)를 제공합니다.

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = optim.Adam(model.parameters(), lr=3e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=1000, eta_min=1e-5)

for step in range(1000):
    loss = train_step(model, optimizer)
    scheduler.step()

워밍업 + 코사인은 HuggingFace의 get_cosine_schedule_with_warmup을 자주 사용합니다.

from transformers import get_cosine_schedule_with_warmup

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=2000,
    num_training_steps=100000,
)

HuggingFace의 이 함수는 Llama와 GPT 미세 조정 스크립트 대부분이 사용하는 표준입니다. 어떤 스케줄을 골라야 할지 확신이 없다면 워밍업 + 코사인을 사용하고, 워밍업을 전체 스텝의 3-5%로 시작합니다. 거의 모든 상황에서 잘 작동합니다.

산출물 만들기

이 강의의 산출물은 다음과 같습니다.

• outputs/prompt-lr-schedule-advisor.md: 학습 환경에 맞는 학습률 스케줄과 하이퍼파라미터를 추천하는 프롬프트(prompt)

연습문제

1. (쉬움) 지수 감쇠(exponential decay)를 구현합니다. lr(t) = lr_0 * gamma^t, gamma = 0.999입니다. 원형 데이터셋에서 코사인 어닐링과 비교합니다.
2. (중간) Leslie Smith의 학습률 탐색 시험(learning rate range test)을 구현합니다. 몇백 스텝 동안 학습률을 1e-7에서 1까지 지수적으로 늘리며 손실 대 학습률(loss vs LR) 그래프를 그립니다. 최적 최대 학습률은 손실이 증가하기 직전 값입니다.
3. (중간) 워밍업 + 코사인에서 워밍업 길이를 전체 스텝의 0%, 1%, 5%, 10%, 20%로 바꿔 가장 안정적인 지점(sweet spot)을 찾습니다.
4. (어려움) 따뜻한 재시작이 있는 코사인 어닐링(cosine annealing with warm restarts; SGDR)을 구현합니다. T 스텝마다 학습률을 lr_max로 재설정(reset)한 뒤 다시 감쇠시키고, 표준 코사인과 비교합니다.
5. (어려움) 학습 손실을 관찰하다가 손실이 안정되면 워밍업에서 코사인으로 자동 전환하고, 손실이 오래 정체(plateau)되면 학습률을 낮추는 "스케줄 외과의(schedule surgeon)"를 만듭니다.

핵심 용어

더 읽을거리

• Loshchilov & Hutter, "SGDR" (2017) — 코사인 어닐링과 따뜻한 재시작을 처음 소개한 논문입니다.
• Smith, "Super-Convergence" (2018) — 1cycle 정책을 다룬 핵심 논문입니다.
• Touvron et al., "Llama 2" (2023) — 대규모 모델(large-scale model)에서 사용된 워밍업 + 코사인 스케줄의 실제 사례를 보여 줍니다.
• Goyal et al., "Accurate, Large Minibatch SGD" (2017) — 큰 배치 학습(large batch training)에서의 선형 스케일링 규칙과 워밍업을 다룹니다.