Skip to content
Published on

Attention Is All You Need - Transformer論文の完全分析

Authors
  • Name
    Twitter

1. 論文概要

「Attention Is All You Need」は2017年のNeurIPSで発表された論文で、Google BrainおよびGoogle Research所属のAshish Vaswani、Noam Shazeer、Niki Parmar、Jakob Uszkoreit、Llion Jones、Aidan Gomez、Lukasz Kaiser、Illia Polosukhinが共同執筆した。この論文は、従来のRecurrenceとConvolutionを完全に排除し、AttentionメカニズムのみでSequence-to-Sequenceモデルを構成できることを示した、まさにディープラーニング史の転換点となった研究である。

論文が提案したTransformerアーキテクチャは、WMT 2014 English-to-German翻訳タスクで28.4 BLEU、English-to-Frenchで41.8 BLEUを達成し、既存のすべてのモデルを凌駕した。さらに重要なのは、このアーキテクチャがその後BERT、GPT、T5、ViTなど、現代AIのほぼすべての主要モデルの基盤となったという点である。


2. 論文の背景と動機:RNN/LSTMの限界

2.1 Sequential Processingのボトルネック

Transformer以前、Sequence Modelingの標準はRNN(Recurrent Neural Network)とその変形であるLSTM(Long Short-Term Memory)、GRU(Gated Recurrent Unit)であった。これらは系列をt=1,2,...,nt = 1, 2, ..., nの順序で処理しながら、hidden state hth_tを更新する構造を持つ。

ht=f(ht1,xt)h_t = f(h_{t-1}, x_t)

この逐次的な特性は、2つの根本的な問題を引き起こした。

**第一に、並列化が不可能であった。**各時点の計算が前の時点の結果に依存するため、GPUの並列処理能力を効果的に活用できなかった。系列長が長くなるほど、学習時間が線形に増加した。

**第二に、Long-range Dependency問題。**理論的にはLSTMが長期依存性を学習できるとされていたが、実際には系列が長くなるにつれて、離れたトークン間の関係を捉えることが困難であった。すべての過去情報をhidden stateという固定サイズのベクトルに圧縮しなければならないためである。

2.2 Attentionの登場とその限界

Bahdanau et al.(2014)が提案したAttentionメカニズムは、DecoderがEncoderのすべてのhidden stateに直接アクセスできるようにすることで、Long-range Dependency問題を大幅に緩和した。しかし依然としてRNNの上にAttentionを追加する形態であったため、Sequential Processingのボトルネックはそのままであった。

論文の核心的な問いはまさにこれであった:「Recurrenceなしに、Attentionだけで十分か?」

答えはYesであり、その成果がTransformerである。


3. Self-Attentionメカニズム

3.1 コアコンセプト:Query、Key、Value

Self-Attentionの核心的なアイデアは、系列内の各トークンが他のすべてのトークンとの関係を直接計算するということである。このために、各入力ベクトルを3つの役割に変換する。

  • Query (Q):「私はどのような情報を探しているか?」
  • Key (K):「私が提供できる情報の識別子は何か?」
  • Value (V):「私が実際に伝える情報は何か?」

入力系列 XRn×dmodelX \in \mathbb{R}^{n \times d_{model}}に対して、学習可能な重み行列を通じてQ、K、Vを生成する。

Q=XWQ,K=XWK,V=XWVQ = XW^Q, \quad K = XW^K, \quad V = XW^V

ここでWQ,WKRdmodel×dkW^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}WVRdmodel×dvW^V \in \mathbb{R}^{d_{model} \times d_v}である。

3.2 直感的な理解

情報検索システムに例えると理解しやすい。図書館で本を探すとき、Queryは「ディープラーニング入門書」という検索ワード、Keyは各本のタイトルやタグ、Valueは本の実際の内容である。QueryとKeyの類似度が高い本のValueをより多く取得することが、Self-Attentionの本質である。

Self-AttentionがRNNと決定的に異なる点は、系列内の任意の2つのトークン間のパス長(path length)が常にO(1)O(1)であるということである。RNNはO(n)O(n)、CNNはO(logkn)O(\log_k n)(dilated)またはO(n/k)O(n/k)(通常)である。この短いパス長のおかげで、Long-range Dependencyを効果的に学習できる。


4. Scaled Dot-Product Attention

4.1 数式

論文で提案されたAttention関数の正確な数式は以下の通りである。

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

この数式をステップごとに分解してみよう。

Step 1: 類似度計算(QKTQK^T

QueryとKeyのDot Productを計算する。結果はn×nn \times nサイズのAttention Score行列である。各要素(i,j)(i, j)は、ii番目のトークンのQueryとjj番目のトークンのKeyの間の類似度を表す。

Step 2: スケーリング(1dk\frac{1}{\sqrt{d_k}}

dkd_kが大きいほどDot Product値の分散が大きくなり、Softmaxの勾配が非常に小さくなる問題が発生する。具体的には、qqkkの各成分が平均0、分散1の独立確率変数であれば、qkq \cdot kの分散はdkd_kとなる。dk\sqrt{d_k}で割ることで分散が1に正規化され、Softmaxが安定的に動作する。

論文ではこのスケーリングの重要性を実験的にも確認しており、dkd_kが小さい場合はAdditive AttentionとDot-Product Attentionの性能が同程度であったが、dkd_kが大きい場合はスケーリングなしのDot-Product Attentionの性能が大幅に低下した。

Step 3: Softmax

スケーリングされたスコアにSoftmaxを適用してAttention Weightを得る。各行の合計が1になるため、これはValueに対する加重平均の重みとして機能する。

Step 4: Valueとの加重和

最終的にAttention WeightとValueを行列積すると、各トークンの出力は、すべてのトークンのValueを関連性に比例して合算したベクトルとなる。

4.2 マスキング

DecoderのSelf-Attentionでは、未来のトークンの情報が現在のトークンに漏洩するのを防ぐ必要がある。このために、Softmaxの前に未来の位置に対応するスコアを-\inftyに設定するMasked Attentionを使用する。

MaskedAttention(Q,K,V)=softmax(QKT+Mdk)V\text{MaskedAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T + M}{\sqrt{d_k}}\right)V

ここでMMはUpper Triangular Matrixで、許可されない位置に-\inftyを、許可される位置に00を持つ。


5. Multi-Head Attention

5.1 単一Attentionの限界

1つのAttention関数のみを使用すると、モデルはたった1つの観点からのみトークン間の関係を把握することになる。例えば、「The cat sat on the mat because it was tired」という文で、「it」が「cat」を指すという構文的関係と、「tired」が「cat」の状態を説明するという意味的関係を同時に捉えることが困難になる。

5.2 Multi-Head Attentionの構造

論文はこの問題を複数のAttentionを並列実行することで解決した。

MultiHead(Q,K,V)=Concat(head1,head2,...,headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, ..., \text{head}_h) W^O

ここで各headは以下のように定義される。

headi=Attention(QWiQ,KWiK,VWiV)\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

各headの重み行列はWiQRdmodel×dkW_i^Q \in \mathbb{R}^{d_{model} \times d_k}WiKRdmodel×dkW_i^K \in \mathbb{R}^{d_{model} \times d_k}WiVRdmodel×dvW_i^V \in \mathbb{R}^{d_{model} \times d_v}であり、最終出力プロジェクションはWORhdv×dmodelW^O \in \mathbb{R}^{hd_v \times d_{model}}である。

5.3 論文の設定

論文ではh=8h = 8個のheadを使用し、dk=dv=dmodel/h=64d_k = d_v = d_{model} / h = 64と設定した。全体の次元をhead数で分割するため、Multi-Head Attentionの総計算コストはSingle-Head Attentionとほぼ同一である。

論文のAblation Studyによると、head数が1個の場合BLEUが0.9ポイント低下し、headが多すぎる場合(例:32個)はdkd_kが小さくなりすぎて、かえって性能が低下した。

5.4 3つの使用パターン

TransformerではMulti-Head Attentionが3箇所で使用される。

  1. Encoder Self-Attention:Encoder内で入力系列の各トークンが他のすべてのトークンを参照する。Q、K、Vはすべて前のEncoderレイヤーの出力から生成される。
  2. Decoder Self-Attention(Masked):Decoder内で、これまでに生成されたトークンのみを参照できるようにマスキングされたAttentionである。
  3. Encoder-Decoder Attention(Cross-Attention):DecoderのQueryがEncoderのKey、Valueを参照する。これは従来のSeq2SeqモデルのAttentionに最も類似した部分である。

6. Positional Encoding

6.1 必要性

Self-Attentionは本質的に順序に無関係(permutation invariant)である。入力トークンの順序を変えても、Attentionの出力は(順序が変わるだけで)同一の値を持つ。自然言語において語順は非常に重要な情報であるため、位置情報を明示的に注入する必要がある。

6.2 Sinusoidal Positional Encoding

論文はサインとコサイン関数を利用したPositional Encodingを提案した。

PE(pos,2i)=sin(pos100002i/dmodel)PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i / d_{model}}}\right) PE(pos,2i+1)=cos(pos100002i/dmodel)PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i / d_{model}}}\right)

ここでposposは系列内の位置、iiは次元インデックスである。このEncodingは入力Embeddingに加算(element-wise addition)されてモデルに渡される。

6.3 なぜSinusoidalなのか?

この関数が選択された理由は明確である。

相対位置の表現:任意の固定オフセットkkに対して、PEpos+kPE_{pos+k}PEposPE_{pos}の線形変換として表現できる。これにより、モデルが相対的な位置関係を容易に学習できるようになる。

[sin(posω+kω)cos(posω+kω)]=[cos(kω)sin(kω)sin(kω)cos(kω)][sin(posω)cos(posω)]\begin{bmatrix} \sin(pos \cdot \omega + k \cdot \omega) \\ \cos(pos \cdot \omega + k \cdot \omega) \end{bmatrix} = \begin{bmatrix} \cos(k\omega) & \sin(k\omega) \\ -\sin(k\omega) & \cos(k\omega) \end{bmatrix} \begin{bmatrix} \sin(pos \cdot \omega) \\ \cos(pos \cdot \omega) \end{bmatrix}

学習不要の汎化:学習データに存在しないより長い系列にも自然に拡張できる。学習可能なPositional Embeddingと比較した際、論文では両方式が「ほぼ同一の結果」を示したと報告しており、汎化可能性からSinusoidal方式を最終的に選択した。

周波数スペクトル:低い次元(iiが小さい)ほど波長が短く細かい位置区別を、高い次元ほど波長が長く広い範囲の位置関係をエンコードする。


7. Encoder-Decoder全体アーキテクチャ

7.1 Encoderの構造

EncoderはN=6N = 6個の同一レイヤーで構成される。各レイヤーは2つのSub-layerを持つ。

  1. Multi-Head Self-Attention
  2. Position-wise Feed-Forward Network

各Sub-layerにはResidual ConnectionとLayer Normalizationが適用される。

output=LayerNorm(x+Sublayer(x))\text{output} = \text{LayerNorm}(x + \text{Sublayer}(x))

7.2 Decoderの構造

DecoderもN=6N = 6個の同一レイヤーで構成されるが、Encoderとは異なり3つのSub-layerを持つ。

  1. Masked Multi-Head Self-Attention:Auto-regressive特性を維持するために未来の位置をマスキングする。
  2. Multi-Head Cross-Attention:Encoderの出力をKey、Valueとして使用する。
  3. Position-wise Feed-Forward Network

7.3 全体フロー

入力系列がEmbedding + Positional Encodingを経てEncoderに入り、6個のレイヤーを通過したEncoder出力がDecoderのCross-Attentionに渡される。Decoderは以前に生成されたトークンを入力として受け取り、次のトークンの確率分布を出力し、系列終了トークンが出るまでこのプロセスを繰り返す。すべてのSub-layerの出力次元はdmodel=512d_{model} = 512に統一される。


8. Feed-Forward Network、Layer Normalization、Residual Connection

8.1 Position-wise Feed-Forward Network(FFN)

各Attention Sub-layerの後にはPosition-wise FFNが配置される。「Position-wise」とは、各位置(トークン)に対して独立に、同一の重みを共有して適用されることを意味する。

FFN(x)=max(0,xW1+b1)W2+b2\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2

これは2つのLinear Transformationの間にReLU活性化関数を挟んだ構造である。入力と出力の次元はdmodel=512d_{model} = 512、内部次元はdff=2048d_{ff} = 2048である。つまり4倍に拡張してから元のサイズに縮小するBottleneck構造である。

このFFNは1x1 Convolution 2回と等価であり、各トークンに対して非線形変換を行い、Attentionが捉えた関係情報をより豊かな表現に変換する役割を果たす。

8.2 Residual Connection

各Sub-layerの入力を出力に加算するSkip Connectionである。

output=x+Sublayer(x)\text{output} = x + \text{Sublayer}(x)

この設計はResNetから借用したもので、深いネットワークで勾配が円滑に流れるようにして学習を安定化させる。Residual Connectionが正しく動作するには、加算される2つのテンソルの次元が同一である必要があるため、すべてのSub-layerとEmbeddingの出力次元がdmodel=512d_{model} = 512に統一されている。

8.3 Layer Normalization

各Sub-layerの出力にLayer Normalizationを適用する。Batch Normalizationとは異なり、Layer Normalizationは1つのサンプル内のすべてのFeatureに対して正規化を行うため、バッチサイズに依存しない。

LayerNorm(x)=γxμσ+ϵ+β\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sigma + \epsilon} + \beta

ここでμ\muσ\sigmaは該当レイヤーのすべての次元に対する平均と標準偏差、γ\gammaβ\betaは学習可能なパラメータである。論文ではPost-Norm(Sublayer出力 + Residualの後にLNを適用)方式を採用した。


9. 学習戦略

9.1 OptimizerとLearning Rate Schedule

論文はAdam Optimizerを使用しつつ、独特なLearning Rateスケジュールを適用した。このスケジュールは後に「Noam Scheduler」として広く知られるようになる。

lr=dmodel0.5min(step0.5,  stepwarmup_steps1.5)lr = d_{model}^{-0.5} \cdot \min(step^{-0.5}, \; step \cdot warmup\_steps^{-1.5})

このスケジュールの核心はWarmupである。最初のwarmup_stepswarmup\_steps(論文では4,000ステップ)の間はLearning Rateを線形的に増加させ、その後はステップ数の逆平方根に比例して減少させる。

Warmupが必要な理由は、学習初期にAdamの2次モーメント推定が不安定であるためである。初期にLearning Rateを低く維持することで、パラメータが急激に変化するのを防ぎ、モーメント推定が安定した後に本格的に学習を行うことができる。

Adam Optimizerのハイパーパラメータはβ1=0.9\beta_1 = 0.9β2=0.98\beta_2 = 0.98ϵ=109\epsilon = 10^{-9}である。β2\beta_2が一般的な0.999より小さい0.98に設定されている点が特徴的で、これはAttention Score分布の急速な変化に適応するためと解釈される。

9.2 正則化

Residual Dropout:各Sub-layerの出力にDropout(rate = 0.1)を適用した後、Residual Connectionを行う。またEncoderとDecoderの両方で、Embedding + Positional Encodingの和にもDropoutを適用する。

Label Smoothingϵls=0.1\epsilon_{ls} = 0.1のLabel Smoothingを適用した。これは正解クラスのターゲット確率を1ではなく1ϵls1 - \epsilon_{ls}に、残りのクラスのターゲット確率をϵls/(K1)\epsilon_{ls} / (K - 1)に設定する技法である。論文ではLabel SmoothingがPerplexityを悪化させるが、AccuracyとBLEU Scoreを向上させると報告している。モデルが過度に確信を持つことを防ぎ、汎化性能を高める効果があるためである。

9.3 学習データとハードウェア

  • WMT 2014 English-German:約450万文対、Byte-Pair Encoding(BPE)で約37,000の共有語彙を使用
  • WMT 2014 English-French:約3,600万文対、32,000のWord-piece語彙を使用
  • バッチ:約25,000のSourceトークン + 25,000のTargetトークンを含む
  • ハードウェア:8台のNVIDIA P100 GPU
  • 学習時間:Baseモデル約12時間(100Kステップ)、Bigモデル約3.5日(300Kステップ)

10. 主要な実験結果

10.1 機械翻訳性能

モデルEN-DE BLEUEN-FR BLEUTraining Cost (FLOPs)
Transformer (Base)27.338.13.3×10183.3 \times 10^{18}
Transformer (Big)28.441.82.3×10192.3 \times 10^{19}
従来SOTA(Ensemble含む)26.3641.29-

Transformer BigモデルはEN-DEで従来最高性能を2 BLEU以上上回り、EN-FRでも新たなSOTAを記録した。さらに驚くべきことは、この性能を既存モデルの学習コストの一部分で達成したという点である。

10.2 モデルサイズの比較

ConfigNNdmodeld_{model}dffd_{ff}hhdkd_kParameters
Base6512204886465M
Big6102440961664213M

10.3 Ablation Studyの主要結果

論文のAblation Studyは各設計決定の重要性を明確に示している。

  • Attention Head数h=1h = 1で0.9 BLEU低下、h=16h = 16h=32h = 32ではdkd_kが小さくなりすぎて性能低下
  • dkd_k(Key次元):削減すると品質低下。Dot-Product Attentionの表現力に直接影響
  • dmodeld_{model}(モデル次元):大きくするほど一貫して性能向上
  • Dropout:なしでは過学習が発生し性能が大幅に低下
  • Positional Encoding:学習可能方式とSinusoidal方式がほぼ同一の性能

10.4 English Constituency Parsing

翻訳以外のタスクでの汎化能力を検証するため、English Constituency Parsing(構文解析)にも適用した。WSJデータのみ使用で91.3 F1、半教師あり学習設定で92.7 F1を達成し、タスク特化モデルと競争力のある性能を示した。これはTransformerが機械翻訳に限定されない汎用系列モデルであることを実証した。


11. 後続研究への影響

Transformerアーキテクチャは、現代AIのほぼすべての主要な発展の基盤となった。

11.1 BERT(2018、Google)

TransformerのEncoder部分のみを使用して双方向(Bidirectional)事前学習を行った。Masked Language Modeling(MLM)とNext Sentence Prediction(NSP)という2つの事前学習課題を通じて、11個のNLPベンチマークでSOTAを達成した。BERTはNLPにおけるTransfer Learningパラダイムを確立した。

11.2 GPTシリーズ(2018年〜、OpenAI)

TransformerのDecoder部分のみを使用してAuto-regressive Language Modelingを行った。GPT-1(117M)→ GPT-2(1.5B)→ GPT-3(175B)とスケールアップしながら、Scaling Lawの威力を実証した。GPT-3はFew-shot Learning能力を示してAIの新たな可能性を切り拓き、その後ChatGPT、GPT-4へと続く大規模言語モデル(LLM)革命の出発点となった。

11.3 その後

  • T5(2019):すべてのNLPタスクをText-to-Text形式に統合、Encoder-Decoder全体構造を使用
  • ViT(2020):TransformerをComputer Visionに適用、画像をパッチに分割して系列として処理
  • DALL-E、Stable Diffusion:画像生成にTransformerを活用
  • AlphaFold 2:タンパク質構造予測にAttentionメカニズムを活用

1本の論文がNLPを超えてComputer Vision、生物学、音楽、ロボティクスまで、ほぼすべてのAI分野を変革したのである。


12. PyTorchコア コード例

12.1 Scaled Dot-Product Attention

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def scaled_dot_product_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    mask: torch.Tensor = None,
    dropout: nn.Dropout = None
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Scaled Dot-Product Attentionの実装。
    Args:
        query: (batch, h, seq_len, d_k)
        key:   (batch, h, seq_len, d_k)
        value: (batch, h, seq_len, d_v)
        mask:  Attention mask (optional)
    Returns:
        output: (batch, h, seq_len, d_v)
        attention_weights: (batch, h, seq_len, seq_len)
    """
    d_k = query.size(-1)

    # Step 1 & 2: QK^T / sqrt(d_k)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # マスキング(Decoder Self-Attentionなど)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Step 3: Softmax
    attention_weights = F.softmax(scores, dim=-1)

    if dropout is not None:
        attention_weights = dropout(attention_weights)

    # Step 4: Weighted sum of values
    output = torch.matmul(attention_weights, value)
    return output, attention_weights

12.2 Multi-Head Attention

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int = 512, h: int = 8, dropout: float = 0.1):
        super().__init__()
        assert d_model % h == 0, "d_model must be divisible by h"

        self.d_model = d_model
        self.h = h
        self.d_k = d_model // h

        # Q、K、V、Outputそれぞれのlinear projection
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: torch.Tensor = None
    ) -> torch.Tensor:
        batch_size = query.size(0)

        # 1) Linear projectionの後、(batch, h, seq_len, d_k)にreshape
        Q = self.W_q(query).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)

        # 2) Scaled Dot-Product Attention(全headを並列実行)
        attn_output, attn_weights = scaled_dot_product_attention(
            Q, K, V, mask=mask, dropout=self.dropout
        )

        # 3) Head結果をConcatenate:(batch, seq_len, d_model)
        attn_output = (
            attn_output.transpose(1, 2)
            .contiguous()
            .view(batch_size, -1, self.d_model)
        )

        # 4) Final linear projection
        return self.W_o(attn_output)

12.3 Positional Encoding

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int = 512, max_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # (max_len, d_model)サイズのPositional Encoding行列を生成
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # (max_len, 1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )  # (d_model/2,)

        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数次元:sin
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数次元:cos

        pe = pe.unsqueeze(0)  # (1, max_len, d_model) - batch次元を追加
        self.register_buffer('pe', pe)  # 学習パラメータではなくbufferとして登録

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch, seq_len, d_model) - Embeddingの出力
        Returns:
            (batch, seq_len, d_model) - Positional Encodingが加算された結果
        """
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

12.4 Transformer Encoder Layer

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model: int = 512, h: int = 8, d_ff: int = 2048, dropout: float = 0.1):
        super().__init__()

        # Sub-layer 1: Multi-Head Self-Attention
        self.self_attn = MultiHeadAttention(d_model, h, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)

        # Sub-layer 2: Position-wise FFN
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        # Sub-layer 1: Self-Attention + Residual + LayerNorm
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))

        # Sub-layer 2: FFN + Residual + LayerNorm
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_output))

        return x

上記のコードでself.self_attn(x, x, x, mask)のQ、K、Vがすべて同一の入力xから生成されるため、「Self」-Attentionと呼ばれる。Cross-Attentionの場合はKとVにEncoderの出力を渡せばよい。


13. まとめ

「Attention Is All You Need」は、単に機械翻訳モデル1つを提案した論文ではない。Recurrenceという長年の慣性を打ち破り、Attentionだけで十分であるという大胆な主張を実証的に証明した論文である。

この論文の核心的な貢献を整理すると以下の通りである。

  1. Recurrenceの除去:並列処理可能なアーキテクチャで学習速度を飛躍的に改善
  2. Self-Attention:系列内のすべてのトークン間の関係をO(1)O(1)のパス長で直接モデリング
  3. Multi-Head Attention:多様な観点から同時に関係を捉える
  4. スケーラビリティ:シンプルかつ拡張可能なアーキテクチャ設計により、数十億〜数兆パラメータのモデルまでスケールアップ可能

この論文が発表された2017年以降、TransformerはNLPを超えてVision、Audio、Biology、Roboticsなど、AIのほぼすべての領域に拡散した。論文のタイトル通り、Attentionはまさに必要なすべて(All You Need)であった。


References