학습률 스케줄

큰 모델 학습에서, 학습률 스케줄은 학습 안정성과 최종 성능을 좌우해. Optax가 표준 스케줄을 다 제공해.

주요 스케줄

import optax
import matplotlib.pyplot as plt

# 1. Constant
sched_const = optax.constant_schedule(1e-3)

# 2. Linear warmup
sched_warmup = optax.linear_schedule(
    init_value=0.0,
    end_value=1e-3,
    transition_steps=1000,   # 1000 step 동안 0 → 1e-3
)

# 3. Cosine decay (warm restart 가능)
sched_cosine = optax.cosine_decay_schedule(
    init_value=1e-3,
    decay_steps=10_000,
    alpha=0.1,   # 최저 = init * 0.1
)

# 4. Warmup + cosine (현대 표준)
sched = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=1e-3,
    warmup_steps=1000,
    decay_steps=10_000,
    end_value=1e-5,
)

# 5. Exponential decay
sched_exp = optax.exponential_decay(
    init_value=1e-3,
    transition_steps=1000,
    decay_rate=0.5,
)

# 6. Polynomial
sched_poly = optax.polynomial_schedule(
    init_value=1e-3,
    end_value=1e-5,
    power=2.0,
    transition_steps=10_000,
)

그래프로 보기

steps = jnp.arange(15_000)
lrs = jnp.array([sched(s) for s in steps])

plt.plot(steps, lrs)
plt.xlabel("step"); plt.ylabel("learning rate")
plt.show()

warmup_cosine_decay_schedule의 모양:

peak ──╮
       │   ╲ (cosine)
       ╱    ╲
0 ────╯      ╲___ end_value
      ↑      ↑
   warmup    decay 끝

학습 코드 통합

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=3e-4,
    warmup_steps=1000,
    decay_steps=100_000,
    end_value=3e-5,
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.scale_by_adam(),
    optax.scale_by_schedule(schedule),
    optax.scale(-1.0),
)

# 학습 루프 — schedule 이 자동 적용
@jax.jit
def step(params, opt_state, batch):
    grads = jax.grad(loss_fn)(params, *batch)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state

스케줄의 스텝 카운터, opt_state 안에 자동 보존. 사용자가 따로 추적 안 해도 돼.

합성 스케줄

# 여러 단계 — 처음엔 warmup, 그 후 cosine, 그 후 constant
sched = optax.join_schedules(
    schedules=[
        optax.linear_schedule(0.0, 3e-4, 1000),   # warmup
        optax.cosine_decay_schedule(3e-4, 50_000, alpha=0.1),  # decay
        optax.constant_schedule(3e-5),             # 끝까지 유지
    ],
    boundaries=[1000, 51_000],
)

💡 스케줄 디버깅

새 스케줄을 학습에 쓰기 전에는 항상 그래프로 확인해. steps = jnp.arange(N)으로 스텝 축을 만들고 lrs = schedule(steps)로 한 번에 값을 구할 수 있어. 최고점과 warmup 길이, decay 모양이 의도와 맞는지 확인해. 잘못된 스케줄은 학습을 망가뜨리는 흔한 원인이야.

현대 LLM 학습에서는 warmup_cosine_decay나 warmup과 linear decay의 조합을 표준적으로 사용해. 최고 학습률은 모델 크기와 배치 크기에 따라 달라져 (Chinchilla / Llama 식 scaling rule).

Code

import optax
import jax.numpy as jnp

# Cosine decay: starts at init_value, decays to alpha over decay_steps
schedule = optax.cosine_decay_schedule(
    init_value=1e-3,
    decay_steps=10000,
    alpha=0.0,        # minimum learning rate
)

# Check values at different steps
print(f"Step 0: {schedule(0):.6f}")      # 0.001000
print(f"Step 5000: {schedule(5000):.6f}") # 0.000500
print(f"Step 10000: {schedule(10000):.6f}")# 0.000000

# Warmup + cosine decay (very common in practice)
schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,        # start from 0
    peak_value=1e-3,       # warm up to this
    warmup_steps=1000,     # linear warmup for 1000 steps
    decay_steps=50000,     # total steps including warmup
    end_value=1e-5,        # minimum LR at end
)

# Use schedule with an optimizer
optimizer = optax.adamw(learning_rate=schedule, weight_decay=0.01)

# Or compose with chain
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule),
)

# Exponential decay
schedule = optax.exponential_decay(
    init_value=1e-3,
    transition_steps=1000,
    decay_rate=0.96,
)

# Piecewise constant (manual step schedule)
schedule = optax.piecewise_constant_schedule(
    init_value=1e-3,
    boundaries_and_scales={
        5000: 0.1,   # multiply LR by 0.1 at step 5000
        8000: 0.1,   # multiply again at step 8000
    }
)

# Warm restarts (SGDR)
schedule = optax.sgdr_schedule([
    dict(init_value=1e-3, peak_value=1e-3,
         decay_steps=5000, warmup_steps=500),
    dict(init_value=1e-3, peak_value=5e-4,
         decay_steps=5000, warmup_steps=500),
])

💡 스케줄 디버깅

Code

External links

Exercise

Progress

댓글 0