Transformer

Transformerは、自己注意を通じて各トークンが他のすべてのトークンから情報を取り込めるようにし、単語や画像パッチなどのトークンの系列を処理するニューラルネットワークアーキテクチャです。Googleの研究者は2017年、機械翻訳を目的とした論文Attention Is All You Needでこの手法を発表しました。それ以前の再帰型ニューラルネットワーク（RNN）は文を単語ごとに読み込むため、位置間の並列学習ができませんでした。Transformerはすべての位置を一度に処理し、ベースモデルは8基のGPUで12時間学習した後、英独翻訳で最高水準の性能を達成しました。現在、Transformerは大規模言語モデル（LLM）、Vision Transformer、音声認識、多くの物体検出モデルや画像生成モデルの基盤になっています。

Transformerの仕組み#

入力トークンはまず埋め込みに変換されます。埋め込みは固定幅のベクトルで、元のモデルでは幅が512です。次にTransformerは、同一構造の層を積み重ねて適用します。各層には2つのサブレイヤーがあります。

• マルチヘッド自己注意： 各トークンからクエリ、キー、値のベクトルを生成します。注意機構で説明されるように、トークンの出力は、クエリと各キーの適合度に基づいて重み付けしたすべての値ベクトルの加重平均です。元のモデルでは8つのヘッドが並列に動作するため、異なる関係をそれぞれ追跡できます。
• フィードフォワードネットワーク： 各トークンに個別に適用する2層ネットワークで、注意機構が集めた情報を変換します。

各サブレイヤーを残差接続で囲み、レイヤー正規化によって活性化を安定させます。元のモデルでは、入力文を読み取るエンコーダーにこの層を6層、翻訳文をトークンごとに生成するデコーダーに6層積み重ねていました。各デコーダー層には、未来のトークンを隠すマスク付き自己注意と、エンコーダーの出力を読み取るクロスアテンションも追加されます。

注意機構は内容だけを比較するため、Transformer層自体には順序の概念がありません。入力トークンを並べ替えると、出力も同じように並べ替わるだけです。そのため、元の論文の固定された正弦・余弦パターン、BERTやViTの学習可能な位置ベクトル、または現在の多くのLLMで使われる回転位置埋め込み（RoPE）によって、トークン埋め込みに位置情報を追加します。次のPyTorchコードでは、エンコーダー層を1つ使ってその効果を示します。

import torch
from torch import nn

# Transformerエンコーダー層1つと、5つのトークン埋め込みからなる系列
torch.manual_seed(0)
layer = nn.TransformerEncoderLayer(d_model=32, nhead=4, batch_first=True).eval()
tokens = torch.randn(1, 5, 32)
reverse = torch.arange(4, -1, -1)

# 自己注意だけでは順序を認識できません。入力を逆順にすると、出力も逆順になります
print(torch.allclose(layer(tokens[:, reverse])[:, reverse], layer(tokens), atol=1e-5))

# 各位置に学習可能なベクトルを1つ追加すると、各出力はトークンの位置に応じて変わります
positions = nn.Embedding(5, 32).weight
reordered = layer(tokens[:, reverse] + positions)[:, reverse]
print(torch.allclose(reordered, layer(tokens + positions), atol=1e-5))

True、続いてFalseを出力します。位置ベクトルを追加するまでは、トークンを逆順にすると出力も逆順になるだけですが、位置ベクトルを追加するとトークンの順序が各トークンの出力を変化させます。

エンコーダー、デコーダー、エンコーダー・デコーダーモデル#

元の層構成を使用する方法は3つあり、それぞれ異なる役割を果たします。

• エンコーダーのみ： 各トークンは入力全体に双方向から注意を向けるため、理解タスクに適しています。テキスト用のBERTと画像用のViTは、エンコーダーのみのモデルです。
• デコーダーのみ： 各トークンはそれより前のトークンだけに注意を向けるため、モデルはトークンを1つずつ生成します。GPTモデルや現在の多くのLLMがこの設計を採用しています。
• エンコーダー・デコーダー： エンコーダーが1つの系列を読み取り、デコーダーがその系列を参照しながら別の系列を生成します。翻訳モデル、T5、音声認識モデルWhisperなどがこの方式を採用しています。

コンピュータービジョンにおけるTransformer#

Vision Transformer（ViT）は、画像を16×16ピクセルのパッチ系列として扱い、変更を加えていないTransformerエンコーダーを適用します。Swin Transformerは注意範囲を局所ウィンドウに制限し、高解像度画像にも対応します。DETRは検出を集合予測として捉え直しました。Transformerデコーダーが、学習済みの固定数のオブジェクトクエリからクエリごとに1つの物体のボックスを生成するため、非最大値抑制（NMS）と手設計のアンカーが不要になります。Ultralyticsパッケージは、BaiduのRT-DETRを実行します。これはこの設計のリアルタイム版で、ハイブリッドエンコーダーはストライド32の特徴マップ（640×640画像で400トークン）にのみTransformerエンコーダー層を適用し、6層のデコーダーは300個のオブジェクトクエリを精緻化してNMSなしで最終検出結果を生成します。Transformerは画像も生成します。拡散Transformerは近年、テキスト画像生成モデルや動画モデルでU-Netに取って代わり、視覚言語モデルは画像トークンをデコーダーのみのLLMに入力します。

Transformerと関連アーキテクチャの比較#

• 畳み込みニューラルネットワーク（CNN）： 畳み込みは、すべての位置で同じフィルターを使って小さな近傍範囲を捉えます。この組み込みの前提は、データセットが小さい場合に役立ちます。注意機構は1層内で任意の2つの位置を結び付けられますが、空間構造をデータから学習する必要があるため、TransformerがCNNを上回るには、通常、より多くの学習データまたは事前学習が必要です。
• RNNとLSTM： 固定サイズの状態をステップごとに引き継ぐため、トークンあたりのメモリ使用量は一定ですが、学習は逐次的です。Transformerは過去のすべてのトークンを保持し、並列に学習します。
• 状態空間モデル： Mambaなどのアーキテクチャは、系列長に対して線形にコストが増加する再帰状態を使って系列を処理します。ハイブリッドモデルでは、状態空間モデルと注意層を交互に配置します。

自己注意はすべてのトークンを他のすべてのトークンと比較するため、コストは系列長の2乗に比例して増加します。4,096トークンの系列では、各層の各ヘッドに約1,680万個のクエリ・キーの組み合わせがあり、長さを8,192トークンに倍増すると約6,700万個に4倍増します。FlashAttentionはこの計算のメモリ転送量を削減し、KVキャッシュは過去のトークンのキーと値を保存して、デコーダーが生成時にそれらを一度だけ計算するようにします。