# YOCO：You Only Cache Once 中文精读笔记

- **论文**：You Only Cache Once: Decoder-Decoder Architectures for Language Models。
- **作者**：Yutao Sun、Li Dong、Yi Zhu、Shaohan Huang、Wenhui Wang、Shuming Ma、Quanlu Zhang、Jianyong Wang、Furu Wei；Microsoft Research / Tsinghua University。
- **版本**：[arXiv:2405.05254v2](https://arxiv.org/abs/2405.05254v2)，2024-05-09；获取与研读日期：2026-09-20。
- **材料**：[同版本原文 PDF](../source/paper.pdf)；作者 `main.tex`、引用的 `settings.tex` 与符号宏、原始图件。已阅读正文 §§1–5 和附录 A–G，并用 PDF 核对编号、公式及结果。PDF 共 20 页，以下页码均为从 1 起计的 PDF 页码，与印刷页码一致。未执行作者源码，也未复现实验。

## 论文总览

YOCO 将语言模型分成前后两个因果解码器：前半段 self-decoder 用高效序列模块生成上下文化表示，再投影出**唯一一份全局 KV 缓存**；后半段 cross-decoder 的所有层共享这份缓存，各层只更新自己的 query 和隐状态。这样仍可对全部历史 token 做精确的 softmax attention，却不必给每一层保存独立的全长 KV。

这里的“once”指全局 KV 只保存一份，不是整个模型没有其他状态，也不是生成时只计算一次 attention。self-decoder 仍维护各层的固定大小递推状态或滑窗缓存；每个新 token 都要经过两部分，并把新 KV 追加到共享缓存。

最有力的贡献是计算依赖关系的改变：**历史 token 的 cross-decoder 输出不再是未来生成所需的缓存**，所以构建提示词缓存时可跳过这些输出。默认 3B 配置下，作者报告 1M 长度总推理显存约 12.4 GB，相对 Transformer 降低约 9.4 倍；512K 情境下预填充从约 180 秒降到不足 6 秒。与此同时，3B 下游任务与 160M–13B 的固定 token 预算实验显示质量具有竞争力。这些结果分别来自不同实验，不能合并成“任意规模、任意任务均更快更准”。

## 结构地图

| 原文章节 | 解决的问题 | PDF 定位 |
|---|---|---|
| 1 Introduction | KV 内存和长提示词延迟瓶颈 | p.2 |
| 2 You Only Cache Once (YOCO) | 双解码器、共享缓存、推理依赖 | pp.3–5；图2–3、表1–2、式1–3 |
| 3 Design Choices of Self-Decoder | gRet 的三种形式与滑窗替代 | pp.5–6；式4–8 |
| 4.1 Language Modeling Evaluation | 3B、万亿 token 训练后的质量 | p.7；表3 |
| 4.2 Scalability Compared with Transformers | 160M–13B 扩展曲线 | pp.7–8；图4 |
| 4.3 Long-Context Evaluation | 单针、多针和长序列 NLL | pp.8–9；图5–6、表4 |
| 4.4 Inference Advantages | 显存、容量、prefill、吞吐 | pp.9–11；图7–10 |
| 5 Conclusion | 总结与尚未验证的应用设想 | pp.11–12 |
| A Chunk Parallelism for Long-Sequence Training of YOCO | 分块通信 | p.16；图11 |
| B Chunk-wise Representation of Gated Retention | 等价展开 | pp.16–17；式9–11 |
| C / D / E Hyperparameters… | 3B / 扩展曲线 / 长度扩展配置 | pp.17–18；表5–7 |
| F Pseudo Code of Gated Retention | 三种计算形式的伪代码 | pp.18–19 |
| G Comparisons with Transformer Variants | 联想召回与长任务诊断 | pp.19–20；表8、图12 |

## 一：动机——把逐层存储的历史改成共享记忆

常规自回归 Transformer 缓存每层历史 token 的 K、V，使下一步不用重算前缀。但随上下文长度 $N$ 和层数 $L$ 增长，缓存近似按 $LN$ 增长。GQA 降低 KV 头数，却没有移除层数因子。长提示词还需要在许多层中进行大规模因果 attention，导致很长的首次响应等待。

纯递推模型能够用固定大小状态压缩历史，但精确记住具体历史细节并非无代价。YOCO 的取舍是让前半部分使用高效递推或局部处理，同时为后半部分保留逐 token 的全局可寻址记忆。因此它不是把整个上下文都压成一个状态，也不是普通 encoder-decoder：两个部分都施加因果约束，输入和生成 token 都使用整套参数。（§§1–2，pp.2–4）

## 二：方法——先建一次记忆，再多层查询

### 2.1 数据流与基本块

![原图2：YOCO 架构](figures/arch.png)

来源：图2，PDF p.3，作者原始图件。自下向上看，前 $L/2$ 层产生共享 KV，后 $L/2$ 层的 query 随深度改变，K/V 则来自同一位置的 self-decoder 输出。图中的“共享”不表示各层的 attention 分布相同。

| 符号 | 含义与形状（简化为未展开多头） |
|---|---|
| $X^0$ | $N\times d$ 的 token embedding |
| $L$ | 总层数，前后各半 |
| $M=X^{L/2}$ | self-decoder 输出，$N\times d$ |
| $\hat K,\hat V$ | 全局共享 KV；采用 GQA 时实际宽度由 KV 头数决定 |
| $\hat Q^l$ | cross-decoder 第 $l$ 层独立计算的 query |
| $C$ | 滑动窗口长度，固定超参数 |
| $S_n$ | gRet 单头递推状态，形状 $d_k\times d_v$ |

前半部分是 pre-RMSNorm + efficient self-attention（ESA）+ 残差，再接 RMSNorm + SwiGLU + 残差：

$$Y^l=\operatorname{ESA}(\operatorname{LN}(X^l))+X^l,$$
$$X^{l+1}=\operatorname{SwiGLU}(\operatorname{LN}(Y^l))+Y^l.$$

这里 LN 实际使用 RMSNorm。SwiGLU 为 $(\operatorname{swish}(XW_G)\odot XW_1)W_2$。前半部分的关键约束是推理状态不随 $N$ 增长。（式1，p.3）

交界处只投影一次：

$$\hat K=\operatorname{LN}(M)W_K,\qquad \hat V=\operatorname{LN}(M)W_V.$$

后半部分每层计算：

$$\hat Q^l=\operatorname{LN}(X^l)W_Q^l,$$
$$Y^l=\operatorname{Attention}(\hat Q^l,\hat K,\hat V)+X^l,$$
$$X^{l+1}=\operatorname{SwiGLU}(\operatorname{LN}(Y^l))+Y^l.$$

Attention 使用因果 mask，且支持 GQA。顶层输出接 next-token 分类器。因果 mask 必不可少：训练时虽然把整段 KV 都算出来，位置 $n$ 仍只能读取 $m\le n$ 的缓存。（式2–3，p.4）

### 2.2 默认 self-decoder：门控保留机制（gated retention, gRet）

gRet 给 Retention 加上由当前输入决定的逐头衰减门。忽略 batch 和 head 维度，原文式4定义：

$$Q=(XW_Q)\odot\Theta,\quad K=(XW_K)\odot\bar\Theta,\quad V=XW_V,$$
$$\Theta_n=e^{in\theta},\qquad \gamma=\operatorname{sigmoid}(XW_\gamma)^{1/\tau},$$
$$D_{nm}=\begin{cases}\prod_{i=m+1}^{n}\gamma_i,&n\ge m,\\0,&n<m,\end{cases}\qquad O=(QK^\top\odot D)V.$$

$\Theta$ 表示位置旋转，$\tau$ 调节门值向 1 靠近的程度，越接近 1 越少遗忘。$D$ 是因果衰减矩阵；这里没有 attention 的 softmax。逐头标量门而非逐元素门是作者面向 tensor core 的设计选择。（§3.1，p.5）

同一运算可递推计算（式5，p.5）：

$$S_n=\gamma_nS_{n-1}+K_n^\top V_n,\qquad O_n=Q_nS_n.$$

状态的维度不随序列长度增长，但每层每头都需要一份。**固定内存指固定于长度，不是与层数和头维度都无关。** 展开递推可见，旧 token 的外积恰好乘上从它之后到当前的所有门值，因此与并行形式等价。

教学例子：$S_0=0$ 时，$S_2=\gamma_2K_1^\top V_1+K_2^\top V_2$。第一个 token 被遗忘多少取决于第二步门值；当前 token 没有额外衰减。

多头实现将各头结果拼接，做逐头 GroupNorm，再通过 swish 输出门和 $W_O$。这些归一化与门控也是完整模块的一部分，单写 $Q_nS_n$ 不是完整网络层。（式7，p.6）

### 2.3 分块计算：训练并行与推理递推的桥梁

逐 token 递推不利于 GPU 并行；直接构造整个 $N\times N$ 衰减矩阵又昂贵。分块做法在块内并行、块间递推，prefill 和长序列训练采用这一形式；逐 token generation 使用递推形式。

下面是**本次从式4–5重新推导、统一行列约定后的教学写法**，不是逐字照抄式6。块起点为 $a=(i-1)B$，定义 $\beta_j=\prod_{t=1}^{j}\gamma_{a+t}$，并令 query 对应行 $j$、key 对应列 $k$：

$$D^{(i)}_{jk}=\mathbf1_{k\le j}\frac{\beta_j}{\beta_k},$$
$$O_{[i]}=(Q_{[i]}K_{[i]}^\top\odot D^{(i)})V_{[i]}+\operatorname{diag}(\beta)Q_{[i]}R_{i-1},$$
$$R_i=\beta_B R_{i-1}+\sum_{k=1}^{B}\frac{\beta_B}{\beta_k}K_{a+k}^\top V_{a+k}.$$

第一项是当前块内的因果历史，第二项是前面所有块的压缩状态。状态更新中，每个块内 token 需衰减到块末。这样 $R_i=S_{iB}$，等价性直接来自把总和按当前块与旧块拆开。

**原文核验提醒**：式6（p.6）把 $D_{[i]}(j,k)$ 写成 $j\le k$ 时 $\beta_k/\beta_j$；按紧随其后的 $QK^\top$ 常规行列约定，这与式4的因果方向不一致。附录 B（pp.16–17）还有 $n$ 与 $n-1$ 乘积上界及指标混用；附录 F（p.18）伪代码把已赋值的 `gt` 写成了未定义的 `g`。这些材料适合理解思路，不能直接当作已验证实现。重实现时应以式4–5为依据，对三种形式做数值等价测试；本次未运行作者实现验证。

### 2.4 替代选择：滑动窗口 attention（SWA）

只允许 $i-C<j\le i$ 的历史参与注意力，其余位置设为 $-\infty$。每层缓存最近 $C$ 个 token，所以在固定 $C$ 时不随总长度增长。扩展曲线实验的窗口为 1,024。滑窗限制的是 self-decoder；cross-decoder 仍访问全局共享缓存。（§3.2，p.6；§4.2，p.8）

### 2.5 为什么可以提前结束 prefill

![原图3：预填充和生成路径](figures/inference.png)

来源：图3，PDF p.4，作者原始图件。提示词首先经过 self-decoder 填充共享 KV；历史提示词位置的 cross-decoder 输出无需保存。生成时，当前位置的 query 在上半部分逐层更新并读取全局 KV。

通常 next-token logits 仍需要提示词最后一个位置的 cross-decoder 计算。**省掉的是全部历史位置的上半段批量计算，不是完全不用上半段就能预测第一个 token。** 训练若要对所有位置计算 loss，也仍需所有相应输出。原文 §1 p.2 有一句写成“进入 self-decoder 前退出”，与图3和 §2.3 相矛盾；此处依据具体计算图与 §2.3，正确位置是进入 cross-decoder 前。

| 项目 | Transformer | YOCO | 适用前提 |
|---|---|---|---|
| KV / 状态内存 | $O(LND)$ | 原文表1：$O((N+L)D)$ | 固定窗口、头维度等隐藏常数 |
| 更明确的缓存数量 | $O(LN)$ | $O(N+CL)$ | SWA 情况，且长序列时 $N\gg CL$ |
| prefill attention 计算 | $O(LN^2D)$ | 表2：$O(LND)$ | 使用高效模块和提前退出 |

表1–2（p.4）是简化复杂度，不能把 $D$ 当成包含全部实现细节的精确字节公式。gRet 状态实际为逐头矩阵；若头维度也视为变量，必须显式计入它的乘积。参数、FFN、临时激活和通信同样影响实际成本。YOCO 没有消除 generation 阶段读取长 KV 的成本，也没有把训练中的全局 attention 自动变成线性。

### 2.6 分布式训练（附录 A，p.16）

把序列分块放在不同设备。self-decoder 在设备边界传递递推状态或必要的局部窗口；cross-decoder 的共享 KV 只需一次 all-gather，而非每一层重复收集。图11解释了通信依赖的简化，但正文未提供独立的通信带宽或端到端训练加速消融，不能直接赋予一个训练加速倍数。

## 三：实验——分清四种不同证据

| 实验 | 对象和协议 | 指标及比较角色 |
|---|---|---|
| §4.1 训练 token 扩展 | 3B gRet，训练至1.6T，保留1T checkpoint | 8项 LM Eval Harness zero-shot 分数，越高越好；与公开预训练模型比实用质量 |
| §4.2 参数扩展 | 160M、400M、830M、1.4B、2.7B、6.8B、13B，各10B训练token | validation LM loss，越低越好；相同数据/训练设置的架构比较 |
| §4.3 长上下文 | 3B继续训练到64K、256K、1M；无长指令微调数据 | 检索准确率↑、累积平均NLL↓ |
| §4.4 部署 profiling | 默认采用§4.1的3B配置，32K–1M，最后1,024 token作为生成 | 总推理显存↓、prefill秒数↓、含prefill与生成的吞吐↑ |

3B 模型：26层、隐藏维度3072、FFN8192、24个Q头/8个KV头、head dimension 128，非embedding参数2.83B；训练长度4096、batch 4M token、AdamW $(0.9,0.95)$、峰值学习率 $3.2\times10^{-4}$。**5T 是学习率日程目标，实际只训练1.6T**。（表5，p.17；§4.1）

参数扩展实验不是上述3B训练的缩小版：长度2048、batch 0.25M token、40k steps、AdamW $(0.9,0.98)$，通过 FFN 宽度对齐参数：Transformer 为 $8d/3$，YOCO 为 $3d$；gRet head dimension 为256。图4不能被解释成各规模均训练充分的 compute-optimal scaling 证明。（附录D，pp.17–18）

推理 baseline 已包含 GQA、Flash-Decoding 和 kernel fusion；YOCO 的 gRet 使用 Triton 内核，chunk size 256，在 H100-80GB 硬件上测试。论文未给出足以完全重建每个长度的 batch、并行设备配置、计时重复次数及误差区间。引言的7B、4卡背景例子不能当成§4.4默认3B实验配置。

## 四：结果——有效性与边界一起读

### 4.1 语言模型质量（表3，p.7）

| 模型 / 训练token | 八项平均分↑ | 对应比较 |
|---|---:|---|
| OpenLLaMA-3B-v2 / 1T | 0.619 | 公开Transformer baseline |
| StableLM-base-alpha-3B-v2 / 1T | 0.612 | 公开Transformer baseline |
| YOCO-3B / 1T | **0.634** | 对以上分别高1.5和2.2个百分点 |
| YOCO-3B / 1.6T | 0.636 | 对自身1T高0.2个百分点 |
| YOCO-3B-1M / 长度扩展后 | 0.645 | 又增加了长序列继续训练 |

八项是 ARC-C、ARC-E、BoolQ、HellaSwag、OBQA、PIQA、Winogrande、SciQ；表中为原文分数口径，未充分注明各任务的归一化细节。StableLM-3B-4E1T 的多项结果缺失，因此没有可比的完整平均分，不能补算成0分。

质量并非每项都赢：1T 时 YOCO 的 BoolQ 为0.645，低于 OpenLLaMA 的0.657；HellaSwag 为0.689，低于0.700。1T到1.6T也有个别任务回落。公开模型的数据、tokenizer与训练过程不是严格统一的，所以这些结果支持“有竞争力”，不能单独隔离出架构的因果收益。延长上下文后的0.645也混合了额外训练效应。

### 4.2 参数规模扩展（图4，pp.7–8）

![原图4：参数规模与验证损失](figures/scaling_curve.png)

横轴为非embedding参数规模，纵轴为验证loss，越低越好；不同曲线分别对应 Llama 风格 Transformer、YOCO-SWA 和 YOCO-gRet。相同10B token条件下，YOCO在160M–13B保持竞争力，gRet变体整体更好。作者把收益解释为 attention 与 retention 的归纳偏置互补，并提到1:3交错混合也有类似现象；这不是共享缓存本身带来所有质量收益的严格消融证明。图中未配完整原始点表或多种子误差线，本笔记不从图像臆造精确差值。

### 4.3 长上下文：检索、聚合、预测要分别判断

长度扩展配置（表7，p.18）：

| 长度 | 继续训练token | 学习率 | RoPE θ |
|---|---:|---:|---:|
| 65,536 | 6B | $8\times10^{-5}$ | 640K |
| 262,144 | 4B | $4\times10^{-5}$ | 5M |
| 1,048,576 | 1.5B | $2\times10^{-5}$ | 80M |

共额外11.5B token；上采样超过目标训练长度的长文档。不是直接把4K模型零训练外推到1M。

![原图5：百万长度单针检索](figures/1m_retrieval.png)

来源：图5，p.8。横轴为上下文长度，纵轴为插入深度，颜色为准确率；城市名与magic number组成needle，每个深度/长度配置重复10次取均值。近乎全绿说明此合成单针任务表现接近满分，不等于百万长度的多步推理或文档综合理解已解决。

多针结果是在 **128K** 下测得（表4，p.8），$N$ 为needle数而不是上下文长度：

| 模型 | 1针 | 2针 | 4针 | 8针 |
|---|---:|---:|---:|---:|
| YaRN-Mistral-128K 7B | 0.02 | 0.12 | 0.08 | 0.20 |
| LWM-1M-text 7B | 1.00 | 0.90 | 0.76 | **0.62** |
| MiniCPM-128K 2.4B | 1.00 | 1.00 | 0.54 | 0.56 |
| ChatGLM3-128K 6B | 0.94 | 0.72 | 0.52 | 0.44 |
| YOCO-3B-1M 3B | 0.98 | 0.98 | **0.84** | 0.56 |

YOCO在4针上突出，但8针低于LWM且与MiniCPM持平；因此作者“优于”部分模型的整体概括不能解释为每个needle数都严格更好。不同模型规模、长度训练与数据也不同。

图6（p.9）的书籍和repository-level代码实验选择超过1M token的验证样本，报告随已处理长度变化的**累积平均负对数似然 NLL**，越低越好。两组总体下降，支持模型利用长距离信息的解释；但这不是固定同一预测位置、只改变可见历史的受控曲线，累积平均还混入位置与样本难度的变化。原文段落标题提到perplexity，实际图注指标是NLL，两者不能直接混用。

### 4.4 推理：最大的收益来自长提示词和缓存容量

![原图7a：推理总显存](figures/memory-length.png)

来源：图7a，p.10。横轴长度、纵轴GB；绿色YOCO随长度增长较慢，橙色Transformer增长明显。插图显示32K也有收益。这里是权重、激活和缓存的总量，不是只测KV。

| 条件 | 原文结果 | 位置与注意点 |
|---|---|---|
| 3B、1M，总推理显存 | YOCO 12.4GB；Transformer约9.4倍 | 图7/§4.4 p.10；图标注9.38× |
| 3B、32K，总推理显存 | 约1.95×差距 | 图7插图；不等于26× |
| 65B、每token KV内存 | 约80×减少 | 图8 p.10；是层数相关的缓存指标，不是总内存 |
| 32K，prefill | 2.87×加速 | 图9 p.11 |
| 512K，prefill | 约180秒→不足6秒 | §4.4 p.10，图9标30.3× |
| 1M，prefill | 图9标71.82× | p.11；引言四舍五入为71.8× |
| 512K，总推理吞吐 | 4.5→43.1 token/s，约9.6× | §4.4 p.11；包含prefill与生成 |

65B 情况下，1GB缓存约能容纳YOCO的128K token，对比Transformer的1.6K；这不是说1GB就能加载65B模型权重。总显存收益小于KV倍数，因为权重等部分还存在。

prefill优势来自跳过历史位置的cross-decoder，以及self-decoder高效计算。吞吐优势还来自显存节省允许更大batch。**9.6×不是固定batch的单步decode加速，也不等于每个生成token延迟缩短9.6×。** 论文的测量细节不足以脱离其内核和硬件直接外推服务SLA。

### 4.5 附录诊断：全局记忆为何有价值

附录G.1（p.19）比较160M、12层、隐藏维度768的模型。H3是含两层attention的混合版；Mamba采用其double-SSM布局；RetNet和gRetNet调整value及FFN维度以匹配参数。

| 模型 | Valid. Set↓ | AR-Hit↓ | First-Occur↓ |
|---|---:|---:|---:|
| Mamba | 3.645 | 1.555 | 4.126 |
| RetNet | 3.633 | 1.466 | 4.131 |
| Hybrid H3 | 3.591 | 1.251 | 4.130 |
| gRetNet | 3.600 | 1.354 | 4.116 |
| Transformer | 3.564 | 1.219 | 4.104 |
| YOCO-SWA | 3.553 | 1.202 | 4.094 |
| YOCO-gRet | **3.530** | **1.199** | **4.067** |

原文表8将这些指标称为perplexity，此处保留原名与数值，不自行取指数或改标loss。AR-Hit针对上下文中出现过的bigram续写，衡量联想召回；First-Occur偏向不能直接从上下文重现的预测。YOCO两种分组均有竞争力，且相对纯gRetNet的AR-Hit差距更明显，支持“保留可寻址全局记忆有用”的解释；仍不构成对所有状态空间模型的普遍判断。

附录G.2（p.20）继续训练2B token、长度16,384，考察ZeroSCROLLS的Qasper、GovReport、QMSum、NarrativeQA。图12横轴为输入长度、纵轴为答案perplexity；YOCO与Transformer整体较好。Sparse Transformer保持2,048窗口且不改RoPE base，这个对照与其他模型长度扩展方案不同。评估的是答案语言建模分数，不是四项任务的生成准确率或ROUGE。

## 五：结论、局限与值得继续追问的问题

可靠结论是：共享跨层全局KV在保持全局attention的同时，改变了缓存与prefill的依赖结构；在本论文评估范围内，质量具有竞争力，长序列部署成本显著下降。论文把“高效压缩状态”和“逐token外部记忆”组合起来，而非只选择其中之一。

主要边界如下：

1. **缓存共享约束了表示**：cross-decoder深层可改变query与FFN，但不能写入新的逐层KV。哪些任务特别依赖深层重编码历史，本文没有系统隔离。
2. **全局缓存仍线性增长**：长序列decode仍需访问历史KV；压缩、索引、稀疏检索是未来工作。
3. **训练与推理成本不同**：early exit适用于构建生成所需缓存；完整训练loss并不能据此跳过全部cross-decoder。
4. **质量证据范围有限**：1M单针测试不等于1M多针；公开模型比较不是完全受控；没有普遍的多种子显著性证据。
5. **系统复现需补充条件**：需确认batch、精度、设备数、并行策略、warmup及吞吐统计分子；原文未完整交代。
6. **实现材料需校验**：分块公式与伪代码有前述索引/变量问题，不能直接复制运行。

作者提出YOCO+BitNet+Groq、多模态多self-decoder、KV压缩/索引以及预缓存支持RAG等方向（§5，pp.11–12）。这些是展望，不是本论文已经完成的实验。

本次分析建议的后续实验：在同数据与token预算下单独消融“共享KV”和“gRet混合”；固定batch测prefill、首token与单步decode；在固定目标位置比较不同可见历史长度；用小张量检查并行、递推、分块形式及梯度的一致性。这些建议未在本次研读中执行。

## 研读交付范围

本笔记基于同版本作者源码与PDF，不是全文翻译；英文幻灯片见 [slides.en.pdf](slides.en.pdf)，可编辑源文件见 [slides.en.tex](slides.en.tex)。图件来自作者原始资源，PDF图仅转为PNG供Markdown显示，没有生成或修改实验数据。