【博士论文】《通过提前退出算法加速大语言模型推理》

大语言模型已经展现出非凡的能力，但其实际部署仍受制于巨大的计算成本。尽管诸如提前退出（early-exiting）之类的自适应计算方法有望降低这些成本，但它们也引入了一个根本性的矛盾：为了节省计算而在每个 token 级别引入的动态性，往往会在批量推理中造成系统级瓶颈，反而可能降低吞吐量。本论文通过自适应算法与模型架构的协同设计，解决了动态性与效率之间的这一冲突，以实现最优平衡。为此，我们首先针对传统提前退出方法中的关键开销源，提出了一种高效的并行解码机制。接着，我们证明了深度参数共享不仅能够构建紧凑、参数高效的模型，还能够从根本上缓解动态推理中关键的同步问题。最后，我们提出了一个统一框架，在其中轻量级路由器通过预训练，能够为每个 token 动态分配最优的递归深度。该方法有效地在单一模型内同时优化自适应计算与参数效率，建立了效率与性能之间新的帕累托前沿。 关键词：提前退出，自适应计算，大语言模型，推理效率，并行解码，参数共享，路由机制 大语言模型（LLMs）已经在诸如代码生成与逻辑推理等广泛的基准测试上展现出可与人类相媲美，甚至超越人类的能力。然而，实现这一水平的性能伴随着巨大的计算代价，这主要源于其庞大的参数规模以及自回归解码的固有低效性。在自回归生成过程中，模型每次只能生成一个 token，并且每一步都需要激活整个模型的计算图，从而在推理中造成严重瓶颈。 为提升 LLM 的推理效率，已有研究利用了其计算冗余这一关键洞察：并非所有参数在生成每个 token 时都必不可少。基于此，自适应计算成为了一种有前景的思路，它根据输入需求动态分配计算资源。在计算路径中引入动态性是所有自适应方法的核心。其中，**提前退出（early-exiting）**框架因其“简单而有效”的原理而备受关注：如果模型在某个中间层对预测结果已有足够信心，则可以提前终止推理并直接输出结果，而无需经过剩余层。该框架通过为不同输入创建深度不一的计算路径，从而提升计算效率。 然而，讽刺的是，虽然动态计算路径本意是优化推理，但其过度的动态性以及朴素实现方式却会复杂化批量推理，有时甚至降低实际吞吐量。这一问题主要源于两个阻碍提前退出算法在真实场景部署的核心挑战： 1. 缺失 KV cache 问题：提前退出的 token 会跳过剩余层的 KV（Key-Value）对计算，而这些信息可能在后续 token 解码时仍然需要。若对其进行近似计算，性能会下降；若精确计算，则会削弱自适应带来的效率收益。 1. 批量同步问题：已退出的 token 被迫等待同一 batch 中其他 token 完成深层的前向计算，因为位于不同网络深度的 token 通常无法被正常批量处理。

因此，本论文通过自适应算法与模型架构的协同设计，在动态性与效率之间寻求最优平衡，从而解决这两个主要的推理挑战，同时保持原始模型的性能。 在 第 2 章，我们首先针对单一 batch 场景下的缺失 KV cache 问题提出了并行解码机制。该方法将连续提前退出 token 的隐藏状态累积起来，并与下一个未退出的 token 一起，在一个并行化的前向传播中计算。该机制利用了 Transformer 解码过程的内存带宽受限特性（即延迟主要由内存访问而非计算主导），因此几乎不增加额外开销：它避免了额外的内存访问，仅引入了极小的计算成本（相较于现代 GPU 的 FLOPs 能力可以忽略不计）。最终，该框架通过高效地在深层中精确计算退出 token 的 KV 对，实现了无损生成，同时提升了推理效率。 在 第 3 章，我们通过探索并优化 参数共享 来缓解批量同步问题。这使得连续深度批处理成为可能——一种新颖的推理范式：即便 token 处于不同的模型深度阶段，也能在同一批次中处理，因为它们使用相同的参数块。我们证明该方法在解决同步瓶颈时能够带来显著的理论加速。然而，这一方法也带来了新的挑战：它依赖于一个敏感的两阶段训练过程（先进行参数共享，再进行提前退出训练）；此外，它并未解决缺失 KV cache 问题，仍需结合并行解码机制，从而可能导致实现复杂度增加。 在 第 4 章，我们提出了一个统一框架，能够同时解决提前退出中固有的缺失 KV cache 与批量同步两大挑战。其核心是引入路由器（routers），通过端到端训练为每个 token 动态分配最优递归深度，从而仅在必要时分配计算资源。这不仅消除了以往两阶段训练的复杂性，还保持了模型性能并证明其可扩展性。同时，我们提出了一种递归层级的 KV 缓存机制，仅存储在各递归深度下被实际处理的 token 的 KV 对。该机制不仅解决了缓存缺失问题，还显著降低了内存占用。通过系统性地解决前几章的局限，这一统一架构在效率与性能之间建立了新的帕累托前沿。 总体而言，本论文通过自适应算法与新型模型架构的系统性协同设计，提出了一个统一框架，解决了阻碍动态与自适应 LLM 广泛部署的关键挑战。我们相信，本论文所提出的原理将成为提前退出在真实世界应用中更广泛采用的重要基石。