【機械学習】「Loss is NaN」の絶望から学ぶ勾配爆発と学習率チューニングの真髄

プログラミング

ニューラルネットワークを自作して学習ループ(Train Loop)を回し始めたときの緊張感、たまりませんよね。最初は順調にLoss(損失値)が下がり、「おっ、ちゃんと学習できてる!」と安心した矢先、突如としてコンソール画面に表示される最悪の文字列があります。

Loss: NaN(Not a Number、非数)

一度これが表示されると、モデルの重み(Weight)はすべて破壊され、二度と元の学習ルートには戻れません。まさにAIの突然死です。
本記事では、この心臓に悪い「Loss NaN崩壊」がなぜ起こるのかというメカニズムを解剖し、勾配クリッピング(Gradient Clipping)や学習率(Learning Rate)の適切なスケジューリングといった、トップKagglerたちが息をするように行っている「学習の安定化手法」を、初心者にも分かりやすく深掘りしていきます!

1. なぜ数値が「NaN(非数)」になってしまうのか?

LossがNaNになる原因は、主に以下の3つのパターンに分類されます。
プログラムのバグではなく、数学的な限界を突破してしまったときに発生することがほとんどです。

  • 原因1:ゼロ除算(Zero Division)や対数関数のマイナス入力
    例えば log(0) は数学的に負の無限大(-inf)になります。これがその後の計算式で無理やり足し算・掛け算されると途端にNaNに化けます。これを防ぐため、CrossEntropyLossの実装などでは、ゼロにならないように微小な値(epsilon: 1e-8 など)を足すのが定石です。
  • 原因2:異常なデータ(欠損値)の混入
    入力する画像やデータセット自体に、元から NaNInf のピクセル値が含まれているケースです。前処理の段階で np.isnan().any() 等を使ってチェックしていないと、そのままネットワークに毒が回り、モデル全体を汚染してしまいます。
  • 原因3:【圧倒的に多い】勾配爆発(Exploding Gradients)
    これが最も多い原因です!学習率(Learning Rate)が大きすぎると、重みの更新幅が巨大になりすぎます。誤差逆伝播(バックプロパゲーション)で数百のレイヤーを経由して掛け算が繰り返されるうちに数値がオーバーフローし、Float32の限界(約3.4×10^38)を突破して NaN となってしまう現象です。

2. 「勾配クリッピング」による強制ギプス

勾配爆発(一歩の歩幅が大きすぎてジャンプしてしまい、宇宙の果てまで飛んでいってしまう状態)を防ぐ最も強力かつ物理的なアプローチが「勾配クリッピング(Gradient Clipping)」です。

非常にシンプルで、「もし勾配(更新する歩幅)の大きさが一定の値(例:1.0)を超えたら、強制的に1.0に切り詰める」という荒技です。
一見乱暴に見えますが、RNN(再帰型ニューラルネットワーク)や、現代主流であるTransformerなどの深いモデルでは、この安全装置がないとほぼ確実に爆発するため、実質的に必須のテクニックとなっています。

# PyTorchでの勾配クリッピング実装の基本パターン
outputs = model(inputs)
loss = criterion(outputs, labels)
# 通常通り勾配を計算する(ここで勾配が超巨大になるかもしれない)
optimizer.zero_grad()
loss.backward()
# 【ここが超重要!】
# 重みを更新する「直前」に、モデル内の全勾配のL2ノルムを最大値1.0に制限する!
# これにより、どれだけ誤差が大きくても「歩幅」を安全圏内に強制的に抑え込みます。
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# 歩幅が安全な状態におさまったことを保証してから重みを更新
optimizer.step()

3. ウォームアップ (Warm-up) による学習の「立ち上がり」安定化

学習の「最初期の数エポック」は、最もLossがNaNになりやすい超・危険地帯です。
なぜなら、学習開始直後のモデルの重みパラメータは「ランダムに初期化」されているだけで、言わば目隠しをした状態だからです。この目隠し状態でいきなり大きな学習率(歩幅)で最適化空間を走り回ると、容易に崖から落ちて爆発します。

これを防ぐため、学習の最初の数千ステップだけは学習率を 0 から徐々に(線形に)目標値までジワジワと上げていく「Learning Rate Warm-up」を採用します。
「まずはストレッチから始めて、徐々にスピードを上げていく」というこのアプローチは、HuggingFaceのTransformerライブラリなどでもデフォルトで採用されており、学習序盤の崩壊を劇的に防いでくれます。

まとめ:オプティマイザのブラックボックスを開ける

「LossがNaNになる」という絶望的な現象は、裏を返せば「ニューラルネットワークの数学的限界」にモデルが直面しているという貴重な証拠です。

「とりあえずエラーが出たから学習率を0.001から0.0001に下げてみよう」といった勘に頼る調整(お祈りチューニング)から卒業し、Schedulerの挙動やOptimizer(AdamWなど)の数学的特性をしっかりと理解することが、初心者から一歩上のAIエンジニアへとステップアップする鍵となります。
Deep Learningの裏側で動いている数学や勾配の仕組みをさらに深く知りたい方は、ぜひ専門書での学習をおすすめします!

コメント

タイトルとURLをコピーしました