損失曲線はモデルが何を先に学ぶかを見せる

한국어 원문으로 표시합니다.

한 줄 요약

사전학습은 '다음 토큰 맞히기' 를 수백만 번 되풀이하는 일이다. MiniMind 의 train_pretrain.py 는 AdamW·코사인 학습률·기울기 자르기라는 표준 부품으로 이 루프를 돈다. 이 모듈에서는 그 루프를 CPU 한 대 크기로 줄여 직접 쓰고, 손실이 6.9 에서 1 근처까지 내려가는 곡선을 읽는다 — 어디서 급하게 떨어지고, 어디서 느려지고, 학습률이 크면 어떻게 망가지는지.

왜 이게 필요했나

학습 루프는 대부분 loss.backward(); optimizer.step() 두 줄로 보인다. 그런데 실제로 모델을 한 번이라도 끝까지 학습해 본 사람은 그 두 줄 주변의 결정이 결과를 가른다는 것을 안다. 학습률이 크면 손실이 내려가다 튀어 오르고, 시드를 고정하지 않으면 어제의 결과를 다시 낼 수 없고, 검증 손실을 재지 않으면 외우는 것과 배우는 것을 구별하지 못한다. 손실 곡선은 이 모든 것이 드러나는 유일한 창이다. MiniMind README 는 64M 모델의 사전학습을 3090 한 장으로 한 시간 남짓이라고 적는데, 그 한 시간을 제대로 쓰려면 처음 몇 분의 곡선을 읽을 줄 알아야 한다.

어떻게 동작하나

MiniMind 의 루프에서 요점만 추리면 이렇다.

for step, (input_ids, labels) in enumerate(loader, start=1):
    lr = get_lr(epoch * iters + step, epochs * iters, learning_rate)
    for g in optimizer.param_groups: g["lr"] = lr
    res = model(input_ids, labels=labels)
    loss = (res.loss + res.aux_loss) / accumulation_steps   # aux_loss 는 MoE 균형 손실, 밀집 모델은 0
    scaler.scale(loss).backward()
    if step % accumulation_steps == 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)   # 기본 1.0
        scaler.step(optimizer); scaler.update(); optimizer.zero_grad(set_to_none=True)

학습률. get_lr 은 lr × (0.1 + 0.45 × (1 + cos(π·t/T))) 다. 처음엔 lr 그대로, 끝에는 0.1·lr 로 코사인을 그리며 내려간다. 0 까지 내리지 않는 것이 특징이다. 기본 lr 은 5e-4 인데, 이 코스의 작은 모델은 3e-3 으로 훨씬 크게 잡는다 — 모델이 작을수록 큰 학습률을 견딘다.

기울기 누적과 자르기. 기본 배치 32 에 누적 8 이면 실제로는 256 개를 보고 한 걸음을 간다. clip_grad_norm_ 은 기울기 벡터 전체의 길이가 1 을 넘으면 1 로 줄인다 — 가끔 튀는 배치 하나가 가중치를 망가뜨리지 않게. GPU 에서는 bfloat16 autocast 를 쓰지만 CPU 에서는 nullcontext() 라 fp32 로 돈다.

곡선 읽기. 우리 모델은 ln 1024 ≈ 6.93 에서 출발해 수십 걸음 만에 3 아래로 떨어진다. 이 급한 구간에서 모델은 어떤 토큰이 자주 나오는지를 배운다. 문맥을 전혀 보지 않고 빈도만 아는 모델의 손실은 유니그램 엔트로피(우리 말뭉치에서 약 4.6 나트)다. 곡선이 그 아래로 내려간다는 것은 앞 토큰을 보고 다음을 좁히기 시작했다는 뜻이다. 그 뒤는 느리다 — 문형·조사·마을과 특산물의 짝처럼 긴 규칙은 천천히 배운다.

재현성. 가중치 초기화와 배치 뽑기의 시드를 둘 다 고정해야 한다. MiniMind 의 setup_seed 는 random·numpy·torch 를 한꺼번에 고정한다. 같은 CPU·같은 스레드 수라면 20걸음의 손실이 한 자리도 다르지 않게 나온다.

현장에서 만나는 모습

학습을 몇 시간 돌리기 전에 몇 분짜리 짧은 학습으로 확인할 것이 있다. 첫 손실이 ln(어휘) 근처인가, 수십 걸음 안에 급하게 떨어지는가, 검증 손실이 학습 손실을 따라오는가. 첫 손실이 이상하면 자료나 라벨이 틀렸고, 전혀 안 떨어지면 학습률이 너무 작거나 기울기가 흐르지 않는 것이다. 반대로 손실이 떨어지다 크게 튀거나 어느 값에 붙어 움직이지 않으면 학습률이 너무 크다 — 이 실습에서 lr 을 0.05 로 올려 직접 본다.

손실이 낮다고 좋은 모델인 것도 아니다. 우리 말뭉치는 틀에서 찍어 낸 문장이라 다음 토큰이 거의 정해져 있어 손실이 1 근처까지 내려간다. 실제 말뭉치라면 같은 크기의 모델이 이만큼 내려가지 않는다. 숫자는 같은 자료·같은 토크나이저 안에서만 견줄 수 있다.

다음 실습에서 할 것

MiniMind 의 루프를 본뜬 train.py 를 쓰고, 같은 시드로 두 번 돌려 한 자리도 다르지 않은지 본다. 300걸음을 학습해 체크포인트와 손실 기록을 남기고, 곡선의 모양을 숫자로 뽑고, 유니그램 엔트로피와 견주고, 학습률 0.05 가 무엇을 망가뜨리는지 보고, 학습한 모델로 이어 쓰기를 해 본다.