KV キャッシュの計算量と記憶量

1 ヘッドの層（$W_Q, W_K, W_V, W_O$ は $d \times d$、MLP は $d \to d_{\mathrm{ff}} \to d$）で、スカラーの乗算の回数を数えます（層正規化、活性化関数、softmax の指数は数えません）。$n$ 個の位置 $t = 1, \ldots, n$ を順に処理するとします。

1. キャッシュつきでは、位置 $t$ の 1 層の乗算は $2td + 4d^2 + 2d\, d_{\mathrm{ff}}$ 回で、$n$ 個の合計は

です。 2. キャッシュなしで、位置 $t$ のたびに長さ $t$ の列全体を計算し直すと、1 層の乗算は $2t^2 d + t(4d^2 + 2d\, d_{\mathrm{ff}})$ 回で、合計は

です。 3. $d_{\mathrm{ff}} = 4d$ のとき、$C_{\mathrm{cache}}(n) = d\,n(n+1) + 12 n d^2 = O(n^2 d + n d^2)$、$C_{\mathrm{full}}(n) = \frac{1}{3} d\,n(n+1)(2n+1) + 6 d^2 n(n+1) = O(n^3 d + n^2 d^2)$ です。 4. キャッシュが保存する数の個数は、$L$ 層・$n$ 位置で $2 L n d$ です。