OCTOPUS：基于八面体参数化与最优平方误差量化的 Transformer KV 缓存优化 OCTOPUS: Optimized KV Cache for Transformers via Octahedral Parametrization Under optimal Squared error quantization
用八面体映射对旋转后的三维坐标三元组做联合量化，低比特下显著超越现有 KV 缓存压缩方法。
前置知识
KV 缓存 (KV Cache)
自回归 Transformer 解码时，注意力需要访问所有历史 token 的键 K 和值 V 张量，这些历史 K、V 的缓存即为 KV cache。长上下文下每步都要从显存读取它，成为延迟和批大小的主瓶颈。
OCTOPUS 的全部动机都建立在 KV cache 主导显存带宽这一事实上。
Walsh-Hadamard 变换 (WHT)
一种快速正交变换，可通过蝴蝶网络在 $O(d\log d)$ 内原地完成。配合随机符号翻转后作为旋转算子 $R=H\,\mathrm{diag}(s)$，能把方向向量各坐标变成解析已知的 isotropic 分布。
OCTOPUS 及所有旋转预条件 codec 都靠它把坐标边缘变成可解析的形式。
Lloyd-Max 量化器
给定信源边缘分布与比特数 $b$ 时最小化 MSE 的最优 1-D 标量量化器。通过最近质心分配与质心更新交替迭代训练，码本只依赖 $(d,b)$，无需数据校准。
OCTOPUS 的两个码本 $C_\rho$、$C_\xi$ 都是 Lloyd-Max 码本，是其 MSE 最优性的根基。
旋转预条件量化
先用数据无关的随机正交旋转 $R$ 预处理向量，使每坐标边缘变为已知分布（如对称 Beta），再逐坐标独立做 Lloyd-Max 量化。代表方法 TurboQuant、PolarQuant、QJL。
OCTOPUS 属于这一族，但其核心改造是改为对三元组联合量化。
八面体映射 (Octahedral map)
来自实时渲染的等面积参数化，把单位球面 $S^2$ 上的方向映射到正方形 $[-1,1]^2$ 上的两个标量 $(\xi,\eta)$，encode/decode 为分段线性，每个卦限内雅可比恒定。
它是 OCTOPUS 把三元组方向压成两个标量、用 1-D Lloyd-Max 近似 2-sphere 畸变的关键工具。
QJL 残差 (1-bit Johnson-Lindenstrauss sketch)
用 1 比特 JL 投影给量化残差 $r=u-\hat{u}$ 做草图，存 $\sigma=\mathrm{sign}(R'r)$，得到点积的无偏估计 $\sqrt{\pi/(2d)}\,\gamma_r(R'q)^\top\sigma$，几乎零成本纠正点积偏置。
MSE 最优的标量量化在点积上有偏，QJL 残差以每坐标 1 比特把它压到零。
研究动机
长上下文自回归推理（大语言模型、因果视频生成、音频生成）中，键值（KV）缓存在每个解码步都要从高带宽显存读取，主导了延迟与批大小，因此 KV 压缩是首要优化目标。现有方法分三类：token 驱逐、带残差的逐通道标量量化、以及旋转预条件量化（TurboQuant、PolarQuant、QJL）。后三者虽通过结构化随机旋转把每坐标边缘变成解析已知分布、再用匹配的 Lloyd-Max 量化器逼近最优，但它们都只对单个坐标或单个角度独立量化。在极端压缩下它们会灾难性崩溃：在 Qwen2.5-7B-Instruct-1M 的 WikiText-2 上，2 比特时 TurboQuant-QJL 困惑度退化 +772%，PolarQuant +186.6%，TurboQuant-MSE +63.0%；在 Causal Forcing 视频上 TurboQuant-QJL 的最差 prompt LPIPS 接近 1.0，几乎变成随机噪声。
本文的目标是本文要设计一种旋转预条件的 KV 编解码器，对旋转后的坐标三元组做联合量化，而非逐坐标独立量化，从而在相同比特率下取得更低的 MSE。它必须保持数据无关、在线、给定种子确定性可复现，且码本只依赖 $(d,b_{\mathrm{dir}},b_{\mathrm{nrm}})$ 而无需逐向量缩放或数据校准。同时要能以融合 Triton 核在注意力内部即时重建键、不物化完整 $\hat{K}$ 张量，从而不增加解码期带宽或延迟。作者还要证明该构造不局限于语言模型，任何带注意力的自回归 Transformer（视频、音频）都能受益，并在 b=2 这种压缩质量最关键的区间显著超越所有同类方法。
与已有工作不同的是，作者的独特切入点是把计算机图形学中实时渲染常用的八面体映射首次引入 Transformer 解码，作为 $S^2$ 的等面积参数化，把单位三元组方向压成 $[-1,1]^2$ 上的两个标量，使 1-D Lloyd-Max 能很好地近似 2-sphere 畸变。第二个关键洞察是：旋转预处理把熵均匀铺到各坐标后，子块范数携带的熵随维度 $d$ 升高而渐近消失——方向误差 $\sigma_n^2=O(1)$ 而范数方差 $\sigma_\rho^2=O(d^{-1})$，因此方向应分到更多比特。据此在每三元组平方误差上做拉格朗日优化，得到与 $d$ 和总预算都无关、只差一个 $O(1)$ 常数的非均匀比特分配，在 $d=128$ 恰好对应 $(b+1,b-1)$ 的实现方案。
核心方法
直觉上：旋转预条件后各坐标各向同性，范数熵很少、方向熵多，把三元组拆成范数 $\rho_i$ 与单位方向 $n_i\in S^2$ 分别量化、并给方向多分比特，就能击败逐坐标量化。技术路线是九阶段流水线：键 $k$ 归一化为模 $\gamma$ 与方向 $\tilde{u}$；符号翻转 WHT 旋转 $u=R\tilde{u}$；切成 $\lceil d/3\rceil$ 个连续三元组；拆成范数与方向；八面体映射得 $(\xi_i,\eta_i)$；用 Lloyd-Max 码本配合 MSE 最优的非均匀比特分配量化；局部 $3\times3$ 联合取整；输出 $S(k)=(\gamma,I_{\mathrm{dir}},I_{\mathrm{nrm}})$；可选附加 QJL 残差。解码用 split-K flash 核按 $\gamma\sum_i\hat{\rho}_i q_{\mathrm{rot},i}^\top\hat{n}_i$ 累加，从不物化 $\hat{K}$。
核心创新有三。其一，八面体三元组方向量化器把方向编码为两个 $[-1,1]$ 标量，配合与实现匹配的范数/方向边缘分布，是首个用于 KV cache 的方向量化原语。其二，MSE 最优的非均匀比特分配：在范数-方向耦合的平方误差 $\mathbb{E}\|t_i-\hat{t}_i\|^2\propto 2C_\rho\sigma_\rho^2\,4^{-b_{\mathrm{nrm}}}+(6/d)C_n\sigma_n^2\,4^{-b_{\mathrm{dir}}}$ 上做拉格朗日优化，得到 $b^*_{\mathrm{dir}}-b^*_{\mathrm{nrm}}=\log_4\!\big(3C_n\sigma_n^2/(2d\,C_\rho\sigma_\rho^2)\big)$，代入已知方差后差距是 $O(1)$，$d=128$ 时正是 $(b+1,b-1)$。其三，可选的 1 比特 QJL 残差把点积偏置压到零。与逐坐标量化的本质区别在于：它把旋转熵的结构（子块范数熵递减）显式建模成分离的范数+方向两路比特，并解析地给出最优分配。
方法步骤详情
对键 $k\in\mathbb{R}^d$：(1) 归一化 $\gamma=\|k\|_2$、$\tilde{u}=k/\gamma$（$\gamma$ 存 fp32）。(2) 旋转 $u=R\tilde{u}$，$R=H\,\mathrm{diag}(s)$ 每 head 抽一次。(3) 切 $\lceil d/3\rceil$ 个连续三元组 $t_i$。(4) 拆 $\rho_i=\|t_i\|_2$ 与 $n_i\in S^2$。(5) 八面体映射 $(\xi,\eta)=\mathrm{Oct}(n_i)$。(6) 用训练到 $10^{-10}$ 畸变的码本 $C_\xi$、$C_\rho$ 索引。(7) 局部 $3\times3$ 联合取整，先定方向再选最近 $\rho$ 质心。(8) 输出 $S(k)=(\gamma,I_{\mathrm{dir}},I_{\mathrm{nrm}})$。(9) 可选 QJL 存残差符号 $\sigma=\mathrm{sign}(R'r)$。解码用 split-K flash 核按三元组累加点积，$\hat{K}$ 不物化。
技术新颖性
新颖性体现在四方面。一是把图形学的八面体映射首次用作 Transformer 解码的方向量化器：它对 $S^2$ 等面积、encode/decode 为 $O(1)$ 分段线性、每卦限雅可比恒定，使 1-D Lloyd-Max 能紧贴真实 2-sphere 畸变，而递归极坐标（PolarQuant）需超越函数且诱导麻烦的角度边缘。二是推导出解析可解的 MSE 最优非均匀比特分配，且最优差距与 $d$、与总预算都无关，仅 $O(1)$ 常数——实验上 $d=128$ 各比特下严格落在 $\delta=+1$，相对均匀 $(b,b)$ 降 31–41% MSE。三是局部 $3\times3$ 联合取整：只在 Lloyd 标量种子的相邻九格里搜，在 $10^4$ 个随机三元组上与全网格搜索逐字节相同，几乎零成本拿到联合最优。四是融合 Triton split-K flash 解码，使编解码全程不物化 $\hat{K}$，QJL 残差可作为侧车挂接而不改解码器。
实验结果
四模态规律一致：OCTOPUS 在低比特匹配或击败所有旋转基线，比特越低优势越大。合成实验（$d=128$ 高斯键）：4 比特 MSE 比 TurboQuant-MSE 低 1.3 倍，2 比特比 PolarQuant 低 2.4 倍；OCTOPUS-QJL 的 IP 绝对误差比 TurboQuant-QJL 低 3 倍。长上下文 LM：2 比特 WikiText-2 仅 +34.7%，而 TurboQuant-MSE/PolarQuant/TurboQuant-QJL 分别 +63.0%/+186.6%/+772.0%；多键 NIAH 2 比特时仅 OCTOPUS（0.81）与 OCTOPUS-QJL（0.83）保住召回，基线崩到 0.04/0.01。视频/音频 2 比特分化剧烈：Causal Forcing 上 TurboQuant-QJL 最差 prompt LPIPS=0.997，OCTOPUS 稳在 0.58/0.82；音频基线 LSD 升到 12.6–13.2 dB 且 SNR 为负，OCTOPUS 仅 6.75 dB LSD、+1.07 dB SNR。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 合成高斯键重构 ($d=128$, 2-bit) | per-coord MSE (↓) | OCTOPUS 0.0897 | TurboQuant-MSE 0.1161 / PolarQuant 0.2197 | 比 TurboQuant-MSE 低 1.3×，比 PolarQuant 低 2.4× |
| Qwen2.5-7B-Instruct-1M WikiText-2 (2-bit) | 困惑度退化 Δ% | +34.7% | TurboQuant-MSE +63.0% / PolarQuant +186.6% / TurboQuant-QJL +772.0% | 退化幅度约为最强基线 TurboQuant-MSE 的一半 |
| 多键 needle-in-a-haystack (4k–128k, 2-bit) | 召回 (↑) | OCTOPUS 0.81 / OCTOPUS-QJL 0.83 | PolarQuant 0.04 / TurboQuant-QJL 0.01 | 保留可用召回，基线近乎全崩 |
| Causal Forcing 视频 (2-bit, 最差 prompt) | LPIPS (↓) | 0.82 | TurboQuant-QJL 0.997 | 最差情况仍连贯，基线退化为噪声 |
| AAR 音频 (2-bit) | mean LSD (↓, dB) | 6.75 | TurboQuant-MSE/QJL/PolarQuant 12.6–13.2 | LSD 约为一半，SNR 由负转正 (+1.07 dB) |
局限与改进
作者明确承认：更优的率-质量点并非没有代价——OCTOPUS 解码比标量 Lloyd-Max 多做不少算术，仍慢于 bf16 的 SDPA 路径，因此只在 KV 带宽或容量成为瓶颈时才最有吸引力。QJL 残差只买点积精度、不改善重建路径质量，只适合 score-attention 部署场景。此外，OCTOPUS 在某些高比特区间并非全面领先：4 比特视频上 PolarQuant 在 3% 以内，AAR 在 3 比特均值 LSD/SNR 上 PolarQuant 略好。我的观察：八面体映射和 $3\times3$ 联合取整主要在 $S^2$ 上生效，把三元组推广到更高维子块会改变范数分布与雅可比结构，比特分配公式需重推；且 $\gamma$ 全程以 fp32 存储，在超长上下文（如 1M token）下这一项虽只占 0.25 bpc，但仍是固定开销。
独立分析的弱点
弱点一：wall-clock 开销未被量化到与 bf16 路径的直接对比表，附录虽有延迟数据但正文缺乏端到端吞吐数字，难以判断带宽节省换算的实际收益；改进方向是给出真实服务负载下的延迟/吞吐 Pareto。弱点二：非均匀 $(b+1,b-1)$ 分配是针对 $d=128$ 的离散结论，$d=64$ 的视频键和音频键是否仍最优未在正文验证；改进方向是把比特分配做成按 $d$ 自适应的查表。弱点三：合成实验假设各向同性高斯先验，真实 key 的异向性可能让八面体方向边缘偏离训练码本；改进方向是在线轻量校准码本。弱点四：$\gamma$ 固定 fp32、QJL 残差每坐标 1 比特，在 b=2 这种极紧预算下相对开销偏高；改进方向是对 $\gamma$ 做更低位宽或分组共享。
未来方向
作者提出的方向：把 OCTOPUS 推广到任何带注意力的自回归模态（已验证文本/视频/音频），并与 token 驱逐、稀疏编码等互补技术组合。基于本成果可延伸的方向：其一，把三元组推广为更大的子块（如 4 元或 8 元），用四元数或更高维等面积参数化替代八面体映射，可能进一步压低方向熵，但需重推比特分配的拉格朗日解。其二，把局部 $3\times3$ 联合取整扩展为跨三元组的联合优化或学习型取整网络，进一步逼近全局最优。其三，当前码本离线训练且数据无关，可探索少量在线校准以适配真实 key 的异向性。其四，将融合 Triton 解码核与 PagedAttention、speculative decoding 等推理系统深度集成，量化真实服务下的端到端收益。其五，把方向+范数分离量化的思想迁移到权重或激活量化。
复现评估
复现性较好但有缺口。有利因素：码本数据无关、只依赖 $(d,b_{\mathrm{dir}},b_{\mathrm{nrm}})$，作者声明给定种子确定性可复现，编码与 split-K flash 解码核均以融合 Triton（附录给出 Algorithm 1/2）实现；实验在公开模型上做——Qwen2.5-7B-Instruct-1M、Wan-1.3B（CausVid、Causal Forcing）、AAR，数据集为 WikiText-2、C4、AudioSet-20k、needle-in-a-haystack 合成集，均公开可得；比特分配有解析推导（式 10），码本训练到 $10^{-10}$ 畸变且仅约 160 字节。缺口：正文未给出 Triton 核源码或码本文件的公开链接，关键超参（V 侧分组、residual window、boundary-1 配方）散落文中需仔细拼装；八面体方向边缘的实证码本训练细节主要在附录；算力门槛较高——需多卡跑 7B 长上下文与视频 DiT。项目页 octopus-quant.github.io 给出部分材料，但完整代码未见。
论文图表
