AIスタックの基礎：Transformer、Attention、勾配の流れ

Transformerはトークンの表現を文脈に応じた表現へ変換し、Attentionは位置間で情報を受け渡します。そして勾配の流れは、学習の信号が各パラメータへ届く仕組みを決めます。この3つを理解すると、学習、推論、モデルの適応、検索、評価というAIスタック全体をつなげて考えられます。

個人開発者として、基礎研究とプロダクトの判断を結びつけるためにまとめた技術ノートです。対象は主にdecoder-onlyの言語モデル。encoder-onlyやencoder–decoderでは、マスクや構造の選択が異なります。最新モデルの一覧ではなく、2023年までの基礎的な研究と、その実装上の意味を扱います。

1. ツールを選ぶ前に、スタックを整理する

AIは、数式、モデル、プロダクトの3つの層に分けると考えやすくなります。数式の層は表現や確率、行列演算、微分。モデルの層はそれらを学習可能な構造に組み立てます。プロダクトの層は、モデルをデータ、画面、検索、外部ツール、運用とつなぎます。

ベクトルデータベースを追加しても、Attentionのマスクの誤りは直りません。高速なカーネルも、根拠のない回答を正しくはしません。まず、問題がどの層にあるのかを確かめます。

2. トークン、埋め込み、学習目標

Tokenizerはテキストを整数IDへ変換します。トークンは単語とは限らず、単語の一部、記号、バイトなどにもなります。境界は言語とtokenizerに依存します。埋め込み行列から各IDの行を取り出すと、幅dのベクトルになります。n個の位置を並べた行列Xの形はn × dです。埋め込みは学習した座標であり、人間が読める単語の定義ではありません。

Causalな次トークン予測では、各ターゲットをその前のprefixだけから予測します。正解に割り当てた確率の負の対数を平均したものが損失です。

pθ(t₁, …, tT) = ∏ₜ pθ(tₜ | t<t)
ℒ(θ) = −(1/T) ∑ₜ log pθ(tₜ | t<t)
θ ← Optimizer(θ, ∇θℒ)

入力とターゲットは1位置ずれます。正解トークンを見ながらそれを予測できる実装では、損失が下がっても情報漏えいを学習している可能性があります。Paddingなどを損失から除外するマスクと、参照できる位置を制約するAttentionマスクは役割が異なります。

通常の推論では重みは固定です。プロンプトは計算と出力分布を変えますが、重みを更新しません。また、自然な続きを予測する目標と、現在のデータベースに照らして事実を確認する目標は異なります。

3. Scaled Dot-Product Attentionを形から理解する

1つのheadでは、Xを学習可能な行列でQuery、Key、Valueへ射影します。QueryとKeyが情報の受け渡しのスコアを決め、Valueが混合される内容を持ちます。Batch次元を省くと、形は次のようになります。

X: n × d
WQ, WK: d × dₖ       WV: d × dᵥ
Q = XWQ: n × dₖ
K = XWK: n × dₖ      V = XWV: n × dᵥ
S = QKᵀ / √dₖ: n × n
A = softmax_rows(S + M): n × n
Y = AV: n × dᵥ

Mは許可する位置で0、禁止する位置で負の無限大です。Causal self-attentionでは、位置iはj ≤ iを参照でき、未来は参照できません。SoftmaxはQueryごとの行に対し、Keyの方向へ適用します。全Keyを禁止する行がある場合は、無効なsoftmaxを避ける明示的な処理が必要です。

√dₖで割る理由は、QueryとKeyの各成分が独立、平均0、分散1なら、内積の分散がdₖになるためです。この仮定のもとでスケールを整え、初期状態のsoftmaxが極端に尖る傾向を抑えます。学習後の全分布に対する保証ではありません。

Multi-head attentionは複数の混合を計算し、連結して出力射影を適用します。各headに文法や感情などの意味が事前に割り当てられるわけではありません。Vaswaniらの原著論文はこれらをencoder–decoderとして組み合わせました。Decoder-onlyモデルは、その構成要素を引き継ぎ、変更しています。

4. 小さなAttentionを手で計算する

q = [1, 0]と、3組のKey・Valueを使います。3位置の最後から参照するcausalなprefixなら、3つすべてが見えます。学習済みモデルから取り出した値ではなく、計算のために用意した例です。

q = [1, 0]
k₁ = [1, 0]    v₁ = [1, 0]
k₂ = [0, 1]    v₂ = [0, 2]
k₃ = [1, 1]    v₃ = [2, 1]

s = [1/√2, 0, 1/√2]
a ≈ [0.401112, 0.197776, 0.401112]
y = a₁v₁ + a₂v₂ + a₃v₃
y ≈ [1.203336, 0.796664]

重みの合計は1で、出力はベクトルです。その先には射影、非線形変換、さらに多くの層があります。重みが最大の位置だけで、最終回答の理由は説明できません。JainとWallaceの研究は、Attentionを説明として解釈する際の注意点を実験で示しています。

数値的には、許可されたスコアの最大値を引いてからsoftmaxを計算します。行全体に同じ定数を足しても分布は変わりません。上の実験で動かす温度は、この人工的なAttention分布に対するものです。出力トークンをサンプリングするときの温度とは別の場所で作用します。

5. Attentionを微分する

1行についてa = softmax(s)、y = Σⱼ aⱼvⱼと書きます。SoftmaxのJacobianの要素はaᵢ(δᵢⱼ − aⱼ)。これを加重和へ代入すると、簡潔な式になります。gをyに関する上流の損失勾配とします。

∂aᵢ/∂sⱼ = aᵢ(δᵢⱼ − aⱼ)
∂y/∂sⱼ = aⱼ(vⱼ − y)
∂ℒ/∂sⱼ = aⱼ gᵀ(vⱼ − y)
∂ℒ/∂vⱼ = aⱼg

sⱼ = qᵀkⱼ / √dₖ
∇qℒ = (1/√dₖ) ∑ⱼ (∂ℒ/∂sⱼ) kⱼ
∇kⱼℒ = (∂ℒ/∂sⱼ) q / √dₖ

QueryとKeyを通る勾配は何を参照するかを変え、Valueを通る勾配は渡す内容を変えます。Self-attentionでは3つともXから作るため、X全体の勾配には3つの射影経路が加わります。上のスコア微分では、Valueを固定して経路を分けています。

次のPythonは標準ライブラリだけで実行できます。先ほどの出力を再現し、ℒ = [0.3, −0.2] · yという小さな損失で、解析的なスコア勾配を中心差分と比較します。

import math

def softmax(scores):
    peak = max(scores)
    exps = [math.exp(score - peak) for score in scores]
    return [value / sum(exps) for value in exps]

scores = [1 / math.sqrt(2), 0, 1 / math.sqrt(2)]
values = [[1, 0], [0, 2], [2, 1]]
upstream = [0.3, -0.2]

def output(logits):
    weights = softmax(logits)
    return [sum(a * v[d] for a, v in zip(weights, values))
            for d in range(2)]

def loss(logits):
    return sum(g * y for g, y in zip(upstream, output(logits)))

weights, y = softmax(scores), output(scores)
analytic = [a * sum(g * (v[d] - y[d])
                   for d, g in enumerate(upstream))
            for a, v in zip(weights, values)]
epsilon = 1e-5
numeric = []
for j in range(len(scores)):
    plus, minus = scores.copy(), scores.copy()
    plus[j] += epsilon
    minus[j] -= epsilon
    numeric.append((loss(plus) - loss(minus)) / (2 * epsilon))

assert abs(sum(weights) - 1) < 1e-12
assert max(abs(a - b) for a, b in zip(analytic, numeric)) < 1e-7
print("weights:", [round(a, 6) for a in weights])
print("output:", [round(value, 6) for value in y])
print("score gradients:", [round(value, 6) for value in analytic])

中心差分は局所的な微分を確かめる方法です。倍精度と極端でない入力を使い、摂動幅を選びます。小さすぎる幅は丸め誤差を、大きすぎる幅は曲率の影響を受けます。モデル全体の品質を保証するテストではありません。

6. Transformer blockを組み立てる

Attentionは位置間で情報を混ぜ、位置ごとのMLPは特徴を変換します。MLPの重みは位置間で共有されます。残差加算は、各sublayerが変更できる情報の流れを保ちます。Dropoutなどを省いたpre-normalizationのblockは次の形です。

U = X + MultiHeadAttention(Norm(X))
Xnext = U + MLP(Norm(U))
MLP(z) = W₂ φ(W₁z + b₁) + b₂

MLPの式は1トークンを列ベクトルとして書いています。先ほどのXはトークンを行に並べた表記です。実装時には、選んだ表記に合わせて重みの軸を配置します。

RoPEは位置に応じてQueryとKeyの座標ペアを回転します。内積には(Rᵢq)ᵀ(Rⱼk) = qᵀRᵢᵀRⱼkという相対回転が現れ、位置の関係を与えます。任意の長さへの正しい外挿を保証する仕組みではありません。RoFormer論文で導出を確認できます。

LayerNormは平均を引いて再スケールします。RMSNormは平均を引かず、二乗平均平方根で再スケールします。RMSNorm論文はこの簡略化を調べています。正規化、活性化、位置表現、headの構成はモデルごとに異なります。

7. 深いネットワークを流れる勾配

誤差逆伝播は連鎖律を逆向きに適用します。Jₗを層の入力に関する出力のJacobian、gₗを入力での損失勾配とすると、行列の積が信号を深さ方向に運びます。

hₗ₊₁ = fₗ(hₗ)
Jₗ = ∂hₗ₊₁/∂hₗ
gₗ = Jₗᵀgₗ₊₁
g₀ = J₀ᵀJ₁ᵀ … J_(L−1)ᵀ g_L

すべての層の微分がaというスカラーなら、勾配の倍率はaᴸです。a = 0.9、L = 100では約0.0000266、a = 1.1では約13,781になります。実際の行列では方向と特異値が関わるため、この例は減衰と増幅の直感を示すものです。

残差層hₗ₊₁ = hₗ + Fₗ(hₗ)のJacobianはI + JFₗです。恒等写像の項が直接の微分経路を加えます。深い残差学習の構造的な考え方は、Transformerにも広く使われています。

Residual layer: hₗ₊₁ = hₗ + Fₗ(hₗ)
Gradient: gₗ = (I + JFₗ)ᵀgₗ₊₁
Scalar toy: direct gain = |a|ᴸ
Scalar toy: residual gain = |1 + a|ᴸ

恒等写像があっても、JFₗ = −Iなら打ち消されます。増幅が重なれば発散も起こります。上の勾配の実験では両方を確認できます。実際の安定性には、残差分岐、正規化、初期化、数値精度、最適化が関わります。Pre-normalizationでは、分岐内の正規化の微分もJFに含まれます。

元のTransformerはpost-normalizationを使っていました。Xiongらの研究は、正規化を残差分岐の内側へ移すと初期化時の勾配がどう変わるかを解析しています。Pre-normalizationにすれば常にwarm-upや調整が不要、という一般的な保証ではありません。

8. 学習ループと失敗の切り分け

学習ではbatchを選び、順伝播、損失計算、微分、optimizerの更新を行います。Gradient accumulationは複数のmicrobatchをまとめて更新するため、意図した平均になるよう損失をスケールする必要があります。

Gradient clippingは過大な信号を制限しますが、消えた勾配を戻せず、間違った目的関数も直せません。低精度学習では形式と累積精度が重要です。Activation checkpointingは再計算を増やす代わりに、保存する活性値を減らします。損失曲線と対象タスクの両方で変更を評価します。

9. 何を可能にしたブレークスルーなのか

現代の言語モデルは、異なる制約を解く研究の組み合わせから生まれました。

Chinchillaの結果は、学習計算量をパラメータとデータへどう配分するかに関するものです。推論回数が総費用を支配するプロダクトでは、別の配分が有利な場合があります。普遍的なトークン対パラメータ比で全てが決まるわけではありません。

指示学習と選好最適化は、事前学習後の振る舞いを変えます。DPOでは学習するpolicyと参照policyの対数確率比を使い、同じプロンプトに対する選ばれた回答の相対的な優位を高めます。選好データの品質や範囲は依然として重要で、事実の正しさを独立に保証するものではありません。

10. 推論：prefill、decode、KVメモリ

Prefillではプロンプトを処理し、層ごとのKeyとValueを作ります。Decodeでは生成したトークンの状態を追加し、prefixを参照します。KVキャッシュは以前のKey・Value射影を再利用し、その再計算を避けます。保存するのは系列に依存する活性値で、知識や重みの更新ではありません。

層の構成と系列長が一定のbatchを仮定すると、おおよそのKV payloadは次のようになります。

KV bytes ≈ 2 × B × L × T × Hkv × dh × b
B: batch size       L: layers       T: cached tokens
Hkv: KV heads       dh: head width  b: bytes per element

B=1, L=32, T=8192, Hkv=8, dh=128, b=2
KV payload = 1,073,741,824 bytes = 1 GiB

この例には重み、allocatorのoverhead、一時bufferなどは含みません。同じ条件でKV headを32にすると4 GiBです。Grouped-query attentionはQuery headのグループ間でKeyとValueを共有し、この項を減らします。品質と学習も関わる構造の選択です。

FlashAttentionはtilingとonline softmaxを使い、高帯域メモリへの転送を減らして、n × nの行列全体を保持せずに厳密なAttentionを計算します。密なprefillの二次の演算量は残ります。キャッシュを使うdecodeでは伸びるprefixを参照し、射影とMLPにも費用がかかります。原著論文とPyTorchのSDPAドキュメントで、最適化されたインターフェースやmask、dropoutの仕様を確認できます。

長い入力は最初のトークンまでの待ち時間とメモリを増やすことがあります。Time to first tokenと出力tokenのthroughputは分けて測ります。量子化で重みやキャッシュの容量を減らす際も、モデル、形式、タスクごとの品質を確かめます。

11. プロンプト、RAG、LoRAを使い分ける

元のRAG研究は、パラメトリックな生成と検索可能な索引を組み合わせました。典型的な推論時のRAGアプリでは、検索結果を文脈へ追加して回答し、その文書を通じた逆伝播は行いません。検索の品質、文書の権限、回答が根拠に従っているかは別々に評価します。

LoRAではベース行列W₀を固定し、低ランクの更新を学習します。W₀がd_out × d_inなら、Aはr × d_in、Bはd_out × rです。

W = W₀ + (α/r)BA
Trainable parameters = r(d_in + d_out)
Full matrix parameters = d_in × d_out

rが小さければ学習するパラメータ数は減ります。ただし、ベースモデルや活性値などの全費用がなくなるわけではありません。LoRA論文はこの適応を調べています。根拠を与える検索と、振る舞いを学習する調整を区別し、最後はアプリケーション全体を評価します。

12. 設計演習：Shopifyのサポートアシスタント

商品と配送の質問に答える仕組みを考えます。まず、正しい情報の置き場所を整理します。商品カタログ、規約文書、認証された注文状況APIです。モデルは質問の解釈と文章化を担当し、アプリケーションがアクセスを制御して許可されたツールを実行します。

規約の質問では、関連箇所を文書IDとともに検索します。注文の質問では、ユーザーの権限を確認し、構造化された注文データを取得します。モデルには返された状態を説明させ、状態自体を創作させません。検索文書には誤解を招く指示が含まれる可能性があるため、アプリケーションの指示とは分けて扱います。

構成を増やす前に評価用の質問を作ります。回答可能・不可能な質問、古い規約、矛盾する根拠、検索結果なし、別の顧客の注文へのアクセスなどを含めます。根拠への忠実さ、ツール選択、認可、速度、費用を別々に測ります。トークン損失だけでは、この評価を置き換えられません。

13. 数式と実装をつなぐチェック

1. 因果性： dropoutを無効にした評価モードで未来のトークンを変えても、それ以前の出力が変わらないことを確かめます。
2. 確率の性質： 許可されたAttention重みの合計は1、禁止位置は0、行への定数加算で分布は不変であることを確認します。
3. 微分の一致： 小さな倍精度の例で、解析的・自動微分の結果と中心差分を比較します。
4. 形の確認： Query head数とKV head数、行ベクトルと列ベクトルの表記を区別します。
5. メモリの計算： KV payloadを見積もり、その他の費用も含む実測peakと比較します。
6. タスクの評価： 現実的な要求をholdoutにし、失敗の種類、品質、速度を測ってから構成を変えます。

何が変わったか、どの仮定が崩れたかを観察できるようにします。構造、データ、プロダクトのどこに問題があるのかを区別するための土台になります。

14. 基礎から、次の実験へ

小さなAttentionの出力を再現し、その微分を確かめ、実際のtensorの形で1つのTransformer blockを追うのが次の実験です。同じ例の順伝播と逆伝播を確認してから、規模を大きくしていきます。各節に挙げた原著論文が、ブレークスルーをさらに学ぶ出発点です。

プロダクト側の取り組みは開発実績にまとめています。GradFlowLabの開発者についてもご覧ください。