AI・機械学習
上級

重み共有(Weight Tying)(オモミキョウユウ ウェイトタイイング)

LLMの入力埋め込み層(Embedding Layer)と出力層(Language Model Head)の重み行列を同一パラメータとして共有する手法。パラメータ数を削減しつつ、埋め込み表現と出力予測の一貫性を高める効果がある。GPT-2、T5、Llama 3など多くのモデルで採用されている。

0 回閲覧
0 いいね

重み共有(Weight Tying)とは、大規模言語モデルにおいて入力トークン埋め込み層と最終出力(次トークン予測)層の重み行列を物理的に同一のパラメータとして共有するテクニックである。2016年にPress & Wolfの論文「Using the Output Embedding to Improve Language Models」で体系的に分析された。語彙サイズ V、埋め込み次元 d のモデルでは、V × d のパラメータ(数千万〜数億個)を節約でき、かつ入力表現と出力予測の意味的整合性が向上する。GPT-2(124M〜1.5B)、T5(60M〜11B)、ALBERT、Llama 3 405B など、2026年現在の主要モデルの大半がこの手法を採用している。

仕組み

通常のLLMアーキテクチャでは2つの巨大な行列が存在する:

  1. 入力埋め込み行列 E (V × d): トークンID → 埋め込みベクトル
  2. 出力射影行列 W_o (d × V): 隠れ状態 → 語彙上のロジット

Weight Tying ではこの2つの行列を同一パラメータとし、W_o = E^T(転置)として共有する。つまり出力層は入力埋め込み行列の転置をそのまま使用する。

パラメータ削減効果の具体例:

  • GPT-2 (V=50,257, d=768): 50,257 × 768 = 38.6M パラメータを節約
  • Llama 3 (V=128,256, d=8,192): 128,256 × 8,192 = 1.05B パラメータを節約
  • T5-XXL (V=32,100, d=4,096): 32,100 × 4,096 = 131.5M パラメータを節約

採用モデル一覧

モデルパラメータ数語彙サイズWeight Tying節約量
GPT-2124M〜1.5B50,257あり38.6M〜77.2M
BERT110M/340M30,522あり23.4M
T560M〜11B32,100あり131.5M (XXL)
ALBERT12M〜235M30,000あり(層間共有も)大幅削減
Llama 38B/70B/405B128,2568Bのみ1.05B (8B時)
GPT-3/4175B/推定1.8T100,277なし(別パラメータ)-
Mistral 7B7.3B32,000あり131M

初期化との関係

Weight Tyingを使用する場合、共有パラメータの初期化は特に重要になる。入力埋め込みと出力射影の両方の役割を担うため:

  • 入力埋め込みとしては各トークンを均等に分散させたい → 小さめの分散
  • 出力射影としては適切なロジット範囲を生成したい → 層のサイズに依存

実用的には N(0, 0.02) や Xavier 初期化が多く採用され、出力層での追加スケーリングで調整する。Llama 3 8B では埋め込み層を N(0, 1/√d) で初期化し、出力層として使用する際は RMSNorm で正規化してからロジットを計算する設計となっている。

メリットとデメリット

メリット:

  • パラメータ数の削減(特に語彙サイズが大きいモデルで効果的)
  • 入力表現と出力予測の意味的一貫性
  • 過学習の抑制効果
  • メモリ使用量の削減

デメリット:

  • 入力と出力で最適な表現空間が異なる場合に制約となる
  • 非常に大規模なモデル(GPT-3/4クラス)では相対的な節約率が低下
  • 語彙拡張時に両層を同時更新する必要がある

よくある質問(FAQ)

Q1: GPT-3/GPT-4はWeight Tyingを使っていますか? A: GPT-3以降の大規模モデルでは Weight Tying を使用していない。パラメータ数が数千億規模では、語彙埋め込みの割合が相対的に小さく(全体の0.5%程度)、入力と出力で別々の最適化を行うメリットが上回ると判断されている。

Q2: ALBERTのWeight Tyingは通常のものと何が違いますか? A: ALBERTは入力-出力間のWeight Tyingに加え、Transformer層間でもパラメータを共有する「Cross-Layer Parameter Sharing」を実装している。12層のBERTと同等の構造を、層間共有により12Mパラメータまで圧縮した。

Q3: Weight Tyingはファインチューニング時に解除すべきですか? A: 通常は維持したまま学習する。Hugging Face Transformersでは config.tie_word_embeddings = True(デフォルト)で自動管理される。解除する場合は model.tie_weights() を呼ばないよう設定を変更する。

Q4: 語彙を追加した場合、Weight Tyingはどうなりますか? A: 語彙追加時は model.resize_token_embeddings(new_vocab_size) を呼ぶと、入力埋め込みと出力層の両方が自動的にリサイズされる(Hugging Face Transformers)。追加トークンの重みは小さなランダム値で初期化される。

まとめ

  • Weight Tyingは入力埋め込み層と出力層の重み行列を共有する技術
  • V × d のパラメータを節約し、意味的一貫性を向上
  • GPT-2、BERT、T5、Llama 3 8B、Mistral 7B など広く採用
  • GPT-3/4 クラスの超大規模モデルでは非採用が主流
  • Hugging Face Transformers では tie_word_embeddings で制御可能