一文要約: 学習率は、オプティマイザが1ステップでパラメータをどれほど動かすかを決めます。大きすぎれば振動・発散し、小さすぎればなかなか進みません。ただし「勾配の何パーセント」という意味ではなく、どのモデルにも通用する魔法の値もありません。

17.1 学習率とは何か

17.1.1 まず単純な SGD から考える

訓練では、損失が各パラメータに対してどう変化するかを計算し、オプティマイザがパラメータを更新します。最も単純な勾配降下法（SGD）なら、式は次のとおりです。

ここで:

• $\theta_t$ は更新前のパラメータ;
• $g_t$ は現在の勾配;
• $\eta$（eta）は学習率です。

パラメータを一つだけ取り出して計算してみます。

learning_rate = 0.1
old_weight    = 0.90
gradient      = -0.4

new_weight = 0.90 - 0.1 × (-0.4)
           = 0.94

勾配が負なので、その反対方向へ進むと重みは大きくなります。値は 0.90 から 0.94 に変わりました。

17.1.2 パーセントではない

lr = 0.1 は、勾配にスカラー 0.1 を掛けるという意味です。「勾配の10%を使う」と一般化することはできません。パラメータ、勾配、損失の尺度はモデルのパラメータ化によって変わるため、この 0.1 に生まれつきのパーセントの意味はありません。

また、上の式がそのまま表すのは SGD だけです。Adam や AdamW は、勾配を一次・二次モーメントで処理してから更新します。Transformer の訓練では、学習率はオプティマイザが作る更新の全体的なスケールを制御すると捉えるほうが正確です。

17.2 小さすぎる歩幅、大きすぎる歩幅、使える範囲

霧の中で山を下る場面を想像してください。足元でどちらが急な下りかは分かっても、山全体は見渡せません。

小さすぎる場合:

• 一歩が短く、損失の低下が非常に遅くなる;
• 計算予算が決まっていると「何も学んでいない」ように見える;
• 遅いからといって、必ず局所最小値に捕まったとは限らない。

大きすぎる場合:

• 谷底を飛び越え、両側を行き来する;
• 更新量が膨らみ、損失が発散したり NaN になったりする;
• 混合精度訓練では数値不安定性が早く現れることもある。

使える範囲にある場合:

• 訓練 loss は全体として下がる;
• 更新量が突然爆発しない;
• 同じ計算予算なら、極端に小さい学習率より先へ進める。

ここでは「使える範囲」と呼びます。唯一の「完璧な学習率」があるとは限らないからです。お椀型の曲線は直感のための模型にすぎません。実際の Transformer の損失は、鞍点、平坦な方向、パラメータの対称性を持つ高次元曲面です。2次元の絵から大域的最小値が見えるわけではありません。

Adam なら答えはいつも 3e-4、という昔からの冗談を見たことがあるかもしれません。実際にその値を使うレシピがあるから冗談になります。しかしモデル規模、batch、データ、パラメータ化、スケジュールが変われば、もう同じ問題ではありません。

17.3 実際に更新されるパラメータ

Transformer 全体を訓練する場合、主な学習可能パラメータは次の場所にあります。

1. Token embedding: vocab_size × d_model;
2. Attention の射影: Wq、Wk、Wv、Wo;
3. FFN の射影: 古典的 FFN の2行列、または gated FFN の複数の射影;
4. 正規化: LayerNorm の scale と bias、または RMSNorm の scale;
5. 出力射影: 隠れ状態から語彙 logits への写像。

出力射影が別の独立した行列とは限りません。GPT-2 型のモデルでは token embedding と重みを共有（weight tying）し、二つの場所が同じパラメータを参照することがよくあります。

PyTorch で、ある optimizer.step() によってパラメータが変わるには、通常は次を満たす必要があります。

• requires_grad=True である;
• パラメータから loss までの計算グラフに参加し、勾配を受け取る;
• その optimizer に登録されている。

param.grad is None なら、PyTorch のオプティマイザは通常そのパラメータを飛ばします。一方、値がゼロの勾配テンソルは別です。momentum や AdamW の weight decay によって値が変わる場合があります。

一回の optimizer.step() は全パラメータグループを巡りますが、すべてが同じ一歩を進むわけではありません。グループごとに別の学習率を設定でき、AdamW のモーメント履歴によって各座標の実効更新量も変わります。

17.4 本当に接続された PyTorch の1ステップ

次の小さな分類器は、あえて単純にしています。重要なのは、この model から loss が計算され、optimizer も同じパラメータを持っていることです。

import torch
import torch.nn.functional as F

torch.manual_seed(7)

model = torch.nn.Sequential(
    torch.nn.Linear(4, 8),
    torch.nn.GELU(),
    torch.nn.Linear(8, 5),
)
optimizer = torch.optim.AdamW(
    model.parameters(), lr=1e-3, weight_decay=0.01
)

x = torch.randn(6, 4)
targets = torch.tensor([0, 3, 1, 4, 2, 3])
before = model[0].weight.detach().clone()

optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = F.cross_entropy(logits, targets)
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()

max_change = (model[0].weight.detach() - before).abs().max().item()
assert max_change > 0
print(f"loss={loss.item():.4f}, grad_norm={grad_norm:.4f}")
print(f"max parameter change={max_change:.6f}")

順序にも意味があります。

1. zero_grad() で前のステップから蓄積した勾配を消す;
2. forward で logits と loss を得る;
3. loss.backward() で勾配を計算する;
4. 必要なら gradient clipping で勾配全体の norm を制限する;
5. optimizer.step() で初めてパラメータが変わる。

小数第4位に丸めた重みだけを見て、「更新された」と判断してはいけません。更新前のコピーを残し、後で差分を測れば直接確認できます。

17.5 SGD、Adam、AdamW は同じ更新ではない

17.5.1 SGD: 直感をつかみやすい基準

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

momentum を考えなければ、これは先ほどの $\theta_{t+1}=\theta_t-\eta g_t$ そのものです。学習率の説明には最適ですが、他のオプティマイザが生の勾配をそのまま使うとは限りません。

17.5.2 Adam: 勾配の履歴を使う

Adam は勾配の一次モーメント $m_t$ と二次モーメント $v_t$ を追跡します。bias correction の後、更新はおおよそ次の形です。

これにより、座標ごとに正規化された更新になります。蒔絵筆で金粉を置くとき、同じだけ手を動かしても筆先の粉量や直前の流れによって仕上がりは変わります。現在の方向だけでなく最近の状態を見て加減する、という点では Adam の直感に近いでしょう。ただし「各パラメータに最適な学習率を自動発見する」という意味ではありません。

17.5.3 AdamW: weight decay を分離する

AdamW の要点は、weight decay を適応的に正規化された loss gradient の更新から切り離したことです。実装の細部を省けば、一歩は次のように理解できます。

$\lambda$ は weight decay 係数です。一般的な条件では、SGD に対する L2 正則化と weight decay は等価になります。Adam のような適応的前処理の下では、通常は等価ではありません。

AdamW は Transformer 訓練の強力で一般的な基準ですが、どのタスク・規模でも唯一の正解ではありません。行列状の重みだけを decay し、bias と正規化パラメータを weight_decay=0 のグループへ入れるレシピもよくあります。

17.6 なぜ訓練中に学習率を変えるのか

固定学習率でも訓練できますが、多くの Transformer レシピは warmup + decay を使います。

• Warmup: 小さな値から peak まで徐々に上げる;
• Decay: その後で下げる。cosine decay は一般的な選択肢の一つ;
• 最終学習率: 0 にしてもよく、peak の一部を残してもよい。

Warmup の理由は、初期の勾配が必ず「間違っている」からではありません。開始直後は activation と gradient の尺度が落ち着いておらず、Adam のモーメント推定にも履歴がほとんどありません。いきなり peak を使うと update-to-weight ratio が不安定になることがあります。Warmup は有効な領域へ段階的に入る方法です。

次の関数では境界を明示します。訓練 step は 1 から total_steps まで。warmup_steps でちょうど peak_lr、最終 step でちょうど min_lr になります。

import math

def lr_at_step(step, total_steps, warmup_steps, peak_lr, min_lr):
    assert 1 <= step <= total_steps
    assert 0 < warmup_steps < total_steps

    if step <= warmup_steps:
        return peak_lr * step / warmup_steps

    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (peak_lr - min_lr) * cosine

for step in range(1, total_steps + 1):
    lr = lr_at_step(step, total_steps, warmup_steps, peak_lr, min_lr)
    for group in optimizer.param_groups:
        group["lr"] = lr

    optimizer.zero_grad(set_to_none=True)
    loss = compute_loss(model, batch)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

一つの batch につき optimizer.step() は一度だけです。cosine 区間も、総 step 数ではなく warmup 後に残った step 数で計算しています。

17.7 公開レシピは参考資料であって処方箋ではない

17.7.1 万能の数値を示す公開モデルはない

同じ系列では、モデルが大きいほど peak learning rate が下がる場合があります。しかしパラメータ数だけで値が決まるわけではありません。次は論文が報告した特定の訓練設定です。

数値は GPT-3 の Table 2.1 と LLaMA の Table 2 に基づきます。LLaMA は AdamW、betas=(0.9, 0.95)、weight_decay=0.1、gradient clipping 1.0、2,000 step の warmup、peak の10%まで下げる cosine schedule も報告しています。これらは一つのレシピとして読むべきで、3e-4 だけを抜き出すと条件が失われます。

17.7.2 Batch を2倍にすれば LR も2倍か

Batch を2倍にしたら学習率も2倍にする linear scaling rule は、Goyal らによる ResNet-50、ImageNet、大 batch の実験で、同期 SGD と warmup を組み合わせた条件から来ています。その条件では有用な出発点ですが、AdamW で LLM を訓練する普遍法則ではありません。

Global batch、sequence length、gradient accumulation、並列方式、model width のいずれかを変えたら、再検証すべき新しい設定だと考えましょう。loss だけでなく、gradient norm、update-to-weight ratio、同じ token 予算での validation の進み方も確認します。

17.7.3 正直な出発点

本書で作る小さなモデルなら、次を検証候補として始められます。「大規模モデルの標準値」ではありません。

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    betas=(0.9, 0.95),
    weight_decay=0.1,
)

例えば 1e-4、3e-4、6e-4 を、同じ token 数の短い run で比べます。最終的な最適値を探すためではなく、本番予算を使う前に桁が明らかに外れていないかを調べるためです。

17.8 すべての異常を学習率のせいにしない

Learning-rate finder は小規模な診断には使えます。しかし高価な LLM pretraining で、一度の sweep から永遠に使える「最良値」が得られるわけではありません。短い run で明らかに悪い候補を除き、実際の batch、精度、分散環境で確認します。

17.9 章のまとめ

• SGD では学習率が勾配を直接スケールする。AdamW ではモーメントで正規化した更新をスケールし、decoupled weight decay の強さにも関わる。
• 大きすぎれば振動・発散し、小さすぎれば訓練予算を浪費する。万能定数ではなく、安定して進む範囲を探す。
• パラメータは、訓練可能で、loss に接続され、勾配を受け取り、optimizer に属している必要がある。パラメータグループと適応的正規化により実効 step は異なる。
• Warmup + cosine decay は一般的だが必須ではない。境界、総 step、最終学習率を明示する。
• GPT-3 や LLaMA の学習率は、model scale、batch、残りの optimizer recipe と一緒に読む。

章末チェックリスト

• SGD の式で学習率を説明し、パーセントとは呼ばない
• SGD、Adam、AdamW の更新の違いを説明できる
• forward、backward、clipping、optimizer step が接続されたコードを書ける
• 境界が明確な warmup + cosine decay を実装できる
• 公開レシピを証拠として使い、万能デフォルトにはしない

第4部のまとめ

これで 第4部: 完全なアーキテクチャ は完了です。

Transformer の forward だけでなく、誤差が計算グラフを戻り、最後にパラメータを本当に変えるところまで追えるようになりました。

次章予告

第5部では、アーキテクチャをコードへ変えます。

• 第18章: model.py を書く
• 第19章: train.py を書く
• 第20章: inference.py を書く

まず model 定義から始め、これまでの各ブロックを実行可能なコードへ落とし込みます。