第13章 MLP層とGELU:トークンごとの情報を加工する

こんにちは。ゆうせいです。

第12章までで、Transformerブロックの前半、すなわち注意機構とその周辺の仕組みが揃いました。第13章では、ブロックの後半にあたるMLP層を扱います。Transformer Explainerの図では、注意機構の後ろに配置された、シンプルな二層の構造として表示されている部分です。注意機構と比べると地味な扱いを受けがちですが、実はGPT-2のパラメータの大半は、こちらに使われています。なぜ次元をいったん4倍に拡大してから元に戻すのでしょうか。活性化関数に、よく知られたReLUではなくGELUが使われるのはなぜでしょうか。第13章では、この二つを中心に解説します。

MLP層とは何か

MLPは、Multi-Layer Perceptron(多層パーセプトロン)の略です。Feed Forward Network(FFN、順伝播型ネットワーク)と呼ばれることもあります。資料によって呼び方が異なりますが、同じものを指しています。

構造は、非常に単純です。二つの線形変換の間に、活性化関数を挟んだだけのものです。

MLP(x) = W_2 \cdot GELU(W_1 x + b_1) + b_2

処理の流れは、次の三段階です。

第一に、768次元の入力を、3072次元に拡大します。

第二に、活性化関数GELUを適用します。

第三に、3072次元を、768次元に戻します。

768 \rightarrow 3072 \rightarrow 768

拡大してから元に戻すという、いったん膨らませる構造になっています。

注意機構との役割の違い

MLP層の役割を理解するには、注意機構との対比が有効です。

注意機構は、トークンどうしの情報をやり取りします。あるトークンが、他のトークンから情報を集める処理です。第9章で確認したとおり、複数のトークンのValueが重み付きで合成されます。

MLP層は、各トークンの内部で完結します。1つのトークンのベクトルを受け取り、変換して、返すだけです。他のトークンの情報は、一切参照しません。

この違いは、実装上も明確です。MLP層は、トークンごとに独立に、同じ重みで適用されます。5トークンの文章であれば、同じMLP層が5回、それぞれ独立に適用される形です。

二つの処理の役割分担

この構造から、Transformerブロックの役割分担が見えてきます。

注意機構が、周囲から必要な情報を集めてきます。

MLP層が、集めてきた情報を含めて、そのトークンの表現を加工します。

たとえば、注意機構によって「it」というトークンに「animal」の情報が取り込まれたとします。しかし、取り込まれた情報は、単純な重み付き和にすぎません。ここから「これは生物を指す代名詞である」といった、より抽象的な表現へと加工するのが、MLP層の役割です。

集める処理と、加工する処理。この二つが交互に繰り返されるのが、Transformerの基本構造です。

なぜ次元を4倍に拡大するのか

768次元を3072次元に拡大し、また768次元に戻す。この拡大と縮小には、どのような意味があるのでしょうか。

理由1:表現できる関数の複雑さが増す

線形変換だけを重ねても、結局は一つの線形変換と等価になってしまいます。

W_2(W_1 x) = (W_2 W_1)x

二つの行列を掛けた結果も、また一つの行列です。したがって、活性化関数がなければ、二層に分ける意味がありません。

活性化関数という非線形な処理を挟むことで、初めて複雑な変換が表現できるようになります。

そして、活性化関数を適用する空間の次元が大きいほど、表現できる関数の複雑さが増します。3072次元の空間で非線形変換を行うことで、768次元のまま処理するより、はるかに豊かな変換が可能になります。

理由2:知識の記憶場所として機能する

近年の研究では、MLP層が、モデルが学習した知識を保持する場所として機能しているという見方が示されています。

第一の線形変換(W_1)が、入力に対して「どの知識に該当するか」を判定する役割を果たし、第二の線形変換(W_2)が、該当する知識の内容を出力する、という解釈です。

この見方では、3072という次元数は、記憶できる知識の項目数に対応することになります。次元を拡大するのは、より多くの知識を保持できるようにするため、という説明になります。

この解釈は、モデルの内部を解析する研究の中で提案されているものであり、確立された結論ではありません。ただし、MLP層のパラメータを編集することで、モデルが持つ特定の知識を書き換えられるという実験結果も報告されており、有力な仮説として注目されています。

4という倍率について

拡大の倍率が4倍である点には、理論的な必然性はありません。最初のTransformer論文で採用された値が、その後のモデルでも慣例的に踏襲されているというのが実情です。

近年のモデルでは、この倍率を変更しているものもあります。また、活性化関数の構造を変えた設計では、異なる比率が使われることもあります。

「4倍が最適だと証明されている」わけではなく、「経験的に妥当な範囲として広く使われている」と理解しておいてください。

パラメータ数を確認する

MLP層のパラメータ数を計算します。

第一の線形変換:768 \times 3072 = 2359296

第二の線形変換:3072 \times 768 = 2359296

合計:4718592

約472万個です。バイアス項を含めると、これよりわずかに多くなります。

第10章で計算した注意機構のパラメータ数は、1ブロックあたり約236万個でした。MLP層は、その約2倍です。

1つのブロック全体では、次のようになります。

構成要素パラメータ数割合
注意機構(QKV生成と出力射影)約236万個約33パーセント
MLP層約472万個約67パーセント
レイヤー正規化約3千個約0.04パーセント

Transformerといえば注意機構が主役として語られますが、パラメータの数で見れば、MLP層のほうが2倍の規模を占めています。

モデル全体では、次のようになります。

4718592 \times 12 = 56623104

約5662万個です。GPT-2(small)の全パラメータ数約1億2400万個のうち、約46パーセントがMLP層に使われている計算になります。

活性化関数GELU

MLP層の中央に置かれる活性化関数が、GELU(Gaussian Error Linear Unit)です。

ReLUとの比較

まず、より基本的な活性化関数であるReLUを確認します。

ReLU(x) = \max(0, x)

入力が正であればそのまま通し、負であれば0にする関数です。

単純で計算が速く、勾配消失が起きにくいという利点があります。長らく、ディープラーニングの標準的な活性化関数として使われてきました。

一方で、問題もあります。入力が0のところで、関数が滑らかにつながっていない点です。

グラフを描くと、0を境に折れ曲がった形になります。この折れ曲がりの点では、微分が定義できません。

また、入力が負の領域では、微分値が完全に0になります。この領域に入ったニューロンは、勾配が伝わらず、学習が止まってしまいます。この現象を、ニューロンの死(Dying ReLU)と呼びます。

GELUの定義

GELUは、この問題を緩和する活性化関数です。

GELU(x) = x \cdot \Phi(x)

\Phi(x) は、標準正規分布の累積分布関数です。入力xが、標準正規分布から取り出した値以下になる確率を表します。

この式が意味するところを、直感的に説明します。

入力が大きな正の値であれば、\Phi(x) は1に近くなります。したがって、出力は入力とほぼ同じになります。

入力が大きな負の値であれば、\Phi(x) は0に近くなります。したがって、出力は0に近くなります。

入力が0付近であれば、\Phi(x) は0.5前後になります。したがって、出力は入力の半分程度になります。

つまり、ReLUのような「通すか、遮断するか」の二択ではなく、「どれくらいの割合で通すか」を、入力の大きさに応じて滑らかに決める関数です。

GELUの利点

一つ目の利点は、関数が滑らかであることです。折れ曲がりがないため、すべての点で微分が定義できます。学習が安定しやすくなります。

二つ目の利点は、負の領域でも完全に0にならないことです。小さいながらも値が残るため、勾配が伝わり、ニューロンの死が起きにくくなります。

三つ目の利点は、実験的に、Transformerにおいて良い性能を示すことが確認されている点です。理論的な必然性というより、経験的な選択という側面が強いものです。

実装上の近似式

正確な累積分布関数の計算は、やや重い処理です。そこで、実装では近似式が使われることがあります。

GELU(x) \approx 0.5x\left(1 + \tanh\left(\sqrt{\frac{2}{\pi}}\left(x + 0.044715x^3\right)\right)\right)

複雑に見えますが、tanh関数と多項式で構成されており、正確な計算より高速です。

GPT-2の実装では、この近似式が使われています。PyTorchのnn.GELUでは、approximate という引数で、正確な計算と近似のどちらを使うか指定できます。

数値で比較する

いくつかの入力について、ReLUとGELUの出力を比較します。

入力ReLUGELU(近似値)
-3.00.000-0.004
-1.00.000-0.159
-0.50.000-0.154
0.00.0000.000
0.50.5000.345
1.01.0000.841
3.03.0002.996

負の領域で、GELUがわずかにマイナスの値を返している点に注目してください。また、入力が0.5のとき、ReLUは0.5をそのまま返すのに対し、GELUは0.345と、やや小さい値を返します。

入力が大きくなると、両者の差はほとんどなくなります。違いが現れるのは、主に0付近の領域です。

Transformer Explainerで観察できること

Transformer Explainerでは、MLP層の処理が、次元の拡大と縮小として表示されます。

次の点を確認してみてください。

一つ目は、中間層の次元数が、入力の4倍になっているという点です。表示されるベクトルの長さが、いったん伸びて、また戻ります。

二つ目は、GELUを通した後、多くの値が0に近い小さな値になっているという点です。負の入力が、ほぼ0に抑えられるためです。

三つ目は、この処理がトークンごとに独立していることです。あるトークンのMLP層の処理に、他のトークンの値は関与しません。

初学者がつまずきやすい点

つまずき1:MLP層が単純だから重要でないと考える

構造が単純であることと、重要でないことは別です。

パラメータ数で見れば、MLP層はモデルの約46パーセントを占めています。注意機構の2倍の規模です。

Transformerの解説では、注意機構に説明の大半が割かれるため、MLP層が軽視されがちです。しかし、モデルが持つ知識の多くは、MLP層に蓄えられていると考えられています。

つまずき2:活性化関数の位置を間違える

活性化関数は、二つの線形変換の間に置かれます。

拡大した後、縮小する前です。最後の線形変換の後には、活性化関数は適用されません。

もし最後にも適用してしまうと、残差接続で足し合わせる値が、常に正の値に偏ってしまいます。

つまずき3:GELUが常にReLUより優れていると考える

GELUがTransformerで広く使われているのは事実ですが、あらゆる場面で優れているわけではありません。

ReLUは計算が軽く、実装も単純です。推論速度が重視される場面では、依然として有力な選択肢です。

また、近年のモデルでは、SwiGLUと呼ばれる、より新しい活性化関数の構造が採用されることもあります。活性化関数の選択は、モデルの設計における一つの選択肢であり、絶対的な正解があるわけではない、と理解してください。

演習

演習1:パラメータ数を計算する

次の条件のモデルについて、1ブロックあたりのMLP層のパラメータ数を計算してください。バイアス項は無視して構いません。

埋め込み次元が4096、中間層の拡大倍率が4倍のモデル。

また、この値は、GPT-2(small)のMLP層の何倍になるでしょうか。

演習2:GELUとReLUを比較する

import torch
import torch.nn.functional as F

x = torch.tensor([-3.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 3.0])

relu_out = F.relu(x)
gelu_out = F.gelu(x)
gelu_approx = F.gelu(x, approximate='tanh')

print(f"{'入力':>7} | {'ReLU':>8} | {'GELU':>8} | {'GELU近似':>10}")
print("-" * 42)
for i in range(len(x)):
    print(f"{x[i]:>7.1f} | {relu_out[i]:>8.4f} | {gelu_out[i]:>8.4f} | {gelu_approx[i]:>10.4f}")

# 微分値を比較
print("\n=== 微分値の比較 ===")
x_grad = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0], requires_grad=True)

y_relu = F.relu(x_grad).sum()
y_relu.backward()
grad_relu = x_grad.grad.clone()

x_grad.grad = None
y_gelu = F.gelu(x_grad).sum()
y_gelu.backward()
grad_gelu = x_grad.grad.clone()

print(f"{'入力':>7} | {'ReLUの微分':>12} | {'GELUの微分':>12}")
print("-" * 38)
for i in range(len(x_grad)):
    print(f"{x_grad[i].item():>7.1f} | {grad_relu[i]:>12.4f} | {grad_gelu[i]:>12.4f}")

演習3:MLP層を実装して確認する

import torch
import torch.nn as nn

torch.manual_seed(0)

class MLPLayer(nn.Module):
    def __init__(self, d_model=768, expansion=4):
        super().__init__()
        d_hidden = d_model * expansion
        self.fc1 = nn.Linear(d_model, d_hidden)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(d_hidden, d_model)
    
    def forward(self, x):
        h = self.fc1(x)
        print(f"  拡大後の形: {h.shape}")
        h = self.act(h)
        print(f"  GELU後、0以下の要素の割合: {(h <= 0).float().mean().item():.2%}")
        out = self.fc2(h)
        print(f"  縮小後の形: {out.shape}")
        return out

mlp = MLPLayer()

# パラメータ数を確認
total = sum(p.numel() for p in mlp.parameters())
print(f"MLP層のパラメータ数: {total:,}")

# 5トークンの入力
x = torch.randn(5, 768)
print(f"\n入力の形: {x.shape}")
print("処理の流れ:")
output = mlp(x)

# トークンごとに独立していることを確認
print("\n=== 独立性の確認 ===")
single = mlp.fc2(mlp.act(mlp.fc1(x[0:1])))
print(f"1トークンだけ処理した結果と一致するか: {torch.allclose(output[0], single[0], atol=1e-5)}")

演習の解答例と解説

演習1の解答は、次の通りです。

中間層の次元数は、次のようになります。

4096 \times 4 = 16384

第一の線形変換:4096 \times 16384 = 67108864

第二の線形変換:16384 \times 4096 = 67108864

合計:134217728

約1億3400万個です。

GPT-2(small)のMLP層が約472万個ですから、およそ28.4倍になります。

注目すべきは、この1ブロック分のパラメータ数だけで、GPT-2(small)のモデル全体(約1億2400万個)を上回っている点です。埋め込み次元を増やすと、パラメータ数が2乗で増えることが、この計算から確認できます。

演習2では、GELUが負の領域でわずかにマイナスの値を返すこと、および0付近でReLUより小さい値を返すことが確認できます。

微分値の比較では、ReLUが負の領域で完全に0になるのに対し、GELUは0でない値を持つことが確認できます。この違いが、ニューロンの死を防ぐ効果につながっています。

演習3では、MLP層のパラメータ数が約472万個であること、GELUを通した後に多くの要素が0以下になること、そしてトークンごとに独立して処理されることが確認できます。

講師向けの補足

この章では、パラメータ数の比較を、必ず数値で示してください。

「Transformerといえば注意機構」という先入観を持つ受講者が多いため、実際にはMLP層のほうが2倍のパラメータを持つという事実は、印象に残ります。

活性化関数については、グラフを描いて比較することを推奨します。ReLUの折れ曲がりと、GELUの滑らかさを視覚的に対比させると、理解が早くなります。ホワイトボードに手描きで十分です。

また、MLP層が知識の保持場所として機能するという解釈は、興味を引く話題です。ただし、確立された結論ではない点を明示してください。研究の進展によって理解が変わる可能性がある領域である、と伝えることが、誠実な説明になります。

理解度を確認する問い

第13章の内容を、次の問いで確認してみてください。

注意機構とMLP層の役割の違いを、情報の流れという観点から説明してみてください。

MLP層で、線形変換を二回行う間に活性化関数を挟む必要があるのはなぜでしょうか。

ReLUの持つ「ニューロンの死」という問題は、なぜ起きるのでしょうか。また、GELUではなぜそれが緩和されるのでしょうか。

GPT-2(small)において、MLP層はモデル全体のパラメータの何割程度を占めているでしょうか。

まとめ

MLP層は、二つの線形変換の間に活性化関数を挟んだ、単純な構造の処理です。768次元を3072次元に拡大し、活性化関数を適用してから、768次元に戻します。

注意機構がトークンどうしの情報をやり取りするのに対し、MLP層は各トークンの内部で完結します。集める処理と加工する処理が交互に繰り返されるのが、Transformerの基本構造です。

次元を拡大する理由は、より高次元の空間で非線形変換を行うことで、表現できる関数の複雑さが増すためです。また、MLP層がモデルの知識を保持する場所として機能しているという解釈も提案されています。

パラメータ数は、1ブロックあたり約472万個であり、注意機構の約2倍です。モデル全体では約5662万個、全パラメータの約46パーセントを占めます。

活性化関数には、GELUが使われます。ReLUと異なり、関数が滑らかで微分が常に定義でき、負の領域でも完全に0にならないため、ニューロンの死が起きにくいという利点があります。

次の第14章、この連載の最終章では、出力層を扱います。12個のブロックを通過したベクトルが、どのようにして次のトークンの確率へと変換されるのか、そして温度パラメータが確率分布をどう変えるのかを解説します。第2章から積み上げてきた処理が、ここで最終的な予測へと結実します。

第14章に進む前に、Transformer Explainerで、MLP層の次元が拡大して縮小する様子を確認しておいてください。そして、そのブロックが12回繰り返された先に、出力層があることを、全体図の中で位置づけておくと、最終章の内容が理解しやすくなります。

セイ・コンサルティング・グループでは新人エンジニア研修のアシスタント講師を募集しています。

投稿者プロフィール

山崎講師
山崎講師代表取締役
セイ・コンサルティング・グループ株式会社代表取締役。
岐阜県出身。
海外放浪の末、2000年創業、2004年会社設立。
IT企業向け人材育成研修歴業界歴20年以上。
すべての無駄を省いた費用対効果の高い「筋肉質」な研修を提供します!
この記事に間違い等ありましたらぜひお知らせください。

学生時代は趣味と実益を兼ねてリゾートバイトにいそしむ。長野県白馬村に始まり、志賀高原でのスキーインストラクター、沖縄石垣島、北海道トマム。高じてオーストラリアのゴールドコーストでツアーガイドなど。現在は野菜作りにはまっている。