TransformerモデルのAttentionメカニズムの核心と進化

要点（3行）

• TransformerのAttentionメカニズムは、系列データの長距離依存性を捉え、並列処理を可能にするLLMの基盤技術。

• トークン長に対する二次計算量とメモリ使用量が課題だが、FlashAttentionなどにより効率化、Mamba等のAttentionフリーモデルも登場。

• モデルの性能とコストに直結するため、利用目的に応じたAttention機構の選択と、進化する効率化技術の理解が不可欠。

背景（課題/先行研究/最新動向）

従来のRNNやLSTMモデルは、系列データにおける長距離の依存関係を学習することが困難であり、また本質的に逐次処理であるため、計算の並列化が難しいという課題を抱えていました。この課題に対し、2017年6月に発表された「Attention Is All You Need」論文は、Transformerモデルとその中核であるSelf-Attentionメカニズムを導入し、画期的な解決策を提示しました[1]。Transformerは、RNNを排除し、Attentionメカニズムのみで系列モデリングを行うことで、長距離依存性の効果的な捕捉と高い並列処理能力を実現しました。

しかし、Self-Attentionメカニズムは、入力系列のトークン長 $N$ に対して計算量が$O(N^2)$、メモリ使用量が$O(N^2)$となるため、特に大規模言語モデル（LLM）で非常に長いコンテキスト（数万～数十万トークン）を扱う際に、計算リソースとメモリの大きなボトルネックとなります。このスケーリングの課題が、その後の研究の主要な未解決点として注目されてきました。

最新動向

Mamba（2024年1月発表）: Linear-Time Sequence Modeling with Selective State Spacesが発表され、Attentionフリーでありながら長コンテキスト処理においてTransformerモデルに匹敵、あるいは凌駕する性能を示し、Attentionの代替アーキテクチャとして注目を集めています[3]。

これらの最新動向は、Attentionメカニズムが依然としてLLMの核心であり続ける一方で、その計算効率の課題を克服するための多様なアプローチが研究されていることを示しています。

提案手法 / モデル構造

TransformerのSelf-Attentionメカニズムは、入力シーケンス内の各トークンが、他のすべてのトークンとの関連度を学習することで、文脈情報を捉えることを可能にします。これは、各トークンのQuery (Q)、Key (K)、Value (V) の3つのベクトルを計算し、QとKの内積によって関連度（Attention Score）を算出し、そのスコアに基づいてVを重み付け加算する形で実現されます。

Self-Attentionの基本構造

Scaled Dot-Product Attentionは以下の式で表されます。 $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ ここで、$Q$はクエリ行列、$K$はキー行列、$V$はバリュー行列、$d_k$はキーベクトルの次元数です。この演算を複数回並行して実行し、その結果を結合することで「Multi-Head Attention」が構成され、モデルが異なる表現サブスペースから情報を学習できるようになります。

Efficient Attention: FlashAttention

FlashAttentionは、標準的なSelf-Attentionの計算をGPUのメモリ階層を考慮して最適化することで、高速化とメモリ効率化を実現した手法です[2]。主なアイデアは、Attention行列を全体として具体化（materialize）せずに、GPUの高速なSRAM（オンチップメモリ）上で計算を分割し、入出力（IO）のボトルネックを削減することです。これにより、GPUのグローバルメモリ（HBM）へのアクセスが大幅に減り、スループットが向上し、メモリ使用量も削減されます。

Mermaid図：TransformerエンコーダにおけるAttentionのフロー

graph TD
    A["入力埋め込み"] --> B["位置エンコーディング"]
    B --> C["Self-Attention層"]
    C --> D["Multi-Head Attention"]
    D -->|"Q,K,V生成"| E["ヘッド1のAttention"]
    D -->|"Q,K,V生成"| F["ヘッド2のAttention"]
    D -->|"..."| G[...]
    D -->|"Q,K,V生成"| H["ヘッドNのAttention"]
    E --> I["出力"]
    F --> I
    G --> I
    H --> I
    I --> J["結合"]
    J --> K["線形変換"]
    K --> L["Add & Norm"]
    L --> M["Feed Forward層"]
    M --> N["Add & Norm"]
    N --> O["出力埋め込み"]
    C --- メカニズム --> P["Scaled Dot-Product Attention"]
    P -->|"Q,K,V"| Q[Query]
    P -->|"Q,K,V"| R[Key]
    P -->|"Q,K,V"| S[Value]
    Q & R --> T["内積 (QK^T)"]
    T --> U["スケール (/√dk)"]
    U --> V[Softmax]
    V & S --> W["行列乗算 (Softmax(#quot;QK^T/√dk#quot;)V)"]
    W --> P

擬似コード：Self-Attentionの計算

# Self-Attention 計算の擬似コード

# 入力: Q (Query行列, shape: [batch_size, seq_len, d_k])
#       K (Key行列, shape: [batch_size, seq_len, d_k])
#       V (Value行列, shape: [batch_size, seq_len, d_v])
#       mask (オプション, shape: [batch_size, 1, seq_len, seq_len])
# 出力: Attention_Output (shape: [batch_size, seq_len, d_v])

function ScaledDotProductAttention(Q, K, V, mask=None):
    # 1. QとKの転置を乗算してAttentionスコアを計算
    scores = matmul(Q, transpose(K))

    # 2. スケーリング因子で除算
    d_k = K.shape[-1]
    scaled_scores = scores / sqrt(d_k)

    # 3. マスク適用 (オプション)
    if mask is not None:
        scaled_scores = scaled_scores + mask * (-infinity)

    # 4. Softmaxを適用してAttentionの重みを得る
    attention_weights = softmax(scaled_scores, axis=-1)

    # 5. Attentionの重みをValue行列に乗算
    output = matmul(attention_weights, V)

    return output

# Multi-Head Attentionの擬似コード
function MultiHeadAttention(Input_tensor, num_heads, d_model, mask=None):
    d_k = d_v = d_model / num_heads

    Q = linear_transform(Input_tensor, W_Q)
    K = linear_transform(Input_tensor, W_K)
    V = linear_transform(Input_tensor, W_V)

    Q_heads = split_into_heads(Q, num_heads)
    K_heads = split_into_heads(K, num_heads)
    V_heads = split_into_heads(V, num_heads)

    attention_outputs = []
    for i in range(num_heads):
        output_i = ScaledDotProductAttention(Q_heads[i], K_heads[i], V_heads[i], mask)
        attention_outputs.append(output_i)

    concatenated_output = concatenate(attention_outputs, axis=-1)
    final_output = linear_transform(concatenated_output, W_O)

    return final_output

計算量/メモリ/スケーリング

TransformerのSelf-Attentionは、その高い表現能力の代償として、計算量とメモリ使用量の課題を抱えています。

標準Self-Attention

• 計算量: 入力シーケンス長 $N$ に対して $O(N^2 \cdot d_{model})$。特にAttentionスコア行列 $QK^T$ の計算と、その後の $V$ との乗算がこの二次オーダーのボトルネックとなります。

• メモリ使用量: 最も大きなメモリ消費はAttentionスコア行列の保存であり、これも $O(N^2)$ のメモリを必要とします。また、推論時にはKVキャッシュ（過去のKeyとValueの埋め込み）が $O(N \cdot d_{model})$ のメモリを消費します。

FlashAttentionによる改善[2]

FlashAttentionは、GPUのメモリ階層（SRAMとHBM）に着目し、IO-Awareなアルゴリズム設計によって、計算量とメモリ使用量を劇的に改善しました。

• 計算量: $O(N^2 \cdot d_{model})$ の理論的な計算量は変わりませんが、GPUのSRAMを効率的に利用することでHBMへのアクセス回数を削減し、実測のスループットを大幅に向上させます。

• メモリ使用量: Attentionスコア行列をSRAM上でオンザフライで計算し、HBMには書き出さない「non-materialization」と、Softmaxの正規化項の再計算「recomputation」を行うことで、HBM上のメモリ使用量を $O(N \cdot d_{model})$ に削減します。

参考文献

1. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin. “Attention Is All You Need”. arXiv preprint arXiv:1706.03762, June 2017. https://arxiv.org/abs/1706.03762

2. Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, Christopher Ré. “FlashAttention-2: Fast and Memory-Efficient Exact Attention with IO-Awareness”. arXiv preprint arXiv:2307.08691, July 2023. https://arxiv.org/abs/2307.08691

3. Albert Gu, Tri Dao. “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”. arXiv preprint arXiv:2401.00997, January 2024. https://arxiv.org/abs/2401.00997