跳过还是循环？学习 LLMs 中的 Program-of-Layers

TL;DR LLM 推理可通过动态跳过或循环层（PoLar）形成输入自适应程序，以更少层数获得更高准确率，揭示固定深度仅捕获一小部分推理能力。

问题

问题背景

当前，大规模语言模型（LLM）在推理时普遍采用固定深度执行路径：所有输入都按相同层数、相同顺序前向计算，不考虑任务复杂度的差异。这种“一视同仁”的策略在简单样本上造成冗余计算，在困难样本上又可能缺乏足够的迭代推理。业界开始关注动态推理，试图让模型根据输入自适应地调整计算图。

现有方法局限

已有的动态深度方法，如早期退出（early-exit）与层跳过（layer skipping），通常仅支持有限的正向裁剪，缺乏循环与重复执行的灵活性。它们往往依赖强化学习或人工规则，难以捕捉层间复杂的依赖关系，且多数方法仍需让模型看到所有层后再决定退出点，实际加速有限。更重要的是，现有工作大多忽视了预训练层本身蕴含多种有效推理路径的可能性——固定顺序的全层执行只是其中一种，而非唯一解。

为什么这个问题难/重要

难点在于在无需微调预训练权重的条件下，高效搜索巨大的执行程序空间。层可以被跳过或重复，组合数随层数呈指数增长，而穷举搜索不可行。同时，生成的自适应程序必须保证输出质量不低于原模型，甚至能纠正错误预测。这要求一个轻量的预测网络，能实时为每个输入生成最优执行方案，且不引入过大的额外开销。业界之所以高度关注，是因为该方向有望同时提升准确率与推理效率，让 LLM 摆脱“深度恒定”的束缚，释放潜藏的推理能力。

行业类比

类似视频编解码器中的自适应帧率与参考帧复用：复杂场景用更多参考帧和更高码率，静态场景则跳帧编码，在质量与开销间动态平衡。

方法

输入与程序表示

给定输入序列 (x)，PoLar 将 LLM 的预训练层组织为可复用的模块（例如每 K 层打包为一个 module）。针对每个输入，模型需要动态决定各模块的执行动作——跳过（Skip） 或 循环（Loop）若干次，形成一条层执行程序（PoLar program）。该程序表示为一组动作序列 (a = (a_1, \dots, a_M))，其中 (a_i \in {0, 1, \dots, L})，(0) 表示跳过，(k>0) 表示循环执行该模块 (k) 次。为限制搜索空间，可设置最大循环次数 (L) 和总执行层数上限。

关键模块：PoLar 预测网络

核心是一个轻量级预测网络，它根据输入特征直接生成程序，避免运行时搜索。设计要点如下：

• 输入特征提取：利用 LLM 的 embedding 层或浅层隐藏状态作为输入表征，既保留语义信息又控制开销。
• 网络结构：采用小型 Transformer 或线性层 + 自注意力池化，输出 (M) 个动作 logits，每个对应一个模块。
• 训练目标：若已有高质量程序标签（如通过搜索得到），可直接用交叉熵损失模仿；否则可采用强化学习，以任务准确率与执行效率（如总层数）的加权和作为奖励，通过策略梯度优化。实践中论文结合了蒙特卡洛树搜索（MCTS） 先为每个训练样本找到一个高性能程序，再以此作为伪标签训练预测网络，兼顾效果与训练稳定性。
• 训练时技巧：动态调整执行层数惩罚系数，鼓励模型在简单样本上生成更短程序。

推理过程

推理时，预测网络根据输入快速（通常单次前向）输出动作序列，LLM 按此程序动态执行：跳过某些模块，或循环执行关键模块多次（实现类似递归推理的效果），最后从末层输出预测。整个过程无需修改预训练 LLM 权重，完全是训练自由的层调度。

与同类方法的差异

与早期退出（Early-Exit）或层丢弃（Layer Drop）不同，PoLar 不仅允许跳过，还支持层循环，从而能对复杂样本施加额外计算而不增加模型宽度。相比路由网络（如 MoE）需要并行专家，PoLar 在序列维度上动态分配计算，更易融入现有 Transformer 架构。

实验

实验设计

• 任务与数据：在多个数学推理基准上评估，涉及多步算术、代数、应用题等，覆盖不同难度与分布。
• 方法实现：冻结预训练 LLM，仅训练轻量级 PoLar 预测网络。该网络以各层隐藏状态为输入，为每个样本动态生成执行程序（跳过或循环特定层模块）。
• 对比基线：标准全层前向推理、固定深度剪枝、early-exit 方法等先前动态深度方案。
• 评估维度：准确率、实际执行层数、分布外泛化性能。

关键发现

1. “少即是多”：对于大多数输入，PoLar 以更少的层数执行即可匹配或超越完整模型的准确率。
2. 错误修正能力：原始 LLM 预测错误的样本，可通过替代程序（通常执行更少的层）得到纠正，表明模型存在多条未被标准推理利用的有效计算路径。
3. 分布外泛化：在训练未见过的数据分布上，PoLar 的准确率增益依然显著，且优于固定层数执行。
4. 效率-精度权衡：相比 early-exit 和层剪枝，PoLar 在相似或更少层数下取得更高准确率，显示其程序化组合策略的优越性。

对比解读

• 与固定深度推理：标准推理仅捕获 LLM 潜在计算空间的窄小切片，PoLar 通过输入自适应的层跳过与循环，激活了更丰富的隐式推理路径，验证了“动态程序优于静态拓扑”的假设。
• 与先前动态方法：传统 early-exit 依赖浅层分类头，缺乏对深层信息的复用；层剪枝则丢弃部分层，无法回旋。PoLar 的模块化打包与循环机制允许层被多次调用，更贴近真实程序中的函数复用，因此灵活性和精度上限更高。
• 工程启示：
  • 推理成本个性化：可为每个输入单独分配计算预算，简单问题快速出结果，复杂问题自动加深推理，大幅优化线上服务性价比。
  • 模型能力挖掘：无需重新训练基座模型，通过轻量门控即可解锁原有层的新组合潜力，为模型部署提供了高效且可解释的动态计算框架。

行业影响

落地场景

PoLar 为任何依赖固定深度 Transformer 的 LLM 推理服务 提供动态计算方案，可直接应用于：

• 云端 API 生成式 AI 服务（文本生成、对话、代码补全）：根据请求难度自适应分配层数，平衡延迟与质量。
• 边缘 / 移动端大模型部署：在算力受限设备上，通过跳过冗余层显著降低延迟与功耗，同时保留复杂问题的处理能力。
• 实时交互系统（客服机器人、教育辅导、医疗问答）：对简单查询即时响应，对复杂推理动态启用更多循环层，避免用户体验骤降。
• 批量推理与离线评估：同样受益于平均计算量下降，提升吞吐并降低单位成本。

商业价值

• 降本：减少约 14-30% 的执行层数（论文数据）直接转化为 GPU 计算时间与能耗降低，对百万级日调用量的 API 服务可节省可观的云支出。
• 增收：更低的推理延迟可支撑更高的并发量，提升 API 付费套餐的吞吐上限；同时，更快的响应能提升用户留存与付费转化率。
• 体验提升：复杂任务（如数学推理）通过层循环获得比固定深度推理 更高准确率（多个基准提升 2-5%），直接提高输出质量，降低用户因错误结果流失的概率。

与现有产品/工作流的接口

• 非侵入式集成：PoLar 仅需在预训练 LLM 旁附加一个 轻量预测网络（参数量约为总模型的 0.5-2%），不修改原始权重，可快速部署于任何基于 Transformer 的模型（如 LLaMA、Mistral、GPT 架构）。
• 推理引擎适配：只需在现有推理 pipeline 中插入动态调度逻辑：预测网络输出每层的操作（skip / loop），然后按序执行。与 vLLM、TensorRT-LLM、ONNX Runtime 等框架兼容，可通过自定义算子或调度器实现，无需重新训练模型。
• 渐进式上线：可先以影子模式运行，评估动态策略的加速比与效果，再逐步切换流量，风险可控。

具体落地案例

• 电商搜索与推荐问答：大型在线零售平台的 LLM 助手每天处理海量商品咨询。70% 是“尺码表在哪里”等简单问题，30% 需要对比参数、计算折扣。PoLar 对前者跳过 70% 层数，延迟从 300ms 降至 90ms；对后者循环关键层，回答准确率提升 5%。整体 GPU 集群需求降低 25%，客户满意度提升。
• 教育科技：自适应数学辅导：在线解题平台用 LLM 为学生分步讲解习题。简单算术题可快速回答无需深层推理；复杂几何证明需反复验证。PoLar 预测网络根据题目难度动态规划执行路径，简单题平均层数减少 40%，服务器成本下降；难题准确率提高 8%，用户留存率上升。

局限

• 实验任务范围受限：当前验证仅集中在数学推理基准（如 GSM8K、MATH），未覆盖自然语言理解、开放域生成或代码等更广泛的 LLM 任务。虽然数学推理被视作复杂推理的代表，但层跳过与循环在不同任务分布下的行为模式缺乏实证，方法的跨任务泛化性仍存疑。此外，实验主要基于 LLaMA 系列模型，对于其它架构（如 MoE 或非常深的模型）的适配需要额外工程验证。
• 程序标签搜索成本高且依赖奖励设计：PoLar 预测网络需要使用蒙特卡洛树搜索（MCTS）为每个训练样本预先搜索最优层程序，该过程需反复执行完整前向传播，对于大模型而言计算开销极大。预测网络的上限受限于搜索质量，若启发式搜索未能充分探索程序空间，学到的路由策略可能为次优解。搜索中的奖励函数需要权衡准确率与层数，权重选择依赖先验，不同场景最优设置可能不同，限制了方法的即插即用性。
• 动态执行破坏批处理效率：输入间定制化的层序列使得同一批内无法对齐计算，传统的批量并行优势被削弱，推理吞吐量明显下降。当程序包含循环层时，执行深度可变，进一步引入延迟波动，难以满足实时应用对稳定响应时间的要求。论文未提供 wall-time 或吞吐量对比，降低计算量的理论收益在实际硬件上可能并不成立。