Zenn
🎉

LLM最適化の基礎知識：バッチサイズ、最適化手法、学習率の深い理解

に公開
バッチサイズ
学習率
中野哲平
llm最適化
tech

大規模言語モデル（LLM）の学習を理解するには、最適化手法の背景知識が不可欠です。本記事では、バッチサイズ、最適化アルゴリズム、学習率調整などの基礎概念から、それらがLLM学習においてどのような役割を果たすかまでを詳しく解説します。

バッチサイズとは何か

基本概念の理解

バッチサイズとは、モデルが一度のパラメータ更新で処理するサンプル数のことです。機械学習において、この概念は学習効率と最終性能に大きな影響を与えます。

バッチ処理の種類

フルバッチ（Batch Gradient Descent）

  • 全てのデータを一度に処理
  • 最も安定した勾配推定
  • メモリ使用量が膨大
  • 大規模データセットでは現実的ではない

ミニバッチ（Mini-batch Gradient Descent）

  • データの一部分（バッチ）を処理
  • 現実的なメモリ使用量
  • 勾配推定にノイズが含まれる
  • 現代の深層学習の標準的アプローチ

確率的勾配降下（Stochastic Gradient Descent, SGD）

  • バッチサイズ1（1サンプルずつ処理）
  • 最も高い勾配ノイズ
  • メモリ効率が最高
  • 局所解から脱出しやすい

LLMにおけるバッチサイズの特殊性

LLMの学習では、バッチサイズは単純なサンプル数以上の意味を持ちます：

シーケンス長との相互作用

  • 実際の処理トークン数 = バッチサイズ × シーケンス長
  • メモリ使用量は両方の積に比例
  • 勾配計算の統計的性質にも影響

Attention機構との関係

  • Self-attentionはシーケンス長に対して二次的に計算量が増加
  • バッチ並列化によりこの計算を効率化
  • しかし、小さなバッチでは並列化効果が限定的

最適化アルゴリズムの進化

SGD：最も基本的な最適化手法

数学的定義

θ(t+1) = θ(t) - η * ∇L(θ(t))
  • θ：モデルパラメータ
  • η：学習率
  • ∇L：損失関数の勾配

特徴と利点

  • 単純性：実装が容易で理解しやすい
  • メモリ効率：追加の状態保存が不要
  • 理論的基盤：数学的解析が豊富
  • 汎化性能：適切に調整されれば高い汎化性能

課題

  • 学習率調整の困難さ：手動調整が必要
  • 勾配消失問題：深いネットワークでの学習困難
  • 収束速度：適応的手法に比べて遅い場合がある

Momentum：SGDの改良

基本アイデア
過去の勾配情報を蓄積し、慣性のような効果を導入：

v(t+1) = β * v(t) + (1-β) * ∇L(θ(t))
θ(t+1) = θ(t) - η * v(t+1)

効果

  • 振動の抑制：勾配の方向が安定
  • 加速効果：一貫した方向への移動が加速
  • 局所解からの脱出：慣性により浅い谷を越える

Adam：適応的学習率の代表格

革新的な特徴
Adamは1次モーメント（勾配の平均）と2次モーメント（勾配の分散）の両方を利用：

m(t+1) = β1 * m(t) + (1-β1) * ∇L(θ(t))     # 1次モーメント
v(t+1) = β2 * v(t) + (1-β2) * ∇L(θ(t))²    # 2次モーメント

パラメータの意味

  • β1：1次モーメントの減衰率（通常0.9）
  • β2：2次モーメントの減衰率（通常0.999）
  • ε：数値安定性のための小さな値

β2パラメータの深い理解

β2の役割
β2は過去の勾配の二乗の影響度を制御します。これは実質的に「勾配の分散」の推定に関わります。

時間窓の概念
β2 = 0.999の場合、過去約1000ステップの情報が半減するまでの時間となります。この「記憶の長さ」がバッチサイズと相互作用するのが重要なポイントです。

バッチサイズとの相互作用

  • 大きなバッチ：一度に多くのサンプルを見るため、勾配推定が安定
  • 小さなバッチ：勾配にノイズが多く、適応的調整がより重要

学習率調整の理論と実践

学習率の基本原理

最適化の観点
学習率は最適化の歩幅を決定します。大きすぎると発散し、小さすぎると収束が遅くなります。

バッチサイズとの関係性
従来理論では、バッチサイズと学習率の関係は以下のように考えられていました：

Linear Scaling Rule

η_new = η_base * (batch_size_new / batch_size_base)

Square Root Scaling Rule

η_new = η_base * sqrt(batch_size_new / batch_size_base)

なぜ平方根スケーリングが使われてきたか

統計学的根拠

  • バッチサイズBの勾配推定の分散は1/Bに比例
  • 勾配の標準偏差は1/√Bに比例
  • 同じ「信号対雑音比」を維持するため√Bスケーリング

理論的限界
しかし、この理論は以下を仮定していました：

  • 勾配ノイズが完全にランダム
  • 最適化の風景が単純な凸関数
  • バッチ内のサンプルが完全に独立

実際のLLM学習では、これらの仮定が成り立たない場合が多く、新しいスケーリング法則が必要でした。

勾配蓄積（Gradient Accumulation）

基本概念

技術的メカニズム
勾配蓄積は、メモリ制約下で大きなバッチサイズの効果を模擬する手法です：

  1. 小さなバッチで順伝播・逆伝播を実行
  2. 勾配を蓄積（累積加算）
  3. 設定回数後にパラメータ更新
  4. 蓄積した勾配をクリア

数学的表現

# 通常のミニバッチ
θ(t+1) = θ(t) - η * (1/B) * Σ∇L(x_i, θ(t))

# 勾配蓄積
for i in range(accumulation_steps):
    grad_accumulate += ∇L(mini_batch_i, θ(t))
θ(t+1) = θ(t) - η * grad_accumulate / total_samples

利点と制約

メリット

  • メモリ効率：GPU使用量を制御可能
  • 大バッチ効果：統計的に安定した勾配
  • 実装の柔軟性：ハードウェア制約に対応

潜在的問題

  • 計算オーバーヘッド：追加の同期コスト
  • バッチ正規化への影響：統計推定が不正確になる可能性
  • 学習ダイナミクスの変化：真の大バッチとは微妙に異なる

メモリ効率最適化手法

AdaFactor

設計思想
Adamの機能を維持しながらメモリ使用量を削減：

技術的特徴

  • 因数分解された2次モーメント：行列の低ランク近似利用
  • 相対的ステップサイズ：パラメータの規模に適応
  • 自動学習率調整：手動調整の必要性を削減

メモリ削減効果

Adam: O(2P) メモリ（Pはパラメータ数）
AdaFactor: O(√P) メモリ

その他の効率化手法

8bit Adam

  • パラメータを8bitで量子化
  • 大幅なメモリ削減
  • 精度への影響は最小限

分散最適化状態

  • 最適化状態を複数GPUに分散
  • Zero Redundancy Optimizer (ZeRO)
  • メモリ使用量の線形削減

ファインチューニング手法の比較

LoRA（Low-Rank Adaptation）

基本原理
元のパラメータを固定し、低ランク分解による追加パラメータのみを学習：

W_new = W_frozen + A * B
  • A, Bは小さな行列
  • 学習パラメータ数が大幅削減

利点

  • パラメータ効率：元モデルの1%未満のパラメータ
  • メモリ効率：少ないGPUメモリで実行可能
  • デプロイ効率：複数のタスクに対応可能

限界

  • 表現能力の制約：低ランク仮定による制限
  • タスクによる効果の差：すべてのタスクに最適ではない
  • ハイパーパラメータ調整：ランク数の選択が重要

完全ファインチューニング

従来の認識

  • メモリ使用量が膨大
  • 計算コストが高い
  • しかし最高の性能を実現

新しい視点
小バッチサイズ + メモリ効率最適化により、完全ファインチューニングの実用性が向上。

分散学習における通信ボトルネック

通信パターンの理解

データ並列学習

  • 各ノードが異なるデータで学習
  • 勾配の同期が必要
  • 通信量は勾配サイズに比例

モデル並列学習

  • モデル自体を分割
  • アクティベーションの転送が必要
  • より複雑な通信パターン

バッチサイズと通信効率

大きなバッチサイズ

  • 計算量が多く、通信頻度が低い
  • 計算対通信比が高い
  • 通信効率が良い

小さなバッチサイズ

  • 計算量が少なく、通信頻度が高い
  • 通信がボトルネックになりやすい
  • ネットワーク帯域幅の制約を受けやすい

通信最適化手法

勾配圧縮

  • 勾配の量子化や スparse化
  • 通信量削減、精度とのトレードオフ

非同期更新

  • 同期待ちを削減
  • 収束特性に影響

階層的通信

  • リング型、ツリー型通信
  • ノード数に対するスケーラビリティ向上

ハードウェアとの相互作用

GPU アーキテクチャの理解

メモリ階層

  • Global Memory：容量大、レイテンシ高
  • Shared Memory：容量小、レイテンシ低
  • レジスタ：最高速、容量極小

並列処理特性

  • SIMD（Single Instruction, Multiple Data）
  • バッチ処理による効率化
  • メモリアクセスパターンの最適化

バッチサイズとハードウェア効率

小さなバッチサイズ

  • GPU使用率の低下
  • メモリアクセス効率の悪化
  • しかし、絶対的なメモリ使用量は削減

適切なバッチサイズ

  • ハードウェアの並列性を最大活用
  • メモリ帯域幅を効率利用
  • 計算ユニットの稼働率最大化

実践的な設計指針

バッチサイズ決定のフレームワーク

ステップ1：メモリ制約の確認

  • 利用可能なGPUメモリ
  • モデルサイズとアクティベーションメモリ
  • 最適化状態のメモリ要求

ステップ2：計算効率の評価

  • GPU使用率の測定
  • メモリアクセス効率
  • 計算対通信比

ステップ3：学習安定性の検証

  • 勾配ノイズレベル
  • 収束特性
  • 最終性能への影響

最適化手法選択の指針

シンプルさ重視

  • SGD + Momentum
  • 理解しやすく、調整が容易
  • 汎化性能が良い場合が多い

効率性重視

  • Adam系列（Adam、AdamW）
  • 自動的な学習率調整
  • 収束が早い

メモリ効率重視

  • AdaFactor、8bit Adam
  • 大規模モデルに適用可能
  • 性能への影響を最小化

まとめ：現代LLM最適化の全体像

現代のLLM最適化は、複数の要素が複雑に相互作用する多次元的な問題です。バッチサイズ、最適化アルゴリズム、学習率調整、メモリ効率、ハードウェア特性、分散学習のすべてが連携して最終的な性能を決定します。

重要な洞察

  1. 一律の解はない：タスク、モデル、環境に応じた調整が必要
  2. トレードオフの理解：メモリ、計算、通信、性能のバランス
  3. 実験的検証：理論だけでなく実際の検証が不可欠
  4. 継続的最適化：技術進歩に応じた手法の更新

これらの背景知識を理解することで、LLMの学習における新しい発見や手法改善をより深く理解し、実践に活用できるようになります。

Discussion