平均から導く attention と transformer ブロック

1997年から2017年ごろまで、その要約は recurrent state でした。ベクトル $\mathbf{h}$ を保持し、各 token で $\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t)$ と更新します。固定サイズ、可変入力、形としてはまさに正しいものです。

それは3つの点で失敗し、この章の architecture はその3つすべてに答えます。$T$ ステップを backpropagation すると $T$ 個の Jacobian を掛け合わせるため、gradient は消えるか爆発します。これは Chapter 5 が単一の $\tanh$ node の中で測った病です。LSTM¹ はまさにそれに対抗するために設計され、利用可能な範囲を数十ステップから数百ステップへ押し広げましたが、token 5 の情報が token 500 に届くには495回の逐次更新を生き延びなければならない、という事実は変わりませんでした。source 全体を1本のベクトルに収めなければならなかったのです。sequence-to-sequence 翻訳²では、encoder が入力を最終状態へ圧縮します。Bahdanau、Cho、Bengio は2014年、transformer の3年前にそのボトルネックに名前を付け、自分で計算した重みによって decoder が すべて の encoder states の weighted sum を取れるようにして修正しました。³ 以下のすべてはその考えを、sequence が自分自身に適用し、recurrence を削除したものです。そして 更新は構造上逐次的です。$\mathbf{h}_t$ には $\mathbf{h}_{t-1}$ が必要で、1万 cores を持つ GPU でもそれには何もできません。勝った architecture は明らかに賢かったわけではありません。高価なステップが行列積だったものが勝ったのです。

もう1つの古典的な inductive bias である convolution、つまり小さな filter を入力全体に滑らせ、どこで検出された feature もどこでも検出されるようにするものも、ここでは作りません。これは画像にはほぼぴったりで、vision course に委ねます。recurrence も convolution もこのページ以降には再登場しません。だから章を割きません。Chapter 1 が約束した通り、省略は黙って行うのではなく明示します。

可変個のベクトルを受け取り1本のベクトルを返す、最も明白な関数は 平均 です。

入力数はいくつでもよく、出力サイズは固定で、微分可能で、無料です。embedding テーブル、この平均、vocabulary への linear layer を足せば、15行で完全な language model になります。同時に、それはひどい model でもあります。そして、どうひどいのかが導出全体です。

下の corpus は1 MB の Shakespeare、1,115,394文字を、Chapter 7 で作った種類の byte-level BPE tokenizer、vocabulary 1024で通したものです。459,760 tokens、1 token あたり2.43文字で、90/10に分割しています。すべての model は幅128、128 tokens を見て、batch 64、$10^{-3}$ の AdamW で3000 steps 学習します。perplexity は held-out split 上です。⁴

2行目を二度読んでください。context を平均しても少し助けになるどころか、context を完全に無視するより 4倍悪く なります。理由は2つあり、どちらも経験的というより証明可能です。

平均は順序を見られません。 加算は可換なので、window を shuffle しても summary は変わりません。近似的にではありません。

A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)          # rows of the averaging matrix
y = x[torch.randperm(T)]                # the same tokens, shuffled
print((A[-1] @ x - A[-1] @ y).abs().max().item())

2.9802322387695312e-08

並べ替えた和に浮動小数点 noise が乗るだけで、2つの summaries は同じベクトルです。context への唯一の視界が平均である model は、the dog bit the man と the man bit the dog を区別できません。3行目は、入力に positions を足しても直らないことを示しています。各 token に learned position embedding を足してから平均しても、188点中2.14点しか稼げません。positions は和に入り、和はそれを忘れます。

そして平均は現在を溺れさせます。 position 100では、現在の token は summary の100分の1です。これにはすでに持っている安い修正があります。token を残し、summary をそこへ 足す のです。これは Chapter 6 の residual connection で、4行目がその結果です。希釈を修理すると、一様平均は まったく何も 貢献しません。baseline 59.71に対して60.45です。すべての token はそこにあり、等しく重み付けされています。そして等しい重み付けは、情報がないのと同じです。

問題は平均ではありません。重み です。

伸びていく prefix の平均は loop に見えます。実際には、行和が1の下三角行列による1回の乗算です。そしてそれは、正確に softmax でもあります。

loop = torch.stack([x[:t + 1].mean(0) for t in range(T)])   # the obvious version

A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)
mat = A @ x                                                  # the same thing

S = torch.zeros(T, T).masked_fill(torch.tril(torch.ones(T, T)) == 0, float("-inf"))
soft = F.softmax(S, dim=-1) @ x                              # and the same thing again

loop vs matmul   max |diff| = 5.960464477539063e-08
loop vs softmax  max |diff| = 5.960464477539063e-08

the averaging matrix A (rows sum to 1, upper triangle is zero):
  1.000 0.000 0.000 0.000 0.000 0.000
  0.500 0.500 0.000 0.000 0.000 0.000
  0.333 0.333 0.333 0.000 0.000 0.000
  0.250 0.250 0.250 0.250 0.000 0.000
  0.200 0.200 0.200 0.200 0.200 0.000
  0.167 0.167 0.167 0.167 0.167 0.167

transformer の名前付き部品が3つ、いま画面上にあります。三角形は causal mask で、目的によって強制されます。position $t$ が position $t{+}1$ を見られたら、答えが入力の中にあります。これは Chapter 6 が監査せよと言った leak が、architecture の内側にあるだけです。softmax は mask の実装方法 です。禁じられた entry を $-\infty$ に設定すると、それらは正確にゼロになり、残りが正規化されます。つまり masking と normalising は1つの操作です。（-1e9 ではなく $-\infty$ を使います。それが masking の 意味する 値であり、float16 に cast しても $-\infty$ として生き残り、たまたまいる範囲に対して選んだ定数が十分大きいかを決めずに済みます。これは Chapter 2 の floating-point box が、答えなくてよい問いを投げているだけです。）そして scores は自由な parameter です。一様平均は、許可されたすべての score が同じ数であるときに得られるものです。そこへ任意の数を置けば、softmax がそれを有効な重みに変えます。

この章の残りは1つの問いです。その数はどこから来るのでしょうか。

それらは素の parameters ではあり得ません。learned $T \times T$ matrix はすべての文で同一になり、4 tokens 前を見よとは encode できても、この pronoun が指す noun を見よとは決して言えません。position $t$ と position $i$ を結ぶ重みは、両方 の位置にあるものに依存しなければなりません。relevance は性質ではなく関係だからです。単語 it はそれ自体で relevant なのではなく、何か に対して relevant なのです。

2本のベクトルから数を返す最安の関数は、Chapter 1 の dot product です。position $i$ を position $t$ に対して $\mathbf{x}_t \cdot \mathbf{x}_i$ と score すれば、仕組みは動きます。ただし2つの点でひどく動き、そのせいで残りすべてが必要になります。ベクトルと自分自身の dot product は squared norm なので、各 token はほとんど自分自身に attend します。そして関係が 対称 になります。it が animal に強く attend するなら、animal も it に強く attend することになりますが、言語ではこれは偽です。adjective は noun を、noun が adjective を必要とするよりずっと強く必要とします。

そこで各 token に2つの役割を与えます。それは token の learned linear maps としての2つです。この位置が 探しているもの、$\mathbf{q}_t = W_q\mathbf{x}_t$、つまり query。そして、この位置が 見つけられるために提示するもの、$\mathbf{k}_i = W_k\mathbf{x}_i$、つまり key。$\mathbf{q}_t \cdot \mathbf{k}_i$ と score すれば対称性は消えます。$W_q \neq W_k$ だからです。token はあるものを宣伝し、別のものを探せます。

まだ1つ問題があります。weighted sum は $\mathbf{x}_i$ 自身の上で取っていましたが、それでは コピー されるものが matching に使われるものと同じであることを強制してしまいます。matching には token を識別する features が欲しい一方、copying には下流で役に立つ features が欲しいのです。そこで3つ目の map、$\mathbf{v}_i = W_v\mathbf{x}_i$、つまり value を学習し、それを合計します。

式はもう帳簿付けです。

ここで $M$ は causal mask で、対角線上および下はゼロ、上は $-\infty$ です。code では30行、そのうち20行は shape です。

class Head(nn.Module):
    """One head of causal self-attention."""

    def __init__(self, d_model, d_head, block):
        super().__init__()
        self.q = nn.Linear(d_model, d_head, bias=False)      
        self.k = nn.Linear(d_model, d_head, bias=False)      
        self.v = nn.Linear(d_model, d_head, bias=False)      
        self.d_head = d_head
        self.register_buffer("mask", torch.tril(torch.ones(block, block)).bool())

    def forward(self, x):
        T = x.shape[1]
        q, k, v = self.q(x), self.k(x), self.v(x)
        s = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)          
        s = s.masked_fill(~self.mask[:T, :T], float("-inf"))          
        w = F.softmax(s, dim=-1)                                      
        return w @ v                                                  

Score、mask、normalise、mix。それ以外は projection です。

$\sqrt{d_k}$ のほぼすべての説明は「softmax が saturate しないように」と言います。それは本当ですが、何も説明していません。議論は Chapter 2 の variance から2行です。$\mathbf{q}$ と $\mathbf{k}$ の entries が平均ゼロ、variance 1で独立なら、各積 $q_j k_j$ の variance は1で、独立なものの variances は足し合わされます。

したがって scores の standard deviation は $\sqrt{d_k}$ です。2万組の random pairs で測るとこうなります。

     d     Var(q.k)         std   sqrt(d)
     4        3.975       1.994     2.000
    16       16.071       4.009     4.000
    64       64.249       8.016     8.000
   256      253.065      15.908    16.000
  1024     1015.562      31.868    32.000

これが重要な理由は、softmax が linear layer とは違う形で scale-sensitive だからです。linear layer の入力を2倍にすると出力も2倍になります。softmax の前で scores を10倍すると、soft blend は hard choice に変わります。64 scores の1行を、割り算ありとなしで見ます。

「effective tokens」は entropy の指数です。その行が実際に何 positions の平均を取っているかを表します。割らない場合、$d_k = 256$ では freshly initialised head が64 tokens 中 ちょうど1つ に attend し、それは random draw だけで選ばれます。

これは forward で悪く、backward ではさらに悪いものです。その形は Chapter 5 がすでに $\tanh$ 上で測りました。1つの entry に committed した softmax は derivative がほとんどありません。その Jacobian の対角は $w_i(1-w_i)$ で、両端でゼロです。2,000 random rows ではこうなります。

$d_k = 1024$ では、10行中7行が training 開始前に凍っており、凍った状態で始まった head は何を見るべきかを学べません。割れば、その量はすべての幅で0.96に平坦化し、何も saturate しません。

ここからが誰も論文に載せない部分です。最終的な perplexity は変わるのでしょうか。 割り算を消して、4つの head widths で学習します。

最初の2行は上の3000-step budget から来ています。最後の3行は短い run です。1500 steps、batch 32、one head、projections 前の normalisation なしで、両 variants を同一設定にしています。

$d_k = 32$ では割り算の価値はなく、割らない run がごくわずかに先行しています。だからといってそれを外してよいわけではありません。256では perplexity 18点分、1024では17点分の価値があるからです。仕組みは scores 自体に見えています。

割らない head は回復しません。暴走 します。scores の standard deviation は initialisation 時の21から5147へ行き、attention entropy はゼロへ落ち、99.9 % の rows が weight の0.99超を単一 token に置きます。head が hard selector になった時点で gradient はほぼゼロになり、戻す力がなくなるため、collapse は安定します。割った head は同じ training 後に score standard deviation 3.44にとどまり、まだ変更可能な soft blend です。

Vaswani らはまさにこれだけを述べています。products は $d_k$ が大きい値で「magnitude が大きくなる」と推測し、割ります。⁵ large という語が重みを担っています。そして tables は large がどこから始まるかを示します。32では何もなく、256ではすべてです。

1つの head は position ごとに1つの softmax row なので、「ここで relevant なのは何か」への答えを1つ持ちます。the animal that crossed the wet street の the の後の単語を予測するには、syntactic slot、subject、previous token が同時に必要で、1つの probability distribution を3箇所に集中させることはできません。そこで複数の heads を並列に走らせ、それぞれ幅 $d_{\text{model}}/h$ とし、concatenate して、さらに1つの matrix $W_o$ で mix します。幅を増やしたのではなく、分割したのです。

attention がすることも1つだけです。positions 間 で情報を動かします。上の code にあるすべての operation は feature 軸に沿って linear であり、Chapter 5 は linear maps の stack が何になるかを証明しました。したがって各 block は、各 position に独立に 適用される小さな MLP も持ちます。幅を4倍に広げ、中央に GELU を置いて戻ってきます。分業は覚える価値があります。attention は positions across を mix し、feed-forward network は position within で計算します。

完全な ladder は次の通りで、各行は上の行に1部品を足しています。

learned weights は uniform weights に対して perplexity 14点勝ちます。これがこの章の議論全体を1行にしたものです。4 heads は16,512個の追加 parameters でさらに3点を買います。そして同じ head は、replacing より added のほうが9点分価値があります。attention は情報を持ち込みますが、position が何であるかを決めるものではありません。

次に、parameters が実際にどこにあるかです。図だけを見た人はよく驚きます。

すべての transformer block の3分の2は feed-forward network です。どの scale でも同じです。attention は4つの $d \times d$ matrices を持ち、MLP はそれに相当する8つを持つからです。model が何を知っているにせよ、それを保持する parameters の大半は per-position MLP にあります。

LayerNorm は Chapter 6 で作られ測られ、この章ではそこで残した形のまま使います。residual connections はそこで名前が付けられ ablation され、ここで組み込まれます。上の「added, not replacing」の行は residual connections で、平均では perplexity 188点、one head では9点の価値があります。LayerNorm⁷ は各 example をその features across で normalise します。Chapter 6 は BatchNorm ではなくこれがここで生き残った理由を示しました。batch への依存がなく、running statistics がなく、training と inference で同一で、sequence length に無関心であることです。これらすべては、1人の user のために1 token ずつ生成するときに要件になります。Chapter 13 はそこへ向かいます。cost は768 parameters で、perplexity 1.8点を買います。

class Block(nn.Module):
    def forward(self, x):
        x = x + self.att(self.ln1(x))     
        x = x + self.ff(self.ln2(x))      
        return x

normalisation がどこにあるかを見てください。各 sub-layer の input の上にあり、input から output への residual path は決して normalise されません。これが pre-norm です。2017年の paper は逆、x = LayerNorm(x + Att(x))、つまり post-norm を行い、residual path 自体に LayerNorm を置きます。

Xiong らは initialisation 時の gradient を通じて違いを説明しました。post-norm network ではそれが depth に対してひどく scale されており、元の transformer が学習するだけでも learning-rate warmup を必要とした理由です。⁸ 12 blocks、1000 steps、learning rate $3 \times 10^{-3}$ です。

gradient norm per block at initialisation, before any step
  pre-norm    block 1 0.0498 ... block 12 0.0657   ratio last/first  1.32
  post-norm   block 1 0.0977 ... block 12 0.1613   ratio last/first  1.65

  pre-norm,  no warmup          perplexity   37.82
  pre-norm,  200-step warmup    perplexity   37.62
  post-norm, no warmup          perplexity  308.05
  post-norm, 200-step warmup    perplexity   37.88

warmup なしの post-norm は8倍悪く、warmup ありの post-norm は pre-norm と完全に一致します。warmup はここでは一般的な良い practice ではありません。normalisation の特定の配置に対する patch であり、LayerNorm を移動するとそれは不要になります。だから2019年以降のほぼすべての model は pre-norm であり、2017年の図は specification ではなく歴史として読むべきなのです。

position embeddings を削除しても model は学習します。ただし、何がどこにあるかをまったく判別できません。これは training failure ではなく symmetry です。attention score のどこにも $t$ や $i$ 自体は現れないので、入力を permutation すると出力も permutation します。self-attention は permutation-equivariant です。これは平均の order-blindness が、より良い変装をしたものです。causal mask は各 position が異なる prefix を見るため一部の order を回復しますが、prefix 内ではすべての orderings が同じです。

position を注入する4つの方法を、64-token windows で学習し、見たことのない長さを越えて64、128、256で評価します。

Learned absolute embeddings、つまり position ごとに1本のベクトルを token に足す方法は、学習した長さでは勝ち、その後崖から落ちます。position 100は batch に現れたことがなく、その embedding は始まったときの random vector のままだからです。Sinusoids、元の選択肢は、learn されるのではなく、幾何級数的に間隔を置いた周波数の sines と cosines から計算されます。2017年の paper はそれが extrapolate すると期待しましたが、table はそうではないと言っています。関数は position 200で定義されていますが、model はそこでそれを読むことを学んでいません。RoPE⁹ は何も足さず、代わりに query と key を position に比例した角度で2次元 slices の中で 回転 させます。dot product の両側を同じように回転させても値は変わらないので、score は $t - i$ のみに依存するようになり、position は無料で relative になり、尽きる table がありません。劣化はしますが、劣化するだけです。ALiBi¹⁰ はここで最も単純で、最も奇妙な結果です。distance に比例する score への linear penalty で、head ごとに異なる slope を持ちます。その perplexity は training length を越えて window が伸びるにつれ、44.95から42.49へ 改善 します。penalty はどんな distance でも定義され、各 head は学習したことを続けるからです。

教訓は table より長く残ります。architecture が何かを represent できないことは、その範囲を一度も learned していないこととは別の問題であり、噛みついてくるのは後者です。これはまた、「context を128Kへ拡張した」という発表の背後にある machinery でもあります。それらはほぼ常に rotary encoding の re-scaling であり、Chapter 16 が context limit は消えるのではなく動くと言う理由です。

Dropout も同じように継承されます。softmax 後の attention weights、residual addition 前の各 sub-layer output、embedding sum に現れ、Chapter 6 が説明した通りのことをします。大規模な pretraining runs では、しばしばゼロに設定されます。各 token を一度しか見ない model は、overfit できる立場にないからです。

layer 内の2つの tensors は shape $n \times n$ を持ちます。ここで $n$ は tokens 数です。scores と softmax 後の weights です。それ以外、すべての projection、MLP 全体は $n$ に linear です。

attention layer 1つ、幅512、8 heads、batch 1、float32、laptop GPU 上です。2つの millisecond columns は 比率 だけを読んでください。これは熱くなると1,785 MHz から300 MHz 未満へ throttle する8 GB laptop card の wall clock なので、同じ code の cold run は7〜10倍速く返り、busy run はさらに遅くなります。megabyte columns は allocator byte counts で、動きません。

  tokens   ms total    ms x4   ms projections   attn matrix MB    peak MB    MB x4
     128      2.246        -            1.324              0.5       14.6        -
     256      2.855     1.27            2.113              2.0       19.2     1.31
     512      5.761     2.02            3.105              8.0       34.4     1.79
    1024     16.414     2.85            4.008             32.0       89.1     2.59
    2048     51.573     3.14            9.989            128.0      296.1     3.32
    4096    225.432     4.37           20.176            512.0     1100.1     3.72
    8192    832.838     3.69           40.106           2048.0     4300.1     3.91
   16384   OUT OF MEMORY                                 8192.0

fitted exponent (log-log slope, last four rows):  time ~ n^1.91   memory ~ n^1.87

x4 columns は上の row への ratio で、$n$ を倍にすると time と memory の両方で正確に4へ収束します。最後の step では理論値4に対して3.91です。projections column は control です。1024 tokens で4.0 ms、8192で40.1 ms。8倍に対して10倍です。宣言通り linear です。

そして最後の行です。attention layer 1つ、sequence 1つ、周囲に model なしで、16,384 tokens において8 GB GPU の memory が尽きます。score matrix だけで8 GB になるからです。8 heads × 16,384 × 16,384 × 4 bytes です。model ではありません。1つの layer の中間 tensor 1つです。

これは後の3章の下にある物理的事実です。context window にそもそも limit がある理由であり、Chapter 16 はそれを価格に変えます。FlashAttention が存在する理由でもあります。同じ結果を tiles で計算し、matrix を保存しないのです。それは speed optimization である前に memory optimization です。¹¹ そして long prompt の価格の背後にある arithmetic でもあり、Chapter 24 はそれを agent loop の中で支払います。その章の別の発見、つまり model は long context を 使う のも下手であるという測定とは別問題であり、その原因をこの式のせいにはしません。

2017年の paper は encoder-decoder を説明しています。source を unmasked attention で読む stack、target を causal に生成する2つ目の stack、そして中央に decoder の queries が encoder の keys と出会う3種類目の attention があります。これは入力と出力が2つの sequences である翻訳には正しい形です。

勝ったのは decoder-only の半分でした。1つの stack、全体が causal、入力と出力は同じ sequence です。その理由は elegance ではありません。「次の token を予測する」は どんな text でも動くので、training set は parallel corpus ではなく internet になり、すべてがその1つの task になります。翻訳は source と target を含む document、質問と答えも document、途中に tool call がある会話も document です。Chapter 11 はその最後のものがどう作られるかを扱います。Encoders は消えたわけではありません。encoder は入力全体を一度に見ます。これは job が text を続けるのではなく represent することである場合に欲しい性質であり、Chapter 19 の retrieval embeddings が、chatting している model ではなく encoders から来る理由です。

block が定義されると、model size は arithmetic になります。幅 $d$、4倍 expansion の block あたりでは、GPT-2 のように4つすべてに biases がある $W_q, W_k, W_v, W_o$ に対して $4d^2 + 4d$ です。上の table はそのうち3つの bias を外しているため、$d = 768$ で block あたり2,304少なくなっています。MLP に $8d^2 + 5d$、2つの LayerNorm に $4d$、つまり $12d^2 + 13d$。さらに $V \times d$ の token table、absolute positions なら $n_{\text{ctx}} \times d$ です。GPT-2 small の shape、つまり $d = 768$、12 blocks、vocabulary 50,257、context 1024、output layer が embedding weights を共有する場合はこうです。

  token embeddings     50,257 x 768 = 38,597,376
  position embeddings   1,024 x 768 =    786,432
  one block                             7,087,872
  12 blocks                            85,054,464
  final LayerNorm         2 x 768 =        1,536
  total (weights tied)                124,439,808

これはその model の公表サイズです。式は近似ではありません。それが model です。小さな model のほぼ3分の1が embedding table であることにも注目してください。だから vocabulary size は preprocessing ではなく architectural decision なのです。Chapter 7 が設定した trade-off です。

perplexity は corpus についての数字です。1つの head が何をしているかは別の問いであり、1 MB の Shakespeare で学習した model はそのための道具として不適切です。500,000-parameter model の attention map について正直に言えるのは、その大半は解釈不能だということです。そこで、正解のある言語を使います。

古典的な例は the animal did not cross the street because it was too tired です。このとき it は animal を指します。一方で …because it was too wet では、1語によって referent は street に移ります。これは Winograd schemas¹⁴ です。1語だけが違う文の pair で、その語が pronoun の指示対象を決めます。

ただしそれは cheating で解けます。tutorials が飛ばす部分です。2つの候補が animal と place なら、tired と wet は referent を category で識別し、どの words が存在するかだけを知る model でも、order を何も知らずに正解します。この version の task を、held-out animal/place pairs で測るとこうなります。

uniform causal average           held-out referent accuracy 100.0 %
one transformer block            held-out referent accuracy  91.7 %

bag of words が transformer に勝ちます。その文に基づく demonstration は、attention について何も証明しません。

そこで穴を塞ぎます。両方 の candidates を16 nouns の同じ pool から引き、どちらもどちらの slot にも現れ得るようにします。そして adjectives は category ではなく role で分けます。it を crosser にする4つ（tired, scared, slow, weak）と、crossed にする4つ（wet, wide, busy, steep）です。

the {x} did not cross the {y} because it was too {adj} , so the {ref} waited .

普通の next-token predictor として学習し、1 position、つまり so the の後の word を score します。held-out set は、逆順 が training にあった noun pairs から作ります。すると、2つの nouns が存在することは知っているがどちらが先に来たかを知らないものは、逆向きに答えなければなりません。

存在する2つの nouns の間での chance は50 %です。一様平均は27.9 %に着地し、pair の間違った noun をちょうど半分の確率で 答えます。これは3 sections 前の shuffle test が予測した通り、どの words があるかは知っているが、その order は何も知らないものの signature です。

次に map です。referent を名指さなければならない position における attention を、各 block の4 heads で平均し、1語だけ違う2文について示します。一様平均なら、見えている15 tokens のそれぞれに0.067を置くはずです。

the animal did not cross the street because it was too tired , so the animal waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00 tired:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.00
         because:0.00 it:0.00 was:0.00 too:0.00 tired:1.00 ,:0.00 so:0.00 the:0.00

the animal did not cross the street because it was too wet , so the street waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00   wet:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.03 the:0.00 street:0.49
         because:0.00 it:0.00 was:0.00 too:0.20   wet:0.03 ,:0.00 so:0.00 the:0.25

Block 1 は 両文で同一 です。adjective が何であれ、first noun に0.70を置きます。これは失敗ではなく証明です。first layer では、ある position の query はその position 自身の token と index の関数であり、position 14 の the は両文で同じ token です。first-layer head は、まだ fetch していない word に condition できません。 だから block 1 は利用可能な唯一の有用なことを行い、first noun を前へ引っ張ります。

Block 2 で文が分岐します。8つすべての adjectives にわたって同じ row が、model の見つけた rule を示します。

crosser-adjective では、second block は weight 全体を adjective に費やします。答えはすでに residual stream にあるからです。block 1 がそこへ置いたので、必要なのは確認だけです。crossed-adjective では代わりにもう一方の noun を fetch しに行きます。これは two-hop circuit です。1つの head が candidate を前へ動かし、後の layer の head が、それを保つべきかを決める token を読みます。layers across の composition が mechanism であり、one block が92.7 %、two blocks が100 %に達した理由です。

これはまた、実 model で最もよく記録された circuit の形でもあります。Induction heads、つまり pattern [A][B] … [A] → [B] を完成させる、次の layer の head に feeding する previous-token head は、Anthropic の interpretability work が in-context learning の大部分の背後に特定したもので、pretraining 中の識別可能な時点で形成されます。この章はその analysis を試みません。両論文を references に置いて委ねます。real model から circuits を読み出すことは研究分野であり、section ではないからです。

最後に implementation です。上の30行に、PyTorch 自身の weights をコピーしたものです。

ours vs nn.MultiheadAttention           max |diff| = 1.7881393432617188e-07
ours vs F.scaled_dot_product_attention  max |diff| = 1.7881393432617188e-07

平均 magnitude 0.159の outputs に対して $1.8 \times 10^{-7}$。float32 precision における、順序を変えた同じ arithmetic です。

この course の残りに出てくるすべての model が基にする architecture を手にしました。そしてそれは評判より小さいものです。重みが learned である weighted average、parameters の3分の2を保持する per-position MLP、2つの normalisations と2つの additions、それを stack したものです。

まだ持っていないのは、何かを知っている model です。そして stack するだけではそれは直りません。この corpus では2 blocks が training perplexity 14.49、validation perplexity 40.57に達し、one block の18.77と38.07に対して、見たものでは良く、見ていないものでは悪くなります。これは Chapter 6 の table に transformer が入ったものです。この model と Chapters 14 to 30 が会話する model との距離は architectural ではありません。同じ block が、より多く繰り返され、はるかに多くの text にかけられているだけです。

つまりそれは accounting problem になります。そして accounting は見た目より奇妙です。どれほどの text が必要で、それをどこで手に入れるのでしょうか。どれほどの arithmetic が必要で、お金を使う 前に どう見積もるのでしょうか。固定 budget があるとき、model を大きくする方がよいのか、それとも data を多く見せる方がよいのか。そして正しい答えはあるのでしょうか、それとも流行だけでしょうか。Chapter 10 は3つすべてに測定で答え、この問いの最も安く有用な形に価格を付けます。今日、GPT-2 のような model をゼロから train するには、いくらかかるのでしょうか。

この題材についての3つの説明は、それぞれの目的においてこの章より優れています。この章はそれらと併読されるように書かれています。Jay Alammar の The Illustrated Transformer は、data flow の最良の図解です。Harvard NLP の The Annotated Transformer は、2017年の paper に実行 code を1行ずつ挟んだものです。Andrej Karpathy の Let's build GPT: from scratch, in code, spelled out は、同じ model を2時間で live に構築します。上の ablations の ladder は、別の corpus で測った同じ backbone です。この章が触れるだけにした interpretability の問いについては、primary sources は Anthropic の interpretability group による Elhage et al., A Mathematical Framework for Transformer Circuits (2021) と Olsson et al., In-context Learning and Induction Heads (2022) です。