# T3 星港急速译码局（DIVERSED 验证核）

## 题目背景

星港是一座用于向太空输送物资的集散港。每天都有大量调度口令、维修通报、回收窗口说明和夜班摘要需要被快速整理成正式文本。为此，星港在内部部署了一套双模型推理链路：

1. 小模型先快速起草 token
2. 大模型再验证这些草稿 token
3. 验证器决定哪些草稿可以直接接受，哪些必须回退

这种思路正对应大模型推理加速中的 **投机采样（speculative decoding）**。你在调研星港内部推理加速方案时，读到了一篇最新论文：**DIVERSED: Relaxed Speculative Decoding via Dynamic Ensemble Verification**（arXiv:2604.07622，提交日期：2026-04-08）。本题考察的，就是这篇论文里最关键的一段：**验证阶段如何从“严格验证”升级为“动态 ensemble 验证”**。

## 题目描述

### 什么是投机采样

标准 speculative decoding 的基本流程是：

1. draft model 先给出当前步的草稿分布 `q_t(·)`，并提出一个草稿 token `x̃_t`
2. target model 再给出当前步的更可靠分布 `p_t(·)`
3. 验证器根据 `p_t` 与 `q_t` 决定是否接受 `x̃_t`
4. 若拒绝，则从回退分布中重新选择最终 token `x_t`

它的目标是：**尽量复用小模型的速度，同时尽量保持大模型的质量**。

### 标准 speculative decoding 的问题

标准 speculative decoding 的验证规则比较刚性。它往往要求验证分布尽量严格贴住 `p_t`，因此会拒绝不少“其实也足够合理”的草稿 token。这样会带来两个直接问题：

1. acceptance rate 偏低
2. 小模型提出的大量草稿无法真正转化为加速收益

也就是说，瓶颈往往不只在小模型，而是在**验证器过于保守**。

### DIVERSED 改进了什么

DIVERSED 的核心改进是：不再固定只看 `p_t`，而是先用一个小型 verifier head 根据上下文动态预测权重

`w_t = f_θ(h_t^q, h_t^p)`

然后构造动态 ensemble 分布

`ν_t^θ(x) = w_t p_t(x) + (1 - w_t) q_t(x)`

再用这个更灵活的分布做接受 / 回退决策。直观地说：

1. 当 draft 与 target 很接近时，可以更信任 draft
2. 当二者分歧明显时，就更多依赖 target
3. 因此接受率有机会提高，同时质量不至于失控

### 本题中的动态验证器

本题直接采用论文风格的动态验证器：

1. `w_t` 表示给 target 分布 `p_t` 的权重
2. `verify_probs` 直接对应论文中的 `ν_t^θ`
3. `acceptance_probability` 对应论文 Equation (6) 的 relaxed accept rule
4. `residual_distribution` 对应论文 Equation (6) 的 fallback distribution

其中有一条必须特别注意的评测约定：

1. 论文原文在 rejection 后，是从 `P_{ν_t^θ}` 中做随机采样
2. 本题为了保证函数式评测完全确定、结果可复现，不做随机采样
3. 你必须先按论文定义算出 `residual_distribution`
4. 然后把最终回退 token 固定写成 `argmax(residual_distribution)`

也就是说：**本题保留了论文的 fallback distribution 定义，但把“从分布采样”改成了“取分布最大项”。**

为了让考生把注意力集中在验证逻辑本身，题目随题提供了一个**已经训练好的小型 verifier head**：

1. 文件路径：`T3/ensemble_head.json`
2. 结构：一个小型两层前馈网络
3. 输入：本题导出的低维 `h_q`、`h_p`
4. 输出：当前步动态权重 `w_t`

因此，本题虽然仍是纯 Python 评测，但保留了论文最关键的结构：

1. `w_t = f_θ(h_t^q, h_t^p)`
2. `ν_t^θ(x) = w_t p_t(x) + (1 - w_t) q_t(x)`
3. `b_{ν_t^θ}(x̃_t) = min(1, ν_t^θ(x̃_t) / q_t(x̃_t))`
4. `P_{ν_t^θ}(x) ∝ [ν_t^θ(x) - q_t(x)]_+`

### 为什么还要化简

真实系统中的 `h_t^q`、`h_t^p` 是大维度 hidden states，而且推理过程还是 block-level 的 speculative decoding。为了让题目可评测、可手写、可确定性复现，本题做了两层化简：

1. 不直接给高维 hidden states，而是给导出的低维 `h_q_seq / h_p_seq`
2. 不做 block-level 并行验证，而是把验证阶段改写成按时间步串联
3. 不在 rejection 后做随机采样，而是固定取 `residual_distribution` 的最大项

因此，这道题不是完整复现论文的系统工程版，而是一个**高度贴近论文验证核的纯 Python 版**。

### `T3_real_demo/` 目录是什么

`T3_real_demo/` 是从早期 `T3/run_real_demo/` 移动并整理出来的真实双模型 Demo，用来帮助你把这道题与真实系统对应起来。它不属于 `T3` 正式函数式评测，但会真正加载：

1. `Qwen3-0.6B` 作为 draft model
2. `Qwen3-8B` 作为 target model
3. 一个真实训练得到的 verifier head

而正式考试里的 `submission.py` 只抽出“验证核”本身，方便做函数式评测。二者关系如下：

1. `submission.py`：考察论文验证核
2. `T3_real_demo/`：展示真实系统如何训练 head、加载双模型、运行 benchmark

## 你需要实现

在 `submission.py` 中补全以下 4 个函数：

1. `dynamic_ensemble_weight`
2. `ensemble_distribution`
3. `verify_single_step`
4. `run_dynamic_verification`

模板已经提供：

1. `normalize_probs`
2. `acceptance_probability`
3. `residual_distribution`
4. `_dense_forward`

其中 `ensemble_head.json` 会由评测脚本读入，并作为 `ensemble_head` 参数传给你的函数。

为了把难度集中在论文核心逻辑上，当前模板还额外给出了这几部分骨架：

1. `dynamic_ensemble_weight`
   只保留了 verifier head 的输入输出约定，你需要自己完成：
   拼接、第一层前向、`ReLU`、第二层线性映射（hidden -> scalar logit）以及最终的 `w_t`
2. `ensemble_distribution`
   已给出 `target_probs / draft_probs` 的归一化，你只需按 Equation (4) 写出动态混合
3. `verify_single_step`
   只保留了辅助函数接口，你需要自己补完：
   `w_t`、`verify_probs`、`accept_prob`、`residual_probs`、接受/拒绝分支以及返回字典
4. `run_dynamic_verification`
   已给出长度裁剪和结果列表初始化，你只需补完逐步调用与结果汇总

## 输入格式

本题为函数式评测。

常见输入形式如下：

1. `dynamic_ensemble_weight(h_q, h_p, ensemble_head)`
2. `ensemble_distribution(target_probs, draft_probs, w_t)`
3. `verify_single_step(target_probs, draft_probs, h_q, h_p, drafted_token, uniform_u, ensemble_head)`
4. `run_dynamic_verification(target_probs_seq, draft_probs_seq, h_q_seq, h_p_seq, draft_tokens, uniforms, ensemble_head)`

其中：

1. `target_probs` 对应论文中的 `p_t`
2. `draft_probs` 对应论文中的 `q_t`
3. `h_q / h_p` 对应本题导出的低维 `h_t^q / h_t^p`
4. `ensemble_head` 对应训练好的 verifier head 参数

## 输出格式

各函数返回要求如下：

1. `dynamic_ensemble_weight` 返回浮点数 `w_t`
2. `ensemble_distribution` 返回概率列表 `verify_probs`
3. `verify_single_step` 返回字典，包含：
   `token`、`accepted`、`w_t`、`accept_prob`、`verify_probs`、`residual_probs`
4. `run_dynamic_verification` 返回字典，包含：
   `tokens`、`accepted_mask`、`weights`、`accept_probs`

其中 `weights` 存的是每一步的 `w_t`。

## 输入输出样例

### 样例 1：`dynamic_ensemble_weight`

```python
dynamic_ensemble_weight(
    h_q=[0.98852954, 0.24781233, 0.21888716, 0.02892517, 0.75, 0.49951020],
    h_p=[0.89758609, 0.35890295, 0.31411302, 0.04478993, 0.75, 0.47646760],
    ensemble_head=ensemble_head,
)
```

```python
0.5406200888
```

说明：这里直接按训练好的 verifier head 前向计算 Equation (3) 中的 `w_t`。

### 样例 2：`ensemble_distribution`

```python
ensemble_distribution(
    target_probs=[0.14146129, 0.31411302, 0.11352120, 0.35890295, 0.07200153],
    draft_probs=[0.21888716, 0.20332875, 0.13430597, 0.24781233, 0.19566578],
    w_t=0.5406200888,
)
```

```python
[0.1770291811, 0.2632209545, 0.1230693070, 0.3078701539, 0.1288104035]
```

说明：这里按论文 Equation (4) 做混合：
`ν_t^θ(x) = w_t p_t(x) + (1 - w_t) q_t(x)`。

### 样例 3：`verify_single_step`

```python
verify_single_step(
    target_probs=[0.29078671, 0.06818170, 0.11572419, 0.11590222, 0.10095468, 0.14480821, 0.16364229],
    draft_probs=[0.19036332, 0.03812279, 0.12309340, 0.05884189, 0.22237712, 0.15286530, 0.21433616],
    h_q=[0.93345821, 0.22237712, 0.21433616, 0.00804096, 0.66666667, 0.56678121],
    h_p=[0.95026642, 0.29078671, 0.16364229, 0.12714442, 0.00000000, 0.45950838],
    drafted_token=4,
    uniform_u=0.71563129,
    ensemble_head=ensemble_head,
)
```

```python
{
    "token": 0,
    "accepted": False,
    "w_t": 0.5707159669,
    "accept_prob": 0.6883774446,
    "verify_probs": [
        0.2476765538, 0.0552778902, 0.1188876752, 0.0914071319,
        0.1530793967, 0.1482669914, 0.1854043608
    ],
    "residual_probs": [
        0.5354696651, 0.1602777468, 0.0, 0.3042525882, 0.0, 0.0, 0.0
    ],
}
```

说明：由于 `uniform_u > accept_prob`，因此本步拒绝草稿 token `4`，并从 `residual_probs` 的最大项回退到 token `0`。
这里要再次注意：论文原文会从 `residual_probs` 对应的分布中采样；本题不能这样写，必须固定取 `argmax(residual_probs)`。

### 样例 4：`run_dynamic_verification`

```python
run_dynamic_verification(
    target_probs_seq=[
        [0.06704573, 0.13515031, 0.04974343, 0.17823654, 0.28255892, 0.04860543, 0.23865964],
        [0.29078671, 0.06818170, 0.11572419, 0.11590222, 0.10095468, 0.14480821, 0.16364229],
    ],
    draft_probs_seq=[
        [0.08198483, 0.14908803, 0.06882808, 0.22092038, 0.21743584, 0.09050639, 0.17123646],
        [0.19036332, 0.03812279, 0.12309340, 0.05884189, 0.22237712, 0.15286530, 0.21433616],
    ],
    h_q_seq=[
        [0.95480593, 0.22092038, 0.21743584, 0.00348454, 0.50000000, 0.54986656],
        [0.93345821, 0.22237712, 0.21433616, 0.00804096, 0.66666667, 0.56678121],
    ],
    h_p_seq=[
        [0.90156767, 0.28255892, 0.23865964, 0.04389928, 0.66666667, 0.59576124],
        [0.95026642, 0.29078671, 0.16364229, 0.12714442, 0.00000000, 0.45950838],
    ],
    draft_tokens=[3, 4],
    uniforms=[0.09750454, 0.71563129],
    ensemble_head=ensemble_head,
)
```

```python
{
    "tokens": [3, 0],
    "accepted_mask": [1, 0],
    "weights": [0.4462484883, 0.5707159669],
    "accept_probs": [0.9137807101, 0.6883774446],
}
```

说明：第一步接受草稿 token；第二步拒绝草稿 token，并回退到 residual 分布中概率最大的 token。

## 数据范围

### 可见数据

1. 文件路径：`../data/T3_public_cases.jsonl`
2. 公开集中的每条样本包含：
   `target_probs`、`draft_probs`、`h_q_seq`、`h_p_seq`、`draft_tokens`、`uniforms`
3. 公开集已给出标准结果 `expected`，用于评测脚本直接对比最终输出

### 隐藏数据

1. 本题当前版本不设置隐藏测试点
2. `../data/T3_public_cases.final.jsonl` 是赛后公开的补充案例和答案对照材料，默认 `evaluate.py` 不读取它参与计分

另外，`T3/ensemble_head.json` 中提供了训练好的 verifier head 参数；正式评测与公开自测使用同一份 head。

## 评分说明

总分 `20` 分：

1. A. 4 个核心函数单测 `16` 分
   `dynamic_ensemble_weight`、`ensemble_distribution`、`verify_single_step`、`run_dynamic_verification`
   每个函数 `4` 分
2. B. 公开评测集 `4` 分

### 运行方式

```bash
python T3/evaluate.py
```

### 提示

1. 本题已经统一使用论文原文风格的 `w_t`，不再改写成 `w_draft`
2. `weights` 字段记录的是每一步的 `w_t`
3. 评测脚本不会现场推导参考解，而是直接读取标准结果 `expected`
4. `T3_real_demo/` 只是帮助理解真实系统，不属于正式评分范围
5. 论文原文的 rejection 后是“从 fallback distribution 采样”，本题评分时一律按 `argmax(residual_distribution)` 判定；如果写成随机采样，将无法通过确定性评测
6. 如需进一步查看“代码变量 / 函数 / 公式”对应关系，请阅读：
   `T3/代码与论文公式对照.md`
