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

一言でいうと

事前学習は、「次のトークン当て」を数百万回繰り返す作業です。MiniMindのtrain_pretrain.pyは、AdamW・コサイン学習率・勾配クリッピングという標準的な部品で、このループを回します。このモジュールでは、そのループをCPU1台の大きさに縮小して自分で書き、損失が6.9から1の近くまで下がる曲線を読みます。どこで急に下がり、どこで遅くなり、学習率が大きいとどう壊れるのかを見ます。

なぜ必要なのか

学習ループは、たいていloss.backward(); optimizer.step()の2行に見えます。しかし、実際にモデルを一度でも最後まで学習してみた人は、その2行の周りの決定が結果を分けることを知っています。学習率が大きいと損失が下がったあとに跳ね上がり、シードを固定しないと昨日の結果を再現できず、検証損失を測らないと暗記と学習を区別できません。損失曲線は、これらすべてが表に出る唯一の窓です。MiniMindのREADMEは、64Mモデルの事前学習を、3090のカード1枚で1時間ちょっとと書いていますが、その1時間をきちんと使うには、最初の数分の曲線を読めなければなりません。

どう動くのか

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個を見て1学習ステップ進みます。clip_grad_norm_は、勾配ベクトル全体の長さが1を超えたら1に縮めます。ときどき跳ねるバッチ1つが、重みを壊さないようにするためです。GPUではbfloat16のautocastを使いますが、CPUではnullcontext()なので、fp32で回ります。

曲線の読み方では、私たちのモデルはln 1024 ≈ 6.93から出発して、数十学習ステップで3を下回ります。この急な区間で、モデルはどのトークンがよく出てくるかを学びます。文脈をまったく見ずに頻度だけを知っているモデルの損失は、ユニグラムエントロピー(私たちのコーパスで約4.6ナット)です。曲線がそれより下に下がるのは、前のトークンを見て、次を絞り込みはじめたという意味です。その後は遅くなります。文型・助詞・村と特産物の組み合わせのような長い規則は、ゆっくり学ぶからです。

再現性では、重みの初期化とバッチの選択のシードを、両方固定する必要があります。MiniMindのsetup_seedは、random・numpy・torchを一度に固定します。同じCPU・同じスレッド数なら、20学習ステップの損失が、1桁も違わずに出ます。

現場での姿

学習を数時間回す前に、数分の短い学習で確認することがあります。最初の損失がln(語彙)の近くか、数十学習ステップのうちに急に下がるか、検証損失が学習損失についてくるかです。最初の損失がおかしければデータかラベルが間違っていて、まったく下がらなければ、学習率が小さすぎるか、勾配が流れていません。逆に、損失が下がってから大きく跳ねたり、ある値に張りついて動かなくなったりしたら、学習率が大きすぎます。このラボでは、lrを0.05に上げて直接見ます。

損失が低いからといって、良いモデルとは限りません。私たちのコーパスは、型から打ち出した文なので、次のトークンがほぼ決まっていて、損失が1の近くまで下がります。実際のコーパスなら、同じ大きさのモデルは、ここまで下がりません。数字は、同じデータ・同じトークナイザーの中でしか比べられません。

次のラボですること

MiniMindのループを真似たtrain.pyを書き、同じシードで2回回して、1桁も違わないかを確かめます。300学習ステップを学習してチェックポイントと損失の記録を残し、曲線の形を数字で取り出し、ユニグラムエントロピーと比べ、学習率0.05が何を壊すかを見て、学習したモデルで続きを書かせてみます。