🤖 machine learning

TIDE: Token-Informed Depth Execution for Per-Token Early Exit in LLM Inference

TIDE 是一种无需重新训练即可适配任意 HuggingFace 因果语言模型的推理后处理系统，它通过在周期性检查点层附加轻量级路由器，实现针对每个 Token 的自适应早期退出，从而在保持输出质量的同时显著降低推理延迟并提升吞吐量。

原作者： Jaber Jaber, Osama Jaber

发布于 2026-03-24

📖 1 分钟阅读☕ 轻松阅读

原作者： Jaber Jaber, Osama Jaber

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在指挥一支拥有 32 层楼的超级智能工厂（这就是大型语言模型，LLM）。每当有人问一个问题，工厂里的每一个“工人”（也就是每一个生成的字或词，比如“的”、“是”这种简单的词，或者是复杂的数学推导步骤）都要从 1 楼一直走到 32 楼，完成所有的工作，最后才能把答案交给你。

这听起来很公平，但仔细想想有点浪费：

当工人处理像“的”这样简单的词时，其实走到第 8 楼就已经想明白了，剩下的 24 层楼纯属“磨洋工”。
但当处理复杂的数学题时，确实需要走到第 32 楼才能得出正确答案。

TIDE 系统（Token-Informed Depth Execution）就是为了解决这个“一刀切”的浪费问题而发明的。它就像给每个工人发了一副智能眼镜和一张特制的通行证。

TIDE 是如何工作的？（三个简单步骤）

1. 岗前培训（校准阶段）

在工厂正式开工前，TIDE 会先让工厂“热身”一下。它用一些普通的文章（比如维基百科上的短文）让模型跑一遍，观察每个字在每一层楼的状态。

核心发现：TIDE 发现，对于大多数简单的字，它们的状态在走到某一层时，就已经和走到第 32 楼时的状态几乎一模一样了（就像你穿好衣服出门，走到小区门口和走到公司，衣服其实没变）。
训练“小考官”：TIDE 在特定的楼层（比如第 8 层、第 16 层、第 24 层）安装了一些极小的“小考官”（路由器）。这些小考官的任务是看一眼工人手里的“状态卡”，判断：“嘿，这个字已经想明白了，不需要再往上走了！”

2. 智能放行（推理阶段）

现在，工厂开始正式处理用户的请求了。

全层运行，但“心里有数”：为了保持系统的稳定性（不弄乱记忆缓存），TIDE 依然会让所有字走完所有楼层。
关键时刻的“截停”：但是，当字走到某一层时，那个楼层的“小考官”会立刻检查。如果小考官说：“这个字（比如‘的’）已经懂了，分数够了！”系统就会立刻截停这个字，直接把它送到出口，不再让它去后面那些没必要的楼层。
结果：
- 简单的词（如“的”）可能只走了 8 层就下班了。
- 复杂的词（如数学推导中的关键步骤）会老老实实走到第 32 层。
- 就像坐电梯：以前所有人必须坐到顶楼再下来；现在，想去 8 楼的人，电梯到了 8 楼就开门放人，不用管后面还有没有 24 层。

3. 无需重造工厂（无需重新训练）

这是 TIDE 最厉害的地方。以前的方法可能需要把整个工厂拆了重建（重新训练模型），或者给工厂装个复杂的“草稿模型”（推测解码）。

TIDE 不需要动工厂的砖瓦（模型参数），也不需要重新培训工人。
它只是给现有的工厂加了一副“智能眼镜”（那些微小的路由器），而且这副眼镜非常便宜，训练只需要 3 分钟，文件只有 4MB 大小。

实际效果如何？

根据论文测试，在 NVIDIA A100 显卡上：

速度提升：处理输入（预填充）的速度快了约 7%，生成速度（吞吐量）提升了 6% 到 8%。
准确率：即使让 99% 的简单字提前“下班”，模型在解复杂的数学题时，依然能保持 100% 的正确率。
通用性：它像一个万能适配器，无论是 LLaMA、Qwen 还是其他任何 HuggingFace 上的模型，插上就能用。

总结：TIDE 是什么？

如果把大模型比作一个过度热情的导游，以前不管游客问“今天天气好吗”还是“如何造火箭”，导游都要把整个博物馆（32 层楼）从头到尾讲一遍。

TIDE 就是给导游装了一个智能判断系统：

如果游客问“天气”，导游讲到第 8 层发现游客已经懂了，就立刻说：“好嘞，答案就在这，咱们直接出门吧！”
如果游客问“造火箭”，导游就会继续讲到第 32 层，确保细节无误。

TIDE 的核心价值：它让大模型变得更聪明、更懂分寸，不再做无用功，从而在保持聪明的同时，跑得更快、更省电。而且，它不需要你重新训练模型，就像给旧手机装了一个新的省电 APP 一样简单。

TIDE: 大语言模型推理中的 Token 级早期退出技术总结

1. 研究背景与问题 (Problem)

现有的大型语言模型（LLM）在推理过程中，无论 Token 的语义难度如何（例如是简单的功能词"the"还是复杂的数学推理步骤），都会强制通过模型的所有层。这种均匀的计算分配导致了巨大的资源浪费：

计算冗余：研究表明，对于大量 Token，其中间隐藏状态在到达最后一层之前就已经收敛（与最终隐藏状态高度相似）。
成本高昂：在大规模部署中（如 70B 参数模型），这种冗余直接转化为更高的延迟、更低的吞吐量以及更高的能源消耗。
现有方案局限：
- 编码器模型方法（如 DeeBERT）：仅适用于分类任务，无法处理带有 KV Cache 的自回归生成。
- 预训练早期退出（如 LayerSkip）：需要重新训练模型，耗时耗力，不适合直接使用预训练检查点的用户。
- 基于置信度的启发式方法：利用 Softmax 熵作为退出信号，但在生成任务中熵天然较高，导致信号不可靠。

2. 方法论 (Methodology)

TIDE (Token-Informed Depth Execution) 提出了一种无需重新训练模型的后训练（Post-training）系统，通过在每个检查点层附加轻量级路由器（Router），在推理时动态选择每个 Token 的最早收敛层。

核心流程

TIDE 分为两个阶段：

A. 离线校准 (Offline Calibration)

数据收集：使用 2,000 个 WikiText 样本，在冻结的预训练模型上运行，收集每个检查点层（Checkpoint Layers）和最终层的隐藏状态。
收敛标签生成：计算每个 Token 在检查点层 $k$ $k$ 的隐藏状态 $h_k$ $h_{k}$ 与最终层 $h_L$ $h_{L}$ 的余弦相似度。
- 公式： $s_k = \cos(h_k, h_L)$
- 标签：若 $s_k > \tau$ （默认 $\tau=0.98$ ），则标记为已收敛（Label=1）。
路由器训练：为每个检查点层训练一个轻量级二分类 MLP 路由器（Router）。
- 结构：两层 MLP，瓶颈维度 128，输入为隐藏状态，输出为是否退出的概率。
- 耗时：单 GPU 上仅需 3 分钟，模型大小仅约 4MB。

B. 在线推理 (Online Inference)

完整前向传播：在自回归生成过程中，模型仍然运行所有层，但开启 output_hidden_states=True 以保留 KV Cache 的完整性。
后验退出评估 (Post-hoc Evaluation)：前向传播完成后，路由器并行评估每个检查点层的隐藏状态。
选择退出层：对于每个 Token，选择第一个满足路由器阈值 $\theta$ 的层 $k$ 。
** logits 计算**：直接使用该层 $k$ 的隐藏状态（经过 RMSNorm）计算 Logits，而非最后一层。

技术亮点

通用适配器 (Universal Adapter)：自动探测 17 种属性路径，支持 LLaMA, GPT-2, Qwen, Phi, Falcon 等几乎所有 HuggingFace 架构，无需针对特定模型编写代码。
融合 CUDA 内核：实现了 RMSNorm + 路由器评估的融合内核，支持 FP16/BF16，针对常见隐藏维度（2048-8192）进行了模板特化，显著减少 Kernel 启动开销。
KV Cache 完整性：由于所有层都执行了，KV Cache 不会被破坏，解决了早期退出方法中常见的缓存污染问题。

3. 主要贡献 (Key Contributions)

首个无需修改模型的自回归 LLM 早期退出系统：适用于任何 HuggingFace 因果语言模型。
通用架构适配：自动支持 17 种主流模型架构，无需手动适配。
高性能 CUDA 实现：提供融合内核，支持从 V100 到 Blackwell 的 GPU 自动检测。
后验退出策略：在保持 KV Cache 完整性的同时实现 Token 级深度优化。
开源与实证：发布了包含 74 个测试用例的开源包（tide-inference），并在 DeepSeek R1 和 Qwen3 上验证了效果。

4. 实验结果 (Results)

在 NVIDIA A100 上的实验表明：

退出率 (Exit Rates)：
- Prefill 阶段：100% 的 Token 都能找到退出点。在 DeepSeek R1 Distill 8B 上，5% 的 Token 在第 11 层退出，其余在第 31 层退出。
- Decode 阶段：98%–99% 的 Token 能提前退出。
性能提升：
- 延迟：DeepSeek R1 8B 的 Prefill 延迟降低了 5.5% - 7.2%。
- 吞吐量：
  - DeepSeek R1 8B (Batch Size 1)：提升 6.6%。
  - Qwen3 8B (Batch Size 8)：提升 8.1%。
  - 注：在 Batch Size 8 时，DeepSeek R1 吞吐量略有下降，归因于 output_hidden_states 开销随 Batch 增大呈超线性增长。
生成质量：
- 在数学推理任务中，即使 99.6% 的 Token 提前退出（仅在第 31 层），模型仍能正确解决多步数学问题，且唯一 Token 数量与基线模型一致，无质量损失。
资源消耗：校准过程仅需 3 分钟，生成的路由器检查点仅 4MB。

5. 意义与局限性 (Significance & Limitations)

意义

即插即用：用户无需重新训练昂贵的预训练模型，即可通过加载轻量级路由器获得性能提升。
Token 级粒度：打破了“模型级”或“序列级”优化的限制，实现了真正的 Token 级自适应计算。
生态兼容：完美兼容现有的 transformers 库和 KV Cache 机制，易于集成到现有推理服务中。

局限性与未来工作

非物理跳过：目前 TIDE 采用“后验模式”，即所有层仍会执行计算，只是最后选择哪一层的输出。这意味着没有真正减少物理计算时间（Wall-clock time），仅减少了部分归一化和投影的计算开销。真正的“物理跳过”需要解决 KV Cache 不连续的问题，这是未来的方向。
阈值保守：默认阈值 $\tau=0.98$ 过于严格，导致大多数 Token 集中在倒数第二层退出。未来可通过动态调整阈值来挖掘更深层的退出机会。
大 Batch 瓶颈：在超大 Batch 下，收集所有隐藏状态的开销可能成为瓶颈，未来可考虑仅在检查点层收集隐藏状态。

总结

TIDE 证明了在预训练模型基础上，通过轻量级学习路由器实现 Token 级早期退出的可行性。它在保持生成质量不变的前提下，显著降低了推理延迟并提升了吞吐量，为大模型的高效部署提供了一种低成本、高兼容性的新范式。