デコードの各ステップでKVキャッシュ全体をGPU HBMから読み込む。128Kトークンでは数GB、1Mトークンでは60GBを超える。ボトルネックはメモリ帯域幅であって計算量ではない。GPUはクエリに無関係なデータの転送を待たされる。

共有基底スパース注意機構が構造的に破綻している理由

Quest、ShadowKV、RocketKVは、どのKVページが重要かを推定してそのページだけを読み込むことでこのコストを削減する。それらは現在のクエリを、シーケンス全体から構築した共有低ランク基底に射影してページをスコアリングする。根本的な問題は、固定された基底Wには必ず零空間が存在することである。キーの構造がその零空間に存在するページは、クエリとの実際の関連性に関わらずスコアがほぼゼロになる。訓練手順ではこれを修正できない——これは表現上の不可能性である。

LOCKSの命題1はこれを形式化する：すべての固定共有射影は「ページ内容盲目の方向」を持つ。ShadowKVが小さなトークンバジェットで崩壊するのはチューニングの問題ではなく、この構造的破綻が現れた結果である。

LOCKSが基盤とする観察

16連続トークンのページ内では、キー・ベクトルは低次元部分空間に集中する。共起するトークンは類似した注意パターンを必要とする傾向があり、それが類似したキー・ベクトルを生み、d_head次元空間の小さな領域にクラスタリングされる。

シーケンス全体では、異なるページが異なる部分空間を占める——コードのページ、散文のページ、引用のページはそれぞれ異なるキー部分空間の向きを持つ。グローバルなキー行列は高ランクだが、各ページは局所的に低ランクである。

共有基底はこれらすべてを単一の方向セットで同時にカバーしなければならないが、それは不可能である。ページごとの基底であれば、各々が自分のページだけをカバーすればよい。

LOCKSの仕組み

ページが埋まると（16トークンごと）、LOCKSは中心化されたキー偏差に対して一度だけSVDを実行し、以下を保存する：

  • μⱼ (int8): ページの重心——16トークンにわたる平均キー・ベクトル
  • Vⱼ (int4, rank 8): 偏差行列の上位8右特異ベクトル——このページ内での変動の主方向
  • Cⱼ (int8): Vⱼへの各キーの射影係数

合計：元のページKVデータの約10%。これはデコードのクリティカルパス外で実行される。

各デコードステップで、LOCKSはこれらの要約情報のみからすべてのページをスコアリングする——KまたはVテンソルを読み込まない：

$$\hat{s}j(q) = \log \sum{i \in P_j} \exp\left[\frac{q^\top \mu_j + (V_j^\top q)^\top c_i}{\sqrt{d}}\right]$$

q^T μⱼ：クエリがページの平均キー方向とどれだけ整合するか（ページごとに1スカラー）。
(Vⱼ^T q)^T cᵢ：クエリが各キーの重心からの偏差とどれだけ整合するか、ページの局所基底で測定（トークンごとに1スカラー）。

合計してlog-sum-expを適用すると、ページの推定総注意量が得られる。このスコアで上位のページが完全な注意を受け、残りはスキップされる。

定理1は、非量子化ランクr要約が、SVD再構成誤差をηで抑えたとき、少なくともe^(−4η)の厳密質量選択カバレッジを達成することを証明する。補題1は、質量ランキングがすべての決定論的値盲目的選択器の中で最悪ケース最適戦略であることを証明する。

コード 

import torch
import torch.nn.functional as F

SEQ_LEN = 65536
D_HEAD  = 128
PAGE_SIZE = 16
RANK = 8
TOKEN_BUDGET = 2048

K = torch.randn(SEQ_LEN, D_HEAD)
V = torch.randn(SEQ_LEN, D_HEAD)
q = torch.randn(D_HEAD)

# Build page summaries once (off critical path)
summaries = []
for p in range(SEQ_LEN // PAGE_SIZE):
    keys   = K[p*PAGE_SIZE:(p+1)*PAGE_SIZE]
    mu     = keys.mean(0)
    dev    = keys - mu.unsqueeze(0)
    _, _, Vh = torch.linalg.svd(dev, full_matrices=False)
    basis  = Vh[:RANK].T
    coeffs = dev @ basis
    summaries.append({"mu": mu, "V": basis, "C: coeffs})

# Per-step: score all pages from summaries only
scale = D_HEAD ** -0.5
scores = []
for s in summaries:
    global_t = (q @ s["mu"]) * scale
    local_t  = (s["C"] @ (s["V"].T @ q)) * scale
    scores.append(torch.logsumexp(global_t + local_t, dim=0))
page_scores = torch.stack(scores)

# Select top pages and attend
n_top = TOKEN_BUDGET // PAGE_SIZE
top_pages = torch.topk(page_scores, n_top).indices.sort().values
sel_idx = torch.cat([
    torch.arange(p*PAGE_SIZE, (p+1)*PAGE_SIZE) for p in top_pages
])

out_locks = F.scaled_dot_product_attention(
    q.view(1, 1, D_HEAD),
    K[sel_idx].view(1, len(sel_idx), D_HEAD),
    V[sel_idx].view(1, len(sel_idx), D_HEAD),
).squeeze()

out_full = F.scaled_dot_product_attention(
    q.view(1, 1, D_HEAD),
    K.view(1, SEQ_LEN, D_HEAD),
    V.view(1, SEQ_LEN, D_HEAD),
).squeeze()

cos_sim = F.cosine_similarity(out_locks.unsqueeze(0), out_full.unsqueeze(0)).item()
print(f"Attended: {len(sel_idx)}/{SEQ_LEN} ({100*len(sel_idx)/SEQ_LEN:.1f}%)")
print(f"Cosine similarity to FullKV: {cos_sim:.4f}")
# Expected: ~0.98+ at 3% token attendance

Enter fullscreen mode Exit fullscreen mode

結果

InfiniteBench (GLM-4-9B-Chat-1M, 100K+ context):

Method b=512 b=2048
Quest 34.7 35.7
ShadowKV 36.0 40.3
RocketKV 38.5 42.3
LOCKS 41.1 43.6
FullKV 43.0 43.0

b=2048では、LOCKS（43.6）がFullKV（43.0）を上回る。選択されたトークン部分集合だけに注意を向けることで、注意分布を希釈していた無関係なトークンを除外できる。

LongBench-v1 (Llama-3.1-8B): b=64からb=2048までFullKVから約1ポイント以内に収まる。競合手法はb=512以下で崩壊する。

H200デコードレイテンシ：128Kで1.59倍、256Kで1.80倍、1Mで2.0倍。ステップあたりの読み込みバイト数：1Mで9.8倍削減。デコードはこれらのコンテキスト長ではほぼ完全に帯域幅に律速されるため、帯域幅節約はレイテンシ節約にほぼ線形に反映される。

注意点

定理1は非量子化要約に適用される。提供されているint4設定は経験的に検証されたものであり、ステップごとの保証はない。アブレーションではint4でrecallが数ポイント低下し、int2は崩壊することが示されている。実運用前に検証すること。

LOCKSは帯域幅を節約するが、VRAMは節約しない。フルKVキャッシュは常駐したままになる。制約がメモリ容量であって帯域幅ではない場合、これは役に立たない。

単一著者によるプレプリントであり、独立した再現はまだ行われていない。256K以上のコンテキストで最も効果を発揮する——128Kでの1.59倍の利得は統合オーバーヘッドを正当化しない可能性がある。

References