Transformer 内部構造ガイド入力から出力までを図・形状・コードで追う

全体図での位置処理順 6 / 16 番目 / 全 20 ブロック全体図で確認する

共通抽象度: Transformer Block難易度: 入門推定学習時間: 約12分

位置ごとの全結合処理Position-wise Feed-Forward Network

30秒で分かる要約各位置のベクトルを独立にD→Dff→Dと変換する2層の全結合層です。Attentionが位置間で情報を混ぜるのに対し、こちらは位置内で非線形な変換を行います。パラメータの大半を占める部分です。

hidden_states[B, S, D]hidden_states[B, S, D]

Level 1概要役割と処理の流れ

何をする機能か

各位置のベクトルを、いったん広い次元 Dff へ引き伸ばし、非線形関数を通してから、 元の次元 D へ戻します。

D→Dff→DD \rightarrow D_{ff} \rightarrow D

重要なのは「位置ごとに独立」という点です。 位置 1 の keys を変換するときに、位置 4 の table の値はまったく使いません。 同じ重みを、全位置に対して個別に適用します。 バッチ処理のループで、各レコードに同じ関数を適用するのと同じ構造です。

原典の base モデルでは D = 512、Dff = 2048(4 倍)です。

なぜ必要か

Attention は、Value の加重平均を取る処理です。重み付き和は線形な操作なので、 それだけを重ねても表現力が足りません。

Feed-Forward 層は、ReLU などの非線形関数を挟むことで、 「集めてきた情報を組み合わせて、新しい特徴を作る」役割を担います。

Transformer の 1 層は、次の 2 つの役割分担になっています。

  • Attention: 位置間で情報を混ぜる(どこを見るか)
  • Feed-Forward: 位置内で情報を変換する(見たものをどう解釈するか)

全体の中でいつ呼ばれるか

各層の最後のサブ層です。Encoder では Self-Attention → Add & Norm のあと、 Decoder では Cross Attention → Add & Norm のあとに実行され、 その出力はまた 残差接続 + LayerNorm へ渡されます。

Level 2内部処理Tensor形状・内部ステップ・設計上の意味

前段から受け取るもの

hidden_states [B, S, D]。直前の Add & Norm の出力です。

後段へ渡すもの

同じ形の hidden_states [B, S, D]。次の Add & Norm へ渡されます。

内部の3ステップ

  1. 拡張: [B, S, D] に W1W_1 [D, Dff] を掛けて [B, S, Dff] にする
  2. 非線形: ReLU(原典)や GELU、SwiGLU を適用する
  3. 圧縮: W2W_2 [Dff, D] を掛けて [B, S, D] に戻す

入力

[B, S, D] = [1, 8, 512]

内部での形の変化

[1, 8, 2048] へ拡張 → 非線形 → [1, 8, 512] へ圧縮

例文での処理

共通例文

The keys on the table are lost.

概念上のトークン列(S = 8): [The, keys, on, the, table, are, lost, .]

Self-Attention を通った時点で、位置 5 の are のベクトルには 「主語が keys である」といった情報が混ざっています。

  1. The
  2. keys
  3. on
  4. the
  5. table
  6. are
  7. lost
  8. .
位置5の are。Feed-Forward はこの位置のベクトルだけを見て変換します。他の位置の値は参照しません。

Feed-Forward はここで、その混ざった情報を材料に、 「複数形の主語に対応する be 動詞」という、より抽象的な特徴へ変換します。 このとき他の位置を参照しないので、 位置 5 の処理と位置 1 の処理は完全に並列に実行できます。

なお、[B, S, Dff] へ広げる操作は、大きな特徴辞書を引くような働きをすると解釈されることがあります (Dff の各次元が特定のパターンに反応する)。 ただしこれは解釈の一つであり、個々の次元に固定の意味が割り当てられているわけではありません。

EncoderとDecoderでの違い

処理内容は同じです。扱う形が [B, S, D] か [B, T, D] かの違いだけで、 重みは Encoder 側と Decoder 側で別々に学習されます(層ごとにも別です)。

入力と出力

入力

項目名データの意味データ型Tensor形状記号の意味具体例次の渡し先
hidden_states各位置のベクトル列(Attention 通過後)float32 Tensor[B, S, D]B = バッチサイズ(同時に処理する文の数) / S = 入力系列長(Encoder側のトークン数) / D = モデル内部のベクトル次元(d_model)[1, 8, 512]W1 による拡張

出力

項目名データの意味データ型Tensor形状記号の意味具体例次の渡し先
hidden_states位置ごとに非線形変換された特徴float32 Tensor[B, S, D]B = バッチサイズ(同時に処理する文の数) / S = 入力系列長(Encoder側のトークン数) / D = モデル内部のベクトル次元(d_model)[1, 8, 512](形は変わらない)残差接続 + LayerNorm

補助出力

項目名データの意味データ型Tensor形状記号の意味具体例次の渡し先
intermediate中間層の活性(デバッグ・解析用。通常は保持しない)float32 Tensor[B, S, Dff]B = バッチサイズ(同時に処理する文の数) / S = 入力系列長(Encoder側のトークン数) / Dff = Feed-Forward 層の中間次元[1, 8, 2048](本流では W2 を掛けて破棄される)
B
バッチサイズ(同時に処理する文の数)
S
入力系列長(Encoder側のトークン数)
D
モデル内部のベクトル次元(d_model)
Dff
Feed-Forward 層の中間次元

このサイトでは Tensor 形状を PyTorch の `batch_first=True` 相当、つまり `[B, S, D]`(バッチ次元が先頭)で統一しています。実装によっては `[S, B, D]` を採用するものもあります。

学習対象は W1W_1 [D, Dff] と W2W_2 [Dff, D] で、1 層あたり約 2D⋅Dff2 D \cdot D_{ff} パラメータです。 D = 512、Dff = 2048 なら約 210 万で、 同じ層の Attention(約 105 万)の2 倍になります。 Transformer のパラメータの多くは、実は Attention ではなくこの層にあります。

よくある誤解

Level 3数式・Python数式・疑似実装・計算量・実装差

数式

原典(ReLU)の定義:

FFN(x)=max⁡(0,  xW1+b1) W2+b2\mathrm{FFN}(x) = \max(0,\; x W_1 + b_1)\, W_2 + b_2

ここで W1∈RD×DffW_1 \in \mathbb{R}^{D \times D_{ff}}、W2∈RDff×DW_2 \in \mathbb{R}^{D_{ff} \times D} です。 xx は各位置のベクトル ∈RD\in \mathbb{R}^{D} であり、位置ごとに同じ W1,W2W_1, W_2 が適用されます。

近年よく使われる SwiGLU では、ゲート機構を加えます。

SwiGLU(x)=(Swish(xWgate)⊙xWup)Wdown\mathrm{SwiGLU}(x) = \big(\mathrm{Swish}(x W_{gate}) \odot x W_{up}\big) W_{down}

行列が 3 本になるため、同じパラメータ数にそろえる目的で DffD_{ff} を 83D\frac{8}{3}D 程度にすることが多いです。

Pythonで表したインターフェース

import torch
from torch import Tensor, nn
 
 
class FeedForward(nn.Module):
    """[B, S, D] を位置ごとに D -> Dff -> D と変換する。"""
 
    def __init__(self, d_model: int, d_ff: int) -> None:
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)   # [D, Dff]
        self.w2 = nn.Linear(d_ff, d_model)   # [Dff, D]
        self.activation = nn.ReLU()          # 原典は ReLU
 
    def forward(self, hidden_states: Tensor) -> Tensor:
        """
        Args:
            hidden_states: [B, S, D]
 
        Returns:
            [B, S, D]。位置ごとに独立に変換される。
        """
        intermediate = self.activation(self.w1(hidden_states))  # [B, S, Dff]
        return self.w2(intermediate)                            # [B, S, D]
 
 
class SwiGLUFeedForward(nn.Module):
    """近年の LLM で多い SwiGLU 版(参考)。"""
 
    def __init__(self, d_model: int, d_ff: int) -> None:
        super().__init__()
        self.gate = nn.Linear(d_model, d_ff, bias=False)
        self.up = nn.Linear(d_model, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d_model, bias=False)
 
    def forward(self, hidden_states: Tensor) -> Tensor:
        gated = nn.functional.silu(self.gate(hidden_states))    # [B, S, Dff]
        return self.down(gated * self.up(hidden_states))        # [B, S, D]
 
 
def feed_forward(hidden_states: Tensor, layer: FeedForward) -> Tensor:
    """[B, S, D] を位置ごとに D→Dff→D と変換する。"""
    return layer(hidden_states)
 
 
# 最小使用例
B, S, D, D_FF = 1, 8, 512, 2048
layer = FeedForward(d_model=D, d_ff=D_FF)
 
hidden_states = torch.randn(B, S, D)     # [B, S, D]
out = feed_forward(hidden_states, layer)
 
print(out.shape)                          # torch.Size([1, 8, 512]) == [B, S, D]
 
# 位置ごとに独立であることの確認:
# 1位置だけを [1, 1, D] として通しても、まとめて通した結果と一致する
single = feed_forward(hidden_states[:, 5:6, :], layer)
print(torch.allclose(single, out[:, 5:6, :], atol=1e-6))  # True

この例で単純化している点

  • Dropout(原典では活性化のあとに適用)を省略しています。
  • バイアス項は残していますが、近年のモデルでは省くことが多いです。
  • SwiGLU 版は参考として別クラスで示しています。

型と Tensor 形状の見え方を掴むための擬似実装です。そのまま学習・推論に使えるものではありません。

計算量

O(B⋅S⋅D⋅Dff)O(B \cdot S \cdot D \cdot D_{ff}) です。Dff=4DD_{ff} = 4D なら O(8BSD2)O(8 B S D^2) に相当し、 Attention の投影部分(O(4BSD2)O(4 B S D^2))の約 2 倍になります。

系列長 S支配的な処理
S が D より小さいFeed-Forward(SD2S D^2 の項)
S が D と同程度拮抗
S が D より大きいAttention のスコア計算(S2DS^2 D の項)

派生と実装差

  • 活性化関数: 原典は ReLU。BERT や GPT-2 は GELU。Llama などは SwiGLU。 GELU は ReLU を滑らかにしたもので、負の値を完全に 0 にしません。
  • Mixture of Experts (MoE): Feed-Forward を複数用意し、 トークンごとに一部だけを使う方式。パラメータ数を増やしつつ計算量を抑えられます。
  • Dff の比率: 原典は 4 倍。SwiGLU では行列が 3 本になるため、 パラメータ数をそろえる目的で 8/3 倍程度にすることが一般的です。

この記事の位置づけ

前提となる知識

残差接続 + LayerNorm

このページ

位置ごとの全結合処理

次に読むとよい記事

Encoder出力(文脈表現)

関連記事