大语言模型通过潜在蒸馏进行探索 Large Language Models Explore by Latent Distilling

完整流程在 Algorithm 1 中给出，每个 decode step 包含两阶段。Phase 1 是 Forward & Logits Process：先用 LLM embedding + 第一层 transformer 得到浅层隐藏 $H^1_t$（line 6），异步触发 Distiller $f_\phi$ 预测 $\hat{H}^L_t = f_\phi(H^1_t)$（line 8）；同时主 stream 继续跑第 2 到 $L$ 层得到真实深层表征 $H^L_t$（line 11）。然后用 frozen LM head 同时对 $H^L_t$ 和 $\hat{H}^L_t$ 做投影，得到 $\text{logit}_{ref}$ 和 $\text{logit}_{dist}$（line 14）。对 $\text{logit}_{ref}$ 做可选 TopK 处理后，按公式 $\text{logit}_{new} = (1+\beta)\text{logit}_{ref} - \beta \text{logit}_{dist}$ 融合（line 16），softmax 后采样得 $z_t$（line 22）。Phase 2 是 Online Update：以 batch 为单位计算 MSE 损失 $\mathcal{L} = \frac{1}{|B|}\sum_{i\in B} \| h^L_{t,i} - f_\phi(h^1_{t,i}) \|_2^2$（line 19），用 Adam（$\eta=4\times 10^{-4}$、gradient clip 0.5）更新 $\phi$（line 20），最后由 vLLM worker 做后处理与调度（line 24）。异步实现的关键洞见是：蒸馏器只需要 LLM 第一层的输出就可以预测最后一层，因此它的 forward 与 LLM 后续 $L-1$ 层完全并行；而它的 backward 阶段被推迟到 CPU-bound 的 sampling 和 de-tokenization 阶段执行，这部分算力往往闲置。论文在 Llama-3-8B 上测出这中间的"slack"长达 15-20ms，而 2 层 MLP 蒸馏器只需不到 0.5ms 的前向，留出大量余裕。 此外，作者在开源实现中进一步利用 post-filter candidate intervention——只在 vLLM sampler 过滤后的候选 token 上做干预——以及 Triton grouped prediction backend，把 Distiller 前向的 GPU kernel launch overhead 也压到最低。

技术新颖性体现在三个层面：第一，理论层——首次把 RND 的"预测误差 = 新颖性"思想搬进 LLM 内部浅-深表征映射，并严格证明在 vanishing redundancy 假设下，per-step intrinsic reward $r(s,z) = \log \pi_{ref}(z|s) - \log q_{dist}(z|s)$ 恰好是 KL-正则化目标 $J(\pi) = \mathbb{E}[r] - \alpha D_{KL}(\pi\|\pi_{ref})$ 的闭式最优解；附录 A.1 给出完整 Lagrange 推导，附录 A.2 形式化了 vanishing redundancy 定义并证明 sequential Q 函数在此条件下退化为 per-step reward。第二，算法层——通过把 $\Delta \text{logit}_z = \beta \langle w_z, e_t \rangle$（其中 $e_t = h^L_t - \hat{h}^L_t$）拆解为 $\beta \| w_z \| \cdot \| e_t \| \cdot \cos(w_z, e_t)$，清晰区分了"context novelty $\| e_t \|$"和"semantic direction $\cos(w_z, e_t)$"两个独立维度，并据此自然地引导生成走向新颖语义而非表面词汇。第三，系统层——设计了 dual CUDA stream + pre-allocated ring buffer + post-filter candidate intervention + Triton grouped prediction backend 的一整套工程优化，把异步训练推理做到 1.2% 的实测开销，并通过 producer-consumer 抽象封装在 tLLM 框架里以避免维护 vLLM fork。文中噪声消融（Table 10）进一步验证：把 $e_t$ 替换为等大方差的高斯向量后性能塌缩到 vanilla 水平，证明改进源于误差向量的方向结构而非随机扰动。

独立分析三个潜在弱点并给出改进方向。第一，per-prompt 与 shared Distiller 的二选一难题：当前实现按 benchmark 人工切换，Appendix C.10 的实验数据（共享 vs per-prompt 在 AIME 上 Pass@64 差距 3.4pp）说明这是个值得优化的超参。可以借鉴 multi-task learning 的梯度手术方法，对不同 prompt 的梯度做正交投影来自动避免跨题干扰。第二，蒸馏器架构选择空间小：Table 17 仅比较 2 层 Gated SwiGLU、4 层 Gated SwiGLU、4 层 Plain MLP 三种且 Pass@64 都是 80.0%，但 4 层把 Pass@1 从 41.0% 拉到 38.9%。可考虑引入自适应架构，例如根据当前 prefix 长度动态选择宽度，或用 NAS 搜索轻量结构以更充分捕捉层间非线性。第三，Pass@1 与探索强度 $\beta$ 的耦合：Table 2 显示 $\beta=0.5$ 时 Pass@1 从 6.0% 降到 4.8%，说明过度新颖会伤害单样本质量。可以引入 KL 置信度门控，仅当 reference logits 的 top-1 概率低于阈值时才注入 novelty 信号，类似 contrastive decoding 的 amateur/expert gating；或者让 $\beta$ 在生成过程中按 prefix 位置的预测误差自适应衰减。第四，Case study（Appendix E.6）显示 ESamp 会产生"先猜后算"甚至"试错"等非常规路径——这虽然提升了多样性，但也意味着在某些对格式严格的任务（如要求 $\boxed{ANSWER}$ 的 AIME）上可能产生格式偏差，可考虑在 logit 融合前加一个轻量的 format-preserving 约束。第五，对闭源模型不友好：方法依赖访问 LLM 第一层和最后一层的 hidden state，这要求白盒接入；可通过构建可插拔 hook 框架（如 vLLM 的 SpecDecodeWorker 接口）来缓解。

作者在 Conclusion 与附录中提出的方向：1) 把 ESamp 嵌入更广泛的 test-time scaling 流水线，特别是 reward model reranking——附录 C.8 已经展示了 ESamp+SC 的兼容性，未来可以与 Process Reward Model (PRM) 联合优化；2) 自适应 shared/per-prompt Distiller 策略，Appendix C.10 把此作为"有前途的方向"；3) 进一步工程优化——附录 D.2 暗示若放弃 producer-consumer 接口的通用性，吞吐可继续逼近 vanilla。我额外想到的延伸方向：第一，把 Distiller 预测目标从浅-深映射换成"第 $i$ 层到第 $j$ 层"的任意组合，并通过学习最优跨度来自动发现哪些层承载了任务相关语义——这与 mechanistic interpretability 中的 circuit discovery 自然耦合。第二，把 RND 的 novelty 思路扩展到 multi-modal LLM，让 Distiller 同时预测 vision encoder 输出到 LLM 深层表征的迁移，从而在视觉问答等任务上鼓励多样化图像推理路径。第三，与 MCTS 结合：把 ESamp 当作 expansion 阶段的 novelty prior，用 Distiller 的预测误差指导 UCT 公式中的 exploration bonus，从而避免显式多次前向带来的延迟。第四，long-horizon 任务上的稳定性研究——例如 100+ 轮的 agent 任务中 Distiller 的记忆何时会"过拟合"到某个 prompt 模板，是否需要周期性重置或回放机制。第五，将 novelty 信号的粒度从 per-token 升级到 per-phrase 或 per-paragraph，使多样性在更宏观的推理结构层面发生作用。

复现评估整体友好但存在中等门槛。开源情况：代码已发布在 https://github.com/LinesHogan/tLLM，包含 ESamp 完整算法、tLLM 运行时框架、Triton kernel 后端和 vLLM v1 集成；附录 D 详细描述了工程优化（CUDA graph、post-filter candidate intervention、Triton grouped prediction backend）。数据方面，所有 benchmark 均为公开标准集（AIME 2024/2025、GPQA-Diamond、LiveCodeBench v5、BookCorpus 51,442 本书），评估用 lighteval 框架的官方 prompt。算力门槛：主要实验在 NVIDIA A100 GPU 上跑（AIME 等数学题最大上下文 8192），吞吐基准测试在单卡 RTX4090 上完成；Table 20 显示 Qwen2.5-7B + batch=8 + $n=16$ 时吞吐量 5304.855 tok/s（对比优化 vLLM 基线 5370.616 tok/s，比值 0.9878），意味着在普通消费级 GPU 上即可复现主结果。训练配置：Distiller 是 2 层 Gated SwiGLU MLP、hidden dim 384、Adam（$\eta=4\times 10^{-4}$、$\epsilon=1\times 10^{-4}$、gradient clip 0.5），超参 $\beta=0.25$ 是稳健默认（Table 4 显示在 Qwen3-4B/8B/14B 上几乎都是最优）。复现难度：中等偏上——需要熟悉 vLLM v1 的 worker 内部机制、CUDA stream 编程和 vLLM 0.6+ 的 SpecDecodeWorker 接口；附录 B 详细描述了异步 pipeline 的三大支柱（CUDA streams、pre-allocated ring buffer、mixed-batch guardrails），但实际部署仍需较强的系统工程能力。附录 E.2 注明"the code will be released upon the acceptance of this paper"——论文目前为 preprint 状态，仓库应已开放；附录 E.5 列出了所有 benchmark 的 prompt 模板，方便对齐。