SwiGLU

SwiGLU（Swishゲート付き線形ユニット）は、Transformerのフィードフォワード層におけるゲート付きの変種です。各トークンベクトルを2回射影し、一方の射影にSiLU（Swish）活性化関数を適用して、2つを要素ごとに乗算し、その積をモデル幅に戻すよう射影します。FFN(x) = (SiLU(xW) ⊙ xV)W₂。Noam Shazeerは2020年の論文GLU Variants Improve Transformerでこれを提案しました。その後、PaLM、Llama、Mistral、Qwenなどの大規模言語モデル（LLM）でフィードフォワードブロックとして広く使われるようになりました。各数値に独立して作用するReLUやGELUとは異なり、SwiGLUには学習可能な重みが含まれるため、それ自体が小さな層を構成します。

SwiGLUの仕組み#

標準的なTransformerのフィードフォワードネットワーク（FFN）では、1つ目の重み行列で各トークンベクトルの次元を拡張し、活性化関数を適用した後、2つ目の行列で元の次元に戻します。SwiGLUは、1つ目の行列と活性化関数をゲート付き線形ユニット（GLU）に置き換えます。GLUは、同じ入力を2つの並列射影に通します。

• ゲート射影：xWにSiLU（x · σ(x)）を適用します。強い負の入力では値がほぼゼロになり、正の入力では入力値に近くなります。
• 値射影：xVは線形のまま、内容を保持します。
• 要素ごとの積：2つを乗算することで、入力に応じてゲートが各隠れ特徴量を大きく、小さく、またはほぼゼロに調整できます。出力射影W₂は、結果をモデル幅に戻します。

ゲーティングの考え方はDauphinらのGLUに由来します。彼らは2016年に、ゲート付き畳み込み言語モデルでsigmoid関数によるゲートを使用しました。Shazeerの論文では、同じ構造で異なるゲート関数を検証し、それぞれの活性化関数にちなんで変種を命名しました。GLUはsigmoid、ReGLUはReLU、GEGLUはGELU、SwiGLUはSwishを使い、双線形層は活性化関数を使いません。T5のコードベースに倣い、これらの層ではバイアス項を省略します。

LLMでSwiGLUがReLUとGELUに取って代わった理由#

ゲート付きフィードフォワード層は、2つではなく3つの重み行列を使います。そのためShazeerは、パラメーター数と計算量を同等に保つため、隠れ層の幅を2/3に縮小しました。T5-baseモデルでは、3,072ユニットから2,048ユニットに減らしています。同じ規模で比較すると、SwiGLUは524,288回の学習ステップ後にホールドアウトデータで対数パープレキシティ1.636を達成しました。ReLUは1.677、GELUは1.679でした。また、GLUEとSuperGLUEの言語理解ベンチマークでも平均スコアが高くなりました。論文では改善の理論的な説明は示されておらず、結論ではこれらのアーキテクチャの成功を「他のあらゆることと同じく、神の慈悲によるもの」としています。

この成果は大規模モデルにも引き継がれました。GoogleのPaLMは2022年にSwiGLUを採用し、MetaのLLaMAは2023年にReLUをSwiGLUに置き換えました。隠れ層の幅には2/3 × 4dを使い、256の倍数に切り上げています。その結果、次元数4,096の7Bモデルでは隠れユニット数が11,008になります。Llama 2、Llama 3、Mistral、Qwenもこの設計を採用し、Mixtralなどの混合エキスパート（MoE）モデルでは、各エキスパートがSwiGLUフィードフォワードブロックになっています。

PyTorchでのSwiGLUの実装#

以下のPyTorchのコードでは、次元数512のモデル向けに標準的なReLUフィードフォワード層とSwiGLU層を構築し、サイズを比較します。

import torch
import torch.nn.functional as F
from torch import nn

d_model = 512
x = torch.randn(4, d_model)  # four token vectors

# Standard transformer FFN: two matrices around a 4x wider ReLU hidden layer
relu_ffn = nn.Sequential(nn.Linear(d_model, 2048, bias=False), nn.ReLU(), nn.Linear(2048, d_model, bias=False))

# SwiGLU FFN: three matrices, hidden width cut to about 2/3 of 2048
w_gate = nn.Linear(d_model, 1365, bias=False)
w_value = nn.Linear(d_model, 1365, bias=False)
w_out = nn.Linear(1365, d_model, bias=False)
y = w_out(F.silu(w_gate(x)) * w_value(x))  # (SiLU(xW) * xV) W2

relu_params = sum(p.numel() for p in relu_ffn.parameters())
swiglu_params = sum(m.weight.numel() for m in (w_gate, w_value, w_out))
print(f"ReLU FFN:   output {tuple(relu_ffn(x).shape)}, {relu_params:,} parameters")
print(f"SwiGLU FFN: output {tuple(y.shape)}, {swiglu_params:,} parameters")

ReLU FFN: output (4, 512), 2,097,152 parametersとSwiGLU FFN: output (4, 512), 2,096,640 parametersが出力されます。隠れユニット数1,365は2,048のおよそ2/3であり、SwiGLUの3つの行列の重み数はReLUの2つの行列とほぼ同じです。そのため、品質の差があれば、ゲーティングに起因します。多くの実装ではゲート行列と値行列を1つの射影に融合し、その出力を半分に分割します。これにより、小さな行列積を2回行う代わりに、大きな行列積を1回実行できます。

SwiGLUと関連用語の比較#

• SiLU（Swish）：ゲート内で使う要素ごとの関数です。SwiGLUは、3つの射影から成るゲート付き層全体を指します。
• GLUとGEGLU：sigmoidまたはGELUのゲートを使う類似の変種です。Shazeerの評価では、GEGLUは対数パープレキシティ1.633でSwiGLUと同等の結果を示し、GoogleのGemmaモデルで採用されています。
• アテンション：アテンションもデータに応じた重みを生成しますが、トークン間で情報を混合します。SwiGLUは各トークン内の特徴量を個別にゲーティングします。
• ビジョンモデル：ゲート付きフィードフォワード層は、ビジョンモデルやマルチモーダルモデルにも使われます。YOLOEは、Ultralyticsパッケージで動作する、清華大学の研究者が開発したオープンボキャブラリー検出器です。クラスプロンプトのテキスト埋め込みを、残差SwiGLUブロックで精緻化します。このブロックはset_classes()の呼び出しごとに1回実行され、エクスポート時に取り除かれるため、フレームごとの追加コストはありません。