アテンションの計算はそのまま——ただ一度も書き出さない。 アテンションの本当のコストは掛け算ではなく、巨大なスコアの表を遅いメモリへ出し入れすることだ。Flash attentionは計算を一切変えず、それがどこで起きるかをすべて変える。答えは同じ、時間はわずか。キューブを高速で解くように:同じパズル、同じ手順——勝負を分けるのは純粋な方法だ。
ボトルネックは計算ではない。運搬だ。 標準的なアテンションは、すべての単語を他のすべてと突き合わせた表をまるごと作り——その巨大な格子を遅いメモリに書き込み、正規化のために読み戻し、値に重みを付けるためにもう一度読む。モデル最大のものを抱えての三往復、その間チップは計算を終えてただ待つ。広げた絨毯のように:部屋には大きすぎるので廊下に置かれ——一目見るたびにまた一往復だ。
速いチップが、遅い管に飢える。 T≈max (FLOPscompute rate, bytes movedbandwidth)T \approx \max\!\left(\dfrac{\text{FLOPs}}{\text{compute rate}},\ \dfrac{\text{bytes moved}}{\text{bandwidth}}\right) プロセッサは、噛み砕く数を取りに行くよりもはるかに速く計算をこなす。だから全体の時間は、計算と、バイトの移動の、遅いほうで決まる。アテンションはメモリ律速——計算は準備できているのにデータが来ない——だから速くするのは計算を削ることではなく往復を削ることだ。細いホースにつないだ強力なポンプのように:大海も動かせるのに、流れを決めるのは管のほうだ。
ならば全部は載せない。一度に一ブロックずつ進める。 Sij=QiKj⊤d,Qi, Kj∈RB×dS_{ij} = \dfrac{Q_i K_j^{\top}}{\sqrt{d}}, \qquad Q_i,\, K_j \in \mathbb{R}^{B\times d} 解決は単純に始まる。格子全体を手元に持ってこない。Keys と Values を、チップの小さく速いメモ帳に収まるほど小さなブロックに切り分ける。一ブロックを取り込み、その分の仕事をして、手放し、次を取りに行く。平たく言えば、B 行のブロックが速いメモリに収まるので、一度に一組のブロックだけを扱う。薪を運ぶように:薪の山ごと炉まで引きずらない——一抱えずつ運ぶのだ。
だが softmax は行全体を一度に必要とする。 softmax(x)i=e xi−m∑j=1ne xj−m,m=maxjxj\mathrm{softmax}(x)_i = \dfrac{e^{\,x_i - m}}{\sum_{j=1}^{n} e^{\,x_j - m}}, \qquad m = \max_{j} x_j ここに落とし穴がある。スコアを重みに変えるため、softmax はそれぞれをすべての key にわたる合計で割る——そして安全のため、まず行の最大スコアを引き、指数が決して暴走しないようにする。最大値も合計も、行全体を必要とする。だが手元にあるのは一ブロックだけ。まだ見ていない数に対して、どう正規化するのか。一つの水差しから公平に注ぐように:いくつあるか分からないうちは、最初のグラスも満たせない。
妙技は、走る最大値を持ち、進みながら基準を取り直すこと。 mnew=max(mold, m~)ℓnew=e mold−mnew ℓold+∑ie xi−mnew\begin{aligned} m_{\text{new}} &= \max(m_{\text{old}},\, \tilde m) \\ \ell_{\text{new}} &= e^{\,m_{\text{old}} - m_{\text{new}}}\,\ell_{\text{old}} + \sum_{i} e^{\,x_i - m_{\text{new}}} \end{aligned} 最初に行全体は要らない——必要なのは走る集計だ。これまでに見た最大スコアと、走る合計を持つ。新しいブロックがより大きなスコアを持ち込んだら、一つの縮小係数で古い合計を新しい尺度に移し、そこへブロックを織り込む。出力も同じやり方で取り直す。結果は一度に全部やるのとまったく同じだ。走り高跳びのバーのように:より良い跳び手が来たらバーを上げ、これまでの記録をすべて新しい高さで測り直す。
答えは同じ。巨大な格子は存在すらしなかった。 memory: O(n2) ⟶ O(n),oflash=oexact\text{memory: } O(n^2)\ \longrightarrow\ O(n), \qquad \mathbf{o}_{\text{flash}} = \mathbf{o}_{\text{exact}} すべてを合わせると、n×n の表は一度も保存されない——走る最大値、走る合計、そして出力、どれも小さいものだけだ。メモリは二乗から線形へと下がり、速くなるのは計算を削ったからではなく、遅いメモリへの往復が桁違いに減るからだ。しかもこれは近似ではない。結果は厳密なアテンションそのものだ。折り鶴を折るように:一枚の紙、仕上がった形、机に切れ端は残らない。
🌱 遅かったのは、思考そのものではなかったのかもしれない。 Flash attention は新しい考えを何も足さない——ただ計算がどこに宿るかを尊重し、同じ思考を、それを収める体に合わせて並べ直すだけだ。では、心の速さのどれだけが思考そのもので——どれだけが、すでに知っていることへ手を伸ばす速さにすぎないのだろう?
