大規模言語モデル（Large Language Models, LLM）において、「Attentionヘッド」と「Attention層」は、Transformerアーキテクチャの中核をなす重要なコンポーネントです。これらは、モデルが入力データの異なる部分間の関係性を効率的に学習・利用するために設計されています。本解説では、これら二つの概念の定義、機能、相互関係、そしてそれぞれがモデル全体に与える影響について、説明します。
1. Transformerアーキテクチャの概要
Transformerアーキテクチャは、Vaswaniら（2017年）によって提案され、主に自然言語処理（NLP）タスクにおいて卓越した性能を発揮しています。Transformerは主に「エンコーダー」と「デコーダー」から構成されますが、LLMでは通常エンコーダー部分またはデコーダー部分、あるいはその両方がスタックされて使用されます。ここでは、主にエンコーダー側の構造を中心に説明します。
エンコーダーは、複数の「レイヤー（層）」から構成されており、各レイヤーは「マルチヘッド・セルフアテンション（Multi-Head Self-Attention）」機構と「フィードフォワード・ニューラルネットワーク（Feed-Forward Neural Network）」から成ります。各層は、前の層からの出力を受け取り、より高度な抽象化を行うための変換を適用します。
2. Attention層とは
Attention層とは、Transformerアーキテクチャ内で情報の相互関係を捉えるための主要な構成要素です。具体的には、入力シーケンス内の各トークン（単語やサブワードなど）が他のトークンとの関係性を計算し、重要度に応じて重み付けされた情報を集約するプロセスを指します。
Attention層は、通常以下のような構成要素から成り立ちます：
- クエリ（Query）、キー（Key）、バリュー（Value）の生成:
- 入力ベクトルを異なる重み行列 \( W_Q, W_K, W_V \) を用いて線形変換し、クエリ、キー、バリューを生成します。
- これらはそれぞれ、情報の検索・参照・集約に使用されます。
- スコア計算:
- クエリとキーの内積を取り、スケーリング（通常は \(\sqrt{d_k}\) で割る）して、各クエリに対するキーの関連性スコアを計算します。
- ソフトマックス関数を適用して、スコアを確率分布に変換します。
- 重み付けされたバリューの集約:
- ソフトマックスで得られた重みをバリューに乗じ、重み付けされたバリューを合計して出力ベクトルを生成します。
このプロセスにより、Attention層は入力シーケンス内の重要な情報を動的に選択し、文脈に応じた表現を生成します。
3. Attentionヘッドとは
Attentionヘッドは、Attention層内で独立してAttention機構を実行する個々の「サブユニット」を指します。マルチヘッド・アテンション（Multi-Head Attention）は、単一のAttention機構よりも多様な情報を同時に捉えるために、複数のAttentionヘッドを並行して動作させる手法です。
各Attentionヘッドは、独自のクエリ、キー、バリューの重み行列 \( W_Q^{(i)}, W_K^{(i)}, W_V^{(i)} \) を持ち、独立したAttention計算を行います。これにより、各ヘッドは異なる部分や異なる種類の関係性に焦点を当てることが可能となり、モデル全体としてより豊富な表現力を獲得します。
4. Attention層とAttentionヘッドの具体的な違い
Attention層とAttentionヘッドの違いを明確に理解するために、以下のポイントを整理します：
- 階層的な関係性:
- Attention層は、Transformerの各層全体を指し、その中に複数のAttentionヘッドが含まれます。
- Attentionヘッドは、Attention層内の個々のAttention機構を指します。
- 機能的な役割:
- Attention層は、入力情報全体に対して多角的な注意機構を適用し、文脈依存の表現を生成します。これは、複数のAttentionヘッドの出力を統合し、次の層への入力として渡す役割を担います。
- Attentionヘッドは、特定の視点や特定の情報に焦点を当ててAttention計算を実行し、局所的な情報抽出を担当します。
- 計算の並列性:
- Attention層は、全てのAttentionヘッドの計算をまとめ、統合する役割があります。
- Attentionヘッドは、それぞれ独立に計算を行うため、並列処理が可能です。
- パラメータの共有:
- Attention層内の各Attentionヘッドは、独自のパラメータ（重み行列）を持ち、互いにパラメータを共有しません。
- Attentionヘッドは、独自のパラメータセットを持つため、多様な注意機構を実現します。
5. マルチヘッド・アテンションの詳細
マルチヘッド・アテンションは、単一のAttention機構では捉えきれない多様な関係性を同時に学習・利用するために設計されています。以下にその詳細を示します：
- 並行したAttentionヘッド:
- 複数のAttentionヘッドを並行して動作させることで、異なる視点から情報を処理します。例えば、一つのヘッドが文法的な関係性に注目し、別のヘッドが意味的な関連性に注目する、といった具合です。
- 次元の分割:
- 入力ベクトルの次元 \( d_{model} \) を、ヘッド数 \( h \) に応じて \( d_k = d_{model} / h \) に分割します。各ヘッドは \( d_k \) 次元のクエリ、キー、バリューを扱います。
- これにより、各ヘッドが計算負荷を分散し、効率的に並列処理を行うことが可能となります。
- 出力の統合:
- 各ヘッドから得られた出力ベクトルを連結し、再び線形変換（通常は重み行列 \( W_O \) ）を適用して、最終的なAttention層の出力を生成します。
- これにより、各ヘッドの多様な情報が統合され、次の層へと効果的に伝達されます。
6. Attention層とAttentionヘッドの相互関係
Attention層とAttentionヘッドは、階層的かつ相補的な関係にあります。具体的には：
- Attention層は、Transformer内の一つの処理ブロックを形成し、その中で複数のAttentionヘッドが並行して動作します。
- 各Attentionヘッドは、異なる重み行列を用いて独自のAttention計算を行い、その結果をAttention層全体の出力に寄与します。
- 最終的に、Attention層はすべてのヘッドからの情報を統合し、次の層や最終出力に対して豊かな文脈依存の表現を提供します。
このように、Attention層は複数のAttentionヘッドを包括する単位であり、ヘッドはその内部で独立したAttention機構を実行するサブユニットです。
7. なぜマルチヘッド・アテンションが重要か
マルチヘッド・アテンションの導入には、以下のような利点があります：
- 多様な情報の同時抽出:
- 複数のヘッドが異なる部分に焦点を当てることで、単一のAttention機構では捉えきれない多様な情報を同時に抽出できます。
- 表現力の向上:
- 各ヘッドが異なる種類のパターンや関係性を学習することで、モデル全体の表現力が向上します。
- 計算効率の最適化:
- 各ヘッドが小規模な計算を並行して実行するため、全体としての計算効率が向上します。
- 過学習の抑制:
- ヘッド間での情報の分散により、特定のヘッドに依存しすぎることなく、より汎化性能の高いモデルが構築できます。
8. Attention層とAttentionヘッドの具体的な実装例
以下に、典型的なTransformerエンコーダー層におけるAttention層とAttentionヘッドの具体的な実装例を示します。
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_k = d_model // num_heads
self.num_heads = num_heads
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)
def forward(self, Q, K, V, mask=None):
batch_size = Q.size(0)
# Linear変換とヘッドの分割
Q = self.W_Q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
K = self.W_K(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
V = self.W_V(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
# スコア計算
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn = torch.softmax(scores, dim=-1)
# 重み付けされたバリューの集約
output = torch.matmul(attn, V)
# ヘッドの統合
output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
output = self.W_O(output)
return output, attn
class TransformerEncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, dim_feedforward, dropout=0.1):
super(TransformerEncoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads)
self.linear1 = nn.Linear(d_model, dim_feedforward)
self.dropout = nn.Dropout(dropout)
self.linear2 = nn.Linear(dim_feedforward, d_model)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.activation = nn.ReLU()
def forward(self, src, src_mask=None):
# Self-Attention層
attn_output, attn = self.self_attn(src, src, src, src_mask)
src = src + self.dropout1(attn_output)
src = self.norm1(src)
# フィードフォワードネットワーク
ff_output = self.linear2(self.dropout(self.activation(self.linear1(src))))
src = src + self.dropout2(ff_output)
src = self.norm2(src)
return src, attn
# 使用例
d_model = 512
num_heads = 8
dim_feedforward = 2048
encoder_layer = TransformerEncoderLayer(d_model, num_heads, dim_feedforward)
input_tensor = torch.rand(32, 10, d_model) # (バッチサイズ, シーケンス長, d_model)
output, attn = encoder_layer(input_tensor)上記のコード例では、MultiHeadAttention クラスが複数のAttentionヘッドを実装しており、TransformerEncoderLayer クラスがAttention層を形成しています。各ヘッドは独立してクエリ、キー、バリューを計算し、その結果を統合して出力を生成しています。
9. 実際のLLMにおけるスケールと効率
LLMでは、Attention層とAttentionヘッドの数が非常に多く設定されており、これがモデルの表現力と計算負荷に直接影響します。例えば、GPT-3では96層のTransformerブロックが積み重ねられ、それぞれの層に96のAttentionヘッドが存在する場合もあります。これにより、モデルは膨大な数のパラメータを持ち、多様な言語パターンや文脈を学習・再現することが可能となっています。
しかし、ヘッド数や層数の増加は計算リソースの消費を増大させるため、効率的な実装や分散処理技術が求められます。また、最近の研究では、ヘッドの冗長性を削減し、効率性を向上させるための手法（例えば、ヘッドのプルーニングや共有）が提案されています。
10. まとめ
Attention層とAttentionヘッドは、TransformerベースのLLMにおいて、情報の動的な集約と多様な関係性の学習を可能にする重要な構成要素です。Attention層は、複数のAttentionヘッドを統合し、入力シーケンス全体に対する包括的な注意機構を提供します。一方、Attentionヘッドは、異なる視点や情報の側面に焦点を当て、モデルの表現力を高める役割を果たします。
これらのコンポーネントの効果的な設計と組み合わせにより、LLMは複雑な言語パターンを捉え、高度な自然言語処理タスクを遂行する能力を獲得しています。今後の研究や技術の進展により、さらに効率的で強力なAttention機構の開発が期待されます。
