Skip to content

pytorch_basics.md

Latest commit

 

History

History
186 lines (137 loc) · 18.2 KB

pytorch_basics.md

File metadata and controls

186 lines (137 loc) · 18.2 KB
Raw

PyTorchの基本：Parameter、Module、Linear層の仕組み

PyTorchを用いてLLMなどのニューラルネットワークを定義・学習させる上で、ベースとなる重要なクラスの仕様や数学的背景を整理したドキュメントです。

1. PyTorchの重要概念：Parameter と requires_grad

LLMの重み行列（パラメータ）を定義する際、torch.nn.Parameter や requires_grad というPyTorch固有の設定が頻繁に登場します。

① requires_grad (勾配が必要かどうかのフラグ)

  • 意味: そのテンソルの値を**「学習（更新）の対象にするかどうか」**をPyTorchに指示する設定です。
  • 挙動:
    • requires_grad=True (デフォルト): PyTorchは、このテンソルを使用して行われたすべての計算（数式）を裏で記録し、逆伝播（バックプロパゲーション）の際に自動的に勾配（傾き・更新に必要な情報）を計算します。
    • requires_grad=False: 計算プロセスが記録されず、勾配も計算されません。
  • 使い分け: モデルの学習時は当然 True に設定しますが、「まず順伝播の数式の動きを手動で追いたいだけ（デバッグ中）」のときや、学習済みの重みを固定してファインチューニング（フリーズ）したいときは False を指定します。

② torch.nn.Parameter (重みパラメータクラス)

  • 意味: 単なるデータ（テンソル）を、**「ニューラルネットワークの『学習可能な重み』としてPyTorchシステムに登録するためのクラス」**です。
  • なぜ必要なのか？（普通のテンソルとの違い）:
    • 普通のテンソル（torch.rand など）をモデルクラスの変数に代入しても、PyTorchの管理システム（オプティマイザ）はそれを「学習させるべき重み」として認識してくれません。
    • torch.nn.Parameter(...) でラップしてモデルに持たせると、自動的にそのモデルのパラメータリスト（model.parameters()）に登録され、オプティマイザ（Optimizer）が自動でその値を見つけて更新できるようになります。

2. 勾配消失 (Vanishing Gradient)

ディープラーニングにおいて、モデルが正常に学習できなくなってしまう最も有名な問題の1つです。

① 直感的な意味

AIが学習するためのヒント（勾配）が、層（ネットワーク）を遡って伝わっていく過程でどんどん小さくなり、最終的にほぼゼロ（0）になって消えてしまう現象です。

② なぜ発生するのか？（伝言ゲームの例え）

モデルの学習（バックプロパゲーション）は、**「出力（答え）で発生した誤差の反省点を、入力（最初）の層に向かって順番に伝えていく伝言ゲーム」**のようなものです。 数学的には、この伝達プロセスは「掛け算の繰り返し（微分の連鎖律）」で表されます。

  • この伝える途中で、「1未満の小さな数値（例: 0.1）」 を掛け算する処理が何十回も重なると、どうなるでしょうか？ $$0.1 \times 0.1 \times 0.1 \times \dots \to 0$$ 数値は急激にゼロに近づき、最終的には完全に消えてしまいます。
  • 結果として、入力に近い深い層の重みパラメータは、「どう調整すれば誤差が減るのかというヒント（勾配）」が届かなくなり、学習が完全に停止してしまいます。

③ アテンション（Softmax）での勾配消失

アテンションスコアがスケーリング（$\sqrt{d_k}$での割り算）されていない場合、ソフトマックスの出力が [0.999, 0.000, 0.001] のように極端に張り付きます（飽和状態）。 この状態のソフトマックスの傾き（微分）はほぼ 0 であるため、**「伝達プロセスで 0 を掛け算する」**ことになり、そこから前の層への勾配の伝達が完全にシャットアウトされ、勾配消失が発生します。

3. PyTorchの基本クラス：nn.Module と forward 関数の作法

PyTorchで自作のニューラルネットワーク（レイヤーやモデル全体）を記述する際、必ず torch.nn.Module を継承します。

① nn.Module とは何か？

  • 意味: ニューラルネットワークの「モジュール（部品/塊）」を表す、**すべてのモデルのベースとなる基本クラス（親クラス）**です。
  • 継承する恩恵（自動化される機能）:
    1. 重み（パラメータ）の自動管理: 内部で定義した nn.Parameter などを自動で追跡し、model.parameters() でオプティマイザに渡すパラメータを一括取得できるようにしてくれます。
    2. デバイスの一括転送: model.to("cuda") と書くだけで、モデル内のすべての重みパラメータを一発でGPUに転送してくれます。
    3. モード切り替え: model.train() (学習モード) / model.eval() (評価モード) を呼ぶだけで、ドロップアウト等の挙動を一括制御できます。
    4. モジュールの階層化: モデルの中に別のモデル（例: アテンションブロック）をネストして、複雑な巨大モデル（GPTなど）を綺麗に組み上げられます。

② forward 関数とは何か？

  • 役割: モデルに入力データ（x）を渡した際の、**「順伝播（じゅんでんば／Forward propagation）のメイン計算処理」**を定義する関数です。
  • 役割の分担:
    • __init__: 必要なパーツ（重みパラメータ、別レイヤー等）を「準備・初期化」して nn.Module に登録する場所。
    • forward: 登録したパーツを使って、入力データをどう加工して出力するか「計算の流れ」を書く場所。

③ 【超重要】呼び出しの作法（直接呼んではいけない）

  • 定義する際は def forward(self, x): と書きますが、使用（呼び出し）するときは model.forward(x) と直接呼び出してはいけません。
  • 正しい呼び出し方: model(x) のように、インスタンス自体を関数のように呼びします。
# [定義]
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x * 2

# [使用]
model = MyModel()
output = model(inputs) # ⭕ 正しい (裏側で自動的にシステムの下準備をしてから forward が呼ばれる)
# output = model.forward(inputs) # ❌ 直接呼ぶのはNG

なぜ直接呼ぶのがNGなのか？

PyTorchは forward を実行する前後で、自動勾配（バックプロパゲーション）用の記録登録や、評価/学習モードの切り替えなど、**数多くのシステム的な下準備（フック処理）**を行っています。 model(x) と呼ぶことで、これらの下準備が正常に行われてから安全に forward が呼び出されます。

④ 【超重要】モジュール設計：個別インスタンス vs 使い回し（共有）のルール

モデルを定義する際、LayerNorm や Linear はわざわざ場所ごとに個別のインスタンス（例: self.norm1 と self.norm2）を作るのに対し、なぜ活性化関数（GELU や ReLU）や Dropout は1つのオブジェクトを共有したり、関数として直接呼び出してもよいのでしょうか？

この使い分けの決定的なルールは、**「そのモジュールが内部に『学習可能なパラメータ（重みやバイアス）』を持っているかどうか」**です。

A. 個別インスタンスの作成が必須なもの（状態を持つモジュール）

  • 代表例: nn.Linear、LayerNorm、MultiHeadAttention
  • 理由: これらは内部に、学習過程で値が最適化されていくパラメータ（nn.Parameter）を保持しています。
    • 例えば、LayerNormはデータを標準化する際の拡大縮小・平行移動を司る scale と shift パラメータを持っています。Attentionの前の正規化（norm1）と、FFNの前の正規化（norm2）では、最適な拡大率や移動量は全く異なるはずです。
    • もしこれらを1つのインスタンスで共有して使い回してしまうと、2つの異なる場所の学習情報がごちゃ混ぜになり、お互いのパラメータ学習を破壊し合ってモデルが壊れてしまいます。
    • したがって、独自の学習パラメータを持つものは、必ず場所ごとに個別のインスタンスとして生成しなければなりません。

B. 1つのインスタンスを使い回して（共有して）よいもの（状態を持たないモジュール）

  • 代表例: nn.GELU、nn.ReLU、nn.Dropout（または F.gelu や F.dropout などの関数呼び出し）
  • 理由: これらは内部に学習するパラメータ（重み）を一切持っていません。
    • GELUは「入ってきた数値をただ数式で曲げて出力するだけ」、ドロップアウトは「入ってきた数値の一部をただランダムに0にするだけ」の、状態を持たない（ステートレスな）ただの計算機です。
    • どこで呼び出しても自身の状態が変化しないため、1つのインスタンスをモデル内の複数の場所で共有して使い回しても、全く問題ありません。

Tip

  • 個別インスタンスが必要なもの ＝ 「個人の日記帳」。場所や人ごとに別々の記録が書き込まれるため、共有することはできません。
  • 使い回していいもの ＝ 「ハサミ」や「電卓」。誰がどこで使っても道具自体の状態は変化しないため、1つを使い回しても問題ありません。

4. PyTorchの Linear層 と バイアス（切片）

ニューラルネットワークで最も基本となる構成要素である torch.nn.Linear（全結合層 / 線形層）と、そこで指定できる bias（バイアス）の役割について整理します。

① torch.nn.Linear とは？

  • 意味: 入力データに対して「線形変換（行列積と足し算）」を一括で行うための組み込みレイヤーです。
  • 計算の数式: $$y = x W^T + b$$
    • $x$: 入力ベクトル（または行列）
    • $W$: 重み（weight）パラメータ
    • $b$: バイアス（bias）パラメータ
  • メリット:
    1. 重み $W$ の初期化、データの計算、バイアスの加算をすべて自動で行ってくれます。
    2. 裏側で「重みの初期化アルゴリズム（Xavier初期化など）」が自動適用されるため、手動で乱数作成してモデルを作るよりも、はるかに学習が安定しやすくなります。

② バイアス（バイアスユニット）の数学的役割

  • 意味: グラフにおける**「切片（y切片）」**の役割をするパラメータです。
  • なぜ必要なのか？（直感的な解説）:
    • もしバイアス $b$ がなく、重み $W$ だけ（$y = x W^T$）だと、入力 $x$ が 0 のとき、出力 $y$ も絶対に 0 になってしまいます。これは、グラフで言えば「必ず原点 (0, 0) を通る直線」しか描けないのと同じです。
    • バイアス $b$ があることで、グラフを上下（あるいは左右）に平行移動させることができます（原点を通らない直線が引ける）。これにより、モデルの表現力（後述）が大きく向上し、より複雑なデータパターンに対応できるようになります。

③ なぜLLMのアテンション（Q, K, V）では bias=False が多いのか？

  • 現代の潮流: 現代の高性能なLLM（GPT-3、LLaMA、Chinchillaなど）のアテンション射影では、bias=False に設定してバイアスをあえて使用しない設計が主流です。
  • 理由:
    1. LLMのように極めて巨大なモデルでは、バイアスがない方が過学習（特定のデータにだけ適応しすぎること）を防ぎやすく、かつ学習が数値的に安定しやすいという研究報告があるためです。
    2. キーやクエリのベクトルにおいて、単語全体の傾向から「特定の方向（原点からのズレ）」をあえて固定で持たせる必要がないためです。

④ なぜ nn.Linear は重み行列を [out_features, in_features] (転置された形) で保持するのか？

プログラミングで実際に計算を行う数式 $y = x W + b$ （行ベクトル前提）の観点から見ると、**[out_features, in_features] という形は、数学的な意味での「転置された形（転置後）」**になります。

プログラミング上の素直な形状は [in_features, out_features] であるため、重みがあらかじめ転置された状態で格納されているということになります。そのため、実際の計算を行う際は、以下のように W.T（再転置して元の形状に戻す）を行ってから行列積を計算します。 $$y = x W^T + b$$

このように、わざわざ「転置した状態で格納し、計算時にもう一度転置して元に戻す」という一見無駄に見える処理を行っている理由は、主に以下の3つにあります。

  1. 歴史的な数学の表記法（縦ベクトル前提の数式）との整合性: 多くの数学・AIの教科書では、ベクトル $x$ は「縦ベクトル（列ベクトル）」として定義され、線形変換は $y = W x + b$ と表記されます。このとき、行列 $W$ の形状は [out_features, in_features] になります。PyTorchはこの「教科書の表記における重み $W$ の形状」をそのままパラメータの保存形状として採用しています。
  2. メモリ上の連続性による計算の高速化 (Row-Major): PyTorchの裏側で動く C++ / CUDA の高速な行列演算カーネル（GEMM）では、メモリが連続して配置されている方向（行優先: Row-Major）をスキャンしながら掛け算を行います。重みを行列の形状 [out_features, in_features] で持っておき、転置フラグを立てて計算（x @ W.T）を行う方が、ハードウェアのキャッシュメモリを最も効率的に利用でき、演算スピードが極限まで高まるように最適化されています。
  3. ブロードキャスト（バッチ処理）の適合性: 入力データが [バッチサイズ, in_features] であるとき、重みパラメータの末尾の次元を in_features に揃えておくことで、自動で次元を補完して一括計算する「ブロードキャスト」や、多次元テンソル（LLMのアテンションで用いる [バッチ, ヘッド, 単語数, 特徴量] など）に対する行列演算の次元整合性が取りやすくなります。

5. nn.Parameter と nn.Linear の本質的な違い

PyTorchにおけるこれら2つの最大の違いは、「パーツ（原材料）か、完成品（コンポーネント）か」 という抽象度の違いにあります。

① nn.Parameter (原材料・ネジ)

  • 単なる配列（テンソル）を「自動勾配の対象にする」ための最小限のラッパーです。
  • これを使う場合、重みやバイアスの定義、初期乱数の設定（torch.randなど）、さらに forward 内での行列計算（x @ W.T）などのすべての設計図と組み立てを手書きする必要があります。

② nn.Linear (完成品・エンジン)

  • 内部で nn.Parameter(weight) と nn.Parameter(bias) の定義を自動で行ってくれます。
  • さらに、順伝播の数式（x @ W.T + b）も内部に組み込まれています。
  • さらに、以下に述べる「初期値の自動最適化」も最初から適用されています。
  • つまり、**「Linearという完成品の中に、自動的に Parameter が組み込まれている」**という関係性です。

6. 重みの初期化 (Weight Initialization) と Xavier/He初期化

ニューラルネットワークの学習は、「初期値（スタート地点の乱数）」から少しずつパラメータを調整していく旅です。このスタート地点の数値が適切でないと、学習が全く進まなくなります。

① なぜ初期値の調整が必要なのか？

  • 初期値が大きすぎる場合: 活性化関数（SigmoidやSoftmaxなど）の出力が極端になり、前述の勾配消失や、計算値が無限大に吹っ飛ぶ勾配爆発が起きて学習が崩壊します。
  • 初期値が小さすぎる場合: 層を通過するごとに信号（値）がどんどん減衰してゼロになり、やはり学習が進まなくなります。

② Xavier初期化 / He（Kaiming）初期化とは？

  • 目的: 「入力データのばらつき（分散）」と「出力データのばらつき」が、層を通過した後も同じスケールに保たれるように、次元数（入力数・出力数）に応じて乱数の広がり幅（標準偏差）を動的に決定する数理的な手法です。
  • Xavier初期化: 主に Tanh や Sigmoid などの活性化関数を使う場合に適しています。
  • He初期化: 主に ReLU（マイナスの値を0にする関数）を使う場合に適しています（ReLUによって情報が半分消えることを数学的に考慮し、少し大きめの初期値を与える設計）。
  • nn.Linear を呼び出すと、裏側でこれらの賢い初期化（PyTorchではデフォルトでHe初期化の変形版）が自動で実行され、学習が最初からスムーズに開始されるような「良い初期値」がパラメータに自動で埋められます。