← ブログ一覧へ

モデル学習ノート · 第2章

注意機構：情報の流れを工学の視点で考える

Q/K/V と自己注意・交差注意の違いから、KV キャッシュの設計までをつなぎます。

Muxuan Liu · 自分のモデル学習ノートを整理し、式と実装例を確認した版です。

情報のルーティングとして考える

注意機構は、情報をどこから集めるかを決める仕組みとして考えると理解しやすくなります。query は必要な情報、key は各位置が提供する手がかり、value は実際に混ぜ合わせる情報です。query と key を比較して重みを求め、value を集約します。

言語関係の学習はその応用の一つで、画像などにも同じ発想を使えます。一つの注意の重みだけで、文法関係やモデルの判断理由が証明されるわけではありません。

自己注意と交差注意を区別する

自己注意では同じ系列から Q・K・V を作り、query の数だけ出力を返します。翻訳結果を何トークン生成するかを直接決める機構ではありません。交差注意は一方の系列から Q、別の系列から K/V を作るため、両者の長さが異なっていても計算できます。

RNN のエンコーダー・デコーダーも、異なる入出力長を扱えました。初期の注意は、入力全体を固定長ベクトルへ圧縮する負担を和らげます。Transformer の交差注意は両系列をつなぎ、出力長は生成と停止の仕組みが決めます。デコーダーのみのモデルに、独立した交差注意が必ずあるわけではありません。

形状に沿って計算する

一つのヘッドで Q は [Tq, Dh]、K は [Tk, Dh]、V は [Tk, Dv] です。QKᵀ は [Tq, Tk] になり、key の方向へ softmax を適用します。√Dh で割ることでスコアの尺度を調整します。このスケーリングを持つ式は 2017 年の Transformer のもので、以前の Bahdanau 注意は別のスコア関数です。

多頭注意では異なる部分空間の結果を連結して出力投影します。複数のパターンを学べますが、「動詞の担当ヘッド」「時間の担当ヘッド」は説明用の例であり、事前に決めた役割ではありません。

Attention(Q, K, V) = softmax(QKᵀ / √Dh + mask) V MultiHead = Concat(head₁, …, headₕ) Wₒ

並列学習と逐次生成は別の段階

学習時には正解系列が既知なので、因果マスクを用いて複数位置をまとめて計算できます。一方、通常の自己回帰生成は、生成済みトークンを使って次を決めるため逐次的です。

KV キャッシュは既存の接頭部の key/value を保存して再利用します。因果注意と変わらない接頭部・設定の下では、新しいトークンを加えても過去位置の状態は変わりません。重複計算は減らせますが、新しい query と過去位置の注意計算は必要です。

MHA・MQA・GQA・MLA の違い

MHA は各ヘッドに K/V を持ち、MQA は全 query ヘッドで一組を共有します。GQA では h 個の query ヘッドに g 個の K/V ヘッドを対応させます。g=h なら MHA、g=1 なら MQA です。通常の各層の KV 保存量は、他条件を固定すると 2·B·T·g·Dh に比例します。実際の速度には内核や配置も影響します。

2024 年の DeepSeek-V2 が記述する MLA は、K/V の潜在表現を圧縮し、位置情報の部分も別に扱います。単に GQA の g を減らした構造ではありません。比較表は設計の違いを示し、同じ品質や普遍的な速度順位を保証するものではありません。

構造Query ヘッドK/V の表現
MHAhh 個の独立ヘッド
MQAh1 組を共有
GQAhg 組を共有（1 ≤ g ≤ h）
MLA複数圧縮した潜在表現と位置の成分

実装で省けない transpose

重み付き和の後の形状は [B, H, Tq, Dh] です。まず [B, Tq, H, Dh] へ戻してからヘッドをまとめます。直接 view すると、系列とヘッドの配置を誤って混ぜてしまいます。

例のブールマスクは True が許可です。全 key を禁止した行は、−∞ の softmax が未定義になるため拒否します。API ごとにマスクの意味を確認してください。これは理解用の実装で、完全な KV キャッシュや融合カーネルは含みません。

PyTorch の例を開く
import math
import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    """Teaching example. Mask: True = allowed, [Tq,Tk] or [B,1,Tq,Tk]."""
    def __init__(self, d_model, num_heads, dropout=0.0):
        super().__init__()
        if num_heads < 1 or d_model % num_heads:
            raise ValueError('d_model must be divisible by num_heads')
        self.h = num_heads
        self.d = d_model // num_heads
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, allowed=None):
        if k.shape[1] != v.shape[1]:
            raise ValueError('Key and value lengths must match')
        b, tq, width = q.shape
        def split(x, projection):
            return projection(x).reshape(b, -1, self.h, self.d).transpose(1, 2)
        Q = split(q, self.q_proj)  # [B,H,Tq,Dh]
        K = split(k, self.k_proj)  # [B,H,Tk,Dh]
        V = split(v, self.v_proj)  # [B,H,Tk,Dh]
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d)
        if allowed is not None:
            if allowed.dtype != torch.bool:
                raise ValueError('Use a boolean mask: True = allowed')
            if not allowed.any(dim=-1).all():
                raise ValueError('Each query must have at least one allowed key')
            scores = scores.masked_fill(~allowed, float('-inf'))
        weights = self.dropout(scores.softmax(dim=-1))
        out = weights @ V  # [B,H,Tq,Dh]
        out = out.transpose(1, 2).contiguous().reshape(b, tq, width)
        return self.out_proj(out)

参考文献