# T3 代码与论文公式对照

这份文档的目标不是简单列一个“变量翻译表”，而是回答 3 个问题：

1. 论文 **DIVERSED: Relaxed Speculative Decoding via Dynamic Ensemble Verification** 的主线公式到底是怎么推出来的？
2. 本题的 `submission.py` 到底保留了论文里的哪些关键结构？
3. 本题又在哪些地方做了化简，为什么要这么化简？

如果你先读这份文档，再去读 `submission.py`，会更容易看出：  
**这道题并不是完整复现论文系统，而是尽量忠实地抽出了论文“动态验证器”这一核。**

## 1. 先按论文原文主线读一遍

论文的主线可以分成 4 层：

1. **标准 speculative decoding**：先解释原始验证规则是什么
2. **static ensemble**：说明为什么“放松验证”可以提高 acceptance rate
3. **dynamic ensemble**：说明为什么固定权重不够，需要上下文相关的动态权重
4. **training + inference**：说明这个动态权重如何训练，以及真实系统如何运行

本题最直接对应的是第 3 层，但要真正理解第 3 层，最好先把前两层也看明白。

## 2. 标准 speculative decoding 在做什么

论文 Section 2 的起点是标准 speculative decoding。

在第 `t` 步：

1. target model 给出分布 `p_t(·)`
2. draft model 给出分布 `q_t(·)`
3. draft model 提出一个草稿 token `x̃_t`
4. verifier 决定这个 `x̃_t` 是直接接受，还是拒绝后回退

论文里的标准接受概率写成：

`b_t(x̃_t) = min{1, p_t(x̃_t) / q_t(x̃_t)}`

如果拒绝，则从回退分布重采样：

`P_t(x) ∝ [p_t(x) - q_t(x)]_+`

这里的 `[·]_+` 表示正部，也就是 `max(·, 0)`。

这套规则的核心优点是：

1. 接受时非常严格
2. 最终生成分布仍然严格对齐 target model

但它的问题也正是“太严格”：

1. 许多其实还算合理的 draft token 会被拒绝
2. acceptance rate 不够高
3. 因而 wall-clock speedup 也会受到限制

论文就是从这里开始问：  
**能不能把 verifier 稍微放松一点，同时不要让质量掉得太多？**

## 3. static ensemble 为什么会出现

论文 Section 3.1 先讨论一个最自然的放松方式：  
不再总是拿 `p_t` 当验证分布，而是拿 `p_t` 与 `q_t` 的一个固定加权平均。

于是有了 static ensemble：

`ν_t(x) = w p_t(x) + (1 - w) q_t(x)`  （Equation (2)）

这里的 `w` 是常数，不随 token 或上下文变化。

然后接受规则也相应改成：

`b_{ν_t}(x̃_t) = min{1, ν_t(x̃_t) / q_t(x̃_t)}`

回退分布改成：

`P_{ν_t}(x) ∝ [ν_t(x) - q_t(x)]_+`

这一步的重要意义是：

1. verifier 不再死板地完全贴着 target
2. 只要 `ν_t` 更靠近 draft，acceptance rate 通常就会提高
3. 但权重 `w` 一旦固定，就无法区分“当前这个位置到底该更信 target 还是更信 draft”

论文在这里还给出一个理论结论：  
**static ensemble 沿着 acceptance-quality 的 Pareto front 运行。**

换句话说，它已经说明：

1. “放松验证”这件事本身是合理的
2. 问题不在于要不要混合，而在于“混多少”不能固定写死

## 4. dynamic ensemble 是论文真正的核心

论文 Section 3.2 的核心判断是：

**固定权重不够，因为不同 token、不同上下文下，draft 与 target 的分歧程度不同。**

因此论文把固定权重 `w` 推广成动态权重：

`w_t = f_θ(h_t^q, h_t^p)`  （Equation (3)）

这里：

1. `h_t^q` 是 draft model 在第 `t` 步的 hidden representation
2. `h_t^p` 是 target model 在第 `t` 步的 hidden representation
3. `f_θ` 是一个小型前馈网络，也就是 verifier head

然后，动态 ensemble 分布写成：

`ν_t^θ(x) = w_t p_t(x) + (1 - w_t) q_t(x)`  （Equation (4)）

接受规则与回退分布相应变成：

`b_{ν_t^θ}(x̃_t) = min{1, ν_t^θ(x̃_t) / q_t(x̃_t)}`

`P_{ν_t^θ}(x) ∝ [ν_t^θ(x) - q_t(x)]_+`  （Equation (6)）

这一段是本题最重要的地方。  
你可以把它理解成：

1. `Equation (3)` 决定“当前这一步更偏向 target 还是更偏向 draft”
2. `Equation (4)` 把这个偏好变成真正的验证分布
3. `Equation (6)` 再把这个验证分布变成接受 / 回退机制

所以从实现角度看，论文的动态验证器其实就是三连：

1. **先算 `w_t`**
2. **再算 `ν_t^θ`**
3. **再按 `ν_t^θ` 做 accept / fallback**

## 5. 训练目标与真实推理流程是什么

论文并不是只在推理时定义了一个 `w_t`，然后停在这里。  
它还要回答两个更完整的问题：

1. `f_θ` 怎么训练？
2. 真实系统里怎么把它嵌回 speculative decoding？

论文的训练目标写成：

`E[R(x_{1:T}) + λ Σ_t (1 - TV(q_t, ν_t^θ))]`  （Equation (5)）

直观上，这个目标同时优化两件事：

1. 最终生成质量 / 任务奖励
2. acceptance 友好性

因为 `ν_t^θ` 越接近 `q_t`，通常 acceptance rate 越高；但如果一味贴着 `q_t`，质量又可能掉太多。  
所以训练目标本质上是在学习一个更好的速度-质量折中点。

真实推理时，论文的 Algorithm 2 运行的是 block-level speculative decoding：

1. draft model 先提出一个 block
2. target model 并行验证这个 block
3. verifier head 在每个位置给出一个 `w_t`
4. 再按 Equation (6) 对每个位置做 accept / fallback

也就是说，论文的真实系统是：

1. **双模型**
2. **并行 block 验证**
3. **位置相关动态权重**

## 6. 本题到底保留了什么

本题保留了论文最关键的 3 层结构：

1. `w_t = f_θ(h_t^q, h_t^p)`  
   对应 `dynamic_ensemble_weight(...)`
2. `ν_t^θ(x) = w_t p_t(x) + (1 - w_t) q_t(x)`  
   对应 `ensemble_distribution(...)`
3. `b_{ν_t^θ}(x̃_t)` 与 `P_{ν_t^θ}(x)`  
   对应 `acceptance_probability(...)` 与 `residual_distribution(...)`

再往上一层：

4. 把单步验证组织起来  
   对应 `verify_single_step(...)`
5. 把多个时间步串起来  
   对应 `run_dynamic_verification(...)`

也就是说，这道题的 `submission.py` 虽然短，但它其实正好把论文验证器最核心的链条拆成了 4 个可评测的小函数。

## 7. 本题又化简了什么

本题并不是完整复现论文系统。它做了 4 个重要化简。

### 化简 1：高维 hidden states 变成低维导出特征

论文里的 `h_t^q`、`h_t^p` 是真实模型内部的高维 hidden states。  
如果把它们原样放进考试：

1. 数据会很大
2. 函数接口会非常难读
3. 评测和样例展示都会变得很笨重

因此本题给的是低维导出的 `h_q`、`h_p`：

1. 语义上仍然对应 `h_t^q / h_t^p`
2. 但不再是原始大维度 hidden state
3. 这样既保住了 Equation (3) 的结构，又能让题目保持可写性

所以更准确地说：

1. 论文原文：`w_t = f_θ(h_t^q, h_t^p)`
2. 本题实现：`w_t = f_θ(h_q, h_p)`

差别在于 `h_q / h_p` 是压缩后的 hidden feature，而不是完整 hidden representation。

### 化简 2：训练好的小网络直接随题给出

本题没有要求考生自己训练 verifier head，直接提供：

1. `ensemble_head.json`
2. 已经训练好的小型前馈网络参数

这意味着：

1. 你不需要自己设计 heuristic weight
2. 你也不需要自己训练 head
3. 你要做的是：**按给定参数把 Equation (3) 真的前向出来**

### 化简 3：fallback 不做随机采样，而是取 argmax

论文在 rejection 时，会从 `P_{ν_t^θ}` 中采样。  
但如果考试也这么做，评测会有两个问题：

1. 随机性导致答案不稳定
2. 相同代码多次运行可能得到不同 token

所以本题保留了 fallback distribution 的定义，但把最终动作改成：

1. 先按论文式子算出 `P_{ν_t^θ}`
2. 再取其 `argmax`

这不会改变“回退往哪里回”的基本方向，但会让评测完全确定性。

这里需要特别强调：

1. **你仍然必须先正确算出 `P_{ν_t^θ}`**
2. 只有“最后一步如何选 token”被改了
3. 因此本题不是把 fallback 逻辑删掉，而是把“随机采样”替换成“确定性选最大项”
4. 如果直接照论文原文写成随机采样，即使前面的 `ν_t^θ` 和 `P_{ν_t^θ}` 都算对了，最终输出仍然会和标准答案不一致

### 化简 4：block-level inference 改成按时间步串联

论文的真实系统按 block 运行 speculative decoding。  
本题则拆成：

1. `verify_single_step(...)`
2. `run_dynamic_verification(...)`

这样做的原因是：

1. 便于单测
2. 便于题目只聚焦验证器
3. 便于把复杂系统拆成考生可以手写的若干小块

因此，本题不是 Algorithm 2 的逐行翻译，而是把 Algorithm 2 里最关键的 **verification kernel** 抽出来了。

## 8. 论文中的核心记号

这道题最重要的论文记号如下：

1. `p_t(·)`：target model 在第 `t` 步的分布
2. `q_t(·)`：draft model 在第 `t` 步的分布
3. `x̃_t`：draft model 提出的 token
4. `x_t`：最终被接受或回退后的 token
5. `h_t^q`：draft hidden representation
6. `h_t^p`：target hidden representation
7. `w_t = f_θ(h_t^q, h_t^p)`：动态权重（Equation (3)）
8. `ν_t^θ(x)`：动态 ensemble 分布（Equation (4)）
9. `b_{ν_t^θ}(x̃_t)`：接受概率（Equation (6)）
10. `P_{ν_t^θ}(x)`：回退分布（Equation (6)）

## 9. 变量对照

### 9.1 输入分布与 token

| 代码变量 | 论文符号 | 说明 |
| --- | --- | --- |
| `target_probs` | `p_t(·)` | 当前步 target 分布。 |
| `draft_probs` | `q_t(·)` | 当前步 draft 分布。 |
| `drafted_token` | `x̃_t` | 当前步草稿 token。 |
| `target_probs_seq` | `{p_t(·)}` | 多个时间步上的 target 分布序列。 |
| `draft_probs_seq` | `{q_t(·)}` | 多个时间步上的 draft 分布序列。 |
| `draft_tokens` | `{x̃_t}` | 多个时间步上的草稿 token 序列。 |

### 9.2 hidden feature 与 verifier head

| 代码变量 | 论文符号 | 说明 |
| --- | --- | --- |
| `h_q` | `h_t^q` | 低维导出的 draft hidden feature。 |
| `h_p` | `h_t^p` | 低维导出的 target hidden feature。 |
| `h_q_seq` | `{h_t^q}` | 多个时间步上的 `h_t^q` 序列。 |
| `h_p_seq` | `{h_t^p}` | 多个时间步上的 `h_t^p` 序列。 |
| `ensemble_head` | `f_θ` | 训练好的 verifier head 参数。 |

### 9.3 权重与混合分布

| 代码变量 | 论文符号 | 说明 |
| --- | --- | --- |
| `w_t` | `w_t` | 当前步给 target 分布的动态权重。 |
| `verify_probs` | `ν_t^θ(·)` | 当前步动态 ensemble 分布。 |
| `weights` | `{w_t}` | 多个时间步上的动态权重序列。 |

### 9.4 接受 / 回退

| 代码变量 | 论文符号 | 说明 |
| --- | --- | --- |
| `accept_prob` | `b_{ν_t^θ}(x̃_t)` | 当前步接受概率。 |
| `residual_probs` | `P_{ν_t^θ}(·)` | 当前步 fallback distribution；本题最终不从这里随机采样，而是取其 `argmax`。 |
| `accepted` | `A_t` | 当前步是否接受草稿 token。 |
| `accepted_mask` | `{A_t}` | 多个时间步上的接受标记序列。 |
| `token` | `x_t` | 当前步最终 token。 |
| `tokens` | `{x_t}` | 多个时间步上的最终 token 序列。 |

## 10. 函数对照

| 代码函数 | 论文对应 | 说明 |
| --- | --- | --- |
| `dynamic_ensemble_weight(h_q, h_p, ensemble_head)` | Equation (3) | 计算 `w_t = f_θ(h_t^q, h_t^p)`。 |
| `ensemble_distribution(target_probs, draft_probs, w_t)` | Equation (4) | 计算 `ν_t^θ(x) = w_t p_t(x) + (1 - w_t) q_t(x)`。 |
| `acceptance_probability(drafted_token, verify_probs, draft_probs)` | Equation (6) | 计算 relaxed accept rule。 |
| `residual_distribution(verify_probs, draft_probs)` | Equation (6) | 计算 fallback distribution。 |
| `verify_single_step(...)` | 单步 verification | 把 Equation (3)(4)(6) 组织到一个时间步。 |
| `run_dynamic_verification(...)` | 简化版 inference pipeline | 把多个时间步上的 verification 串起来。 |
