跳过一层还是循环它？学习LLM中的Program-of-Layers

# 跳过还是循环？在LLM中学习层程序
来源: https://arxiv.org/html/2606.06574
###### 摘要

大型语言模型 \(LLMs\) 通过遵循固定的深度和顺序——即所有层的非循环执行——来进行推理。我们揭示了无训练、灵活、动态的“层程序 \(PoLar\)”的广泛存在性，其中预训练层可以作为模块打包，然后被跳过或循环，从而为每个输入定制专属程序。对于大多数输入而言，使用显著缩短的程序执行就能达到相同甚至更高的精度，而原始LLM的错误预测则可以通过使用更少层的替代程序来纠正。这些观察表明，推理过程存在多种有效的潜在计算路径，超越了标准的前向传播。为了在实践中高效实现PoLar，我们提出了一个轻量级的PoLar预测网络，该网络学习为每个输入生成动态跳过或重复预训练层的执行程序。在数学推理基准上的实验表明，PoLar在精度上持续优于标准推理和先前的动态深度方法，同时往往执行更少的层，并且这些优势在分布外评估下依然保持。我们的结果表明，固定深度的执行仅捕获了LLM潜在推理能力的一小部分。

机器学习, ICML

## 1 引言

参见图注 图1:针对两个不同输入的层程序 \(PoLar\)。一个预训练LLM中的DD个层定义了DD个函数f0,...,fD−1\{f\_\{0\},…,f\_\{D\-1\}\}。PoLar的动态推理不是按照从f0f\_\{0\}到fD−1f\_\{D\-1\}的静态固定顺序调用它们，而是执行一个*特定于输入的**程序π=\(i1,...,iK\)\\pi=\(i\_\{1\},…,i\_\{K\}\)，该程序通过*层跳过*和*层循环*的方式调用函数。PoLar为不同输入实现了无需训练的动态深度架构，产生了现有方法无法完全涵盖的多样化潜在计算路径。通用基础模型，例如LLMs和VLMs，对所有输入统一采用静态、预定义的架构，尽管这些输入在复杂性和难度上存在巨大差异和高度变化\(Liuet al\.,2020; Xinet al\.,2020; Zhouet al\.,2020; Liuet al\.,2021a\)。相比之下，传统的问题求解程序在算法结构和复杂度上可以更加灵活和自适应。例如，经验丰富的程序员可以在简单任务上节省更多步骤和计算，同时知道如何扩大时间/空间复杂度来解决更具挑战性的问题。然而，这些程序是针对每个问题类别专门设计和优化的，因此不如LLM通用。这引出了一个问题：对所有不同的任务应用相同的架构或“程序”，即按照固定顺序通过所有层的前向传播，是否总是最优且高效的？通用模型能否进一步优化应用于每个输入的“程序”？

在本文中，我们将预训练LLM中的层形式化为一个原子函数库，程序可以以任意顺序、任意次数调用这些函数。这种形式化使我们能够在推理期间，为每个输入将动态模型架构表示为一个层程序 \(PoLar\)，如图1所示。作为此类研究的首次实证探索，我们通过蒙特卡洛树搜索 \(MCTS\) 研究了超越标准前向传播的PoLar，并发现对于每个评估的输入任务，几乎总是存在更好（更准确和/或更短）的程序。与先前关于层跳过/循环、提前退出和循环Transformer的研究\(Liuet al\.,2020; Xinet al\.,2020; Zhouet al\.,2020; Fanet al\.,2019,2024; Yanget al\.,2023\)不同，这些研究仅采用一种操作（要么跳过要么循环）来产生动态深度架构，而我们对MCTS搜索到的程序的实证研究表明，在层跳过/循环的联合空间中进行搜索，往往能发现比在独立空间中搜索要好得多的程序。虽然大多数有效程序可能比默认程序更短，但通过跳过/循环操作增加程序复杂度可以显著提高输出质量，尤其是在更困难的任务上。此外，大多数成功的程序主要由连续的层段组成。这些观察不仅验证了无需任何训练即可广泛存在更好的PoLar，也激发了一种实用的PoLar预测方法，该方法避免了在PoLar的大搜索空间中使用MCTS带来的高昂成本。具体来说，我们的目标是用一种直接的、推理时的机制来生成执行程序，取代基于搜索的程序发现。不是为每个输入枚举或探索执行路径\(Liet al\.,2025\)，我们的目标是预测一个特定于输入的层程序，该程序决定了推理过程中如何执行预训练层。这将程序选择问题从在线搜索问题转变为单次预测问题，从而使得层程序推理能够在LLMs中实际部署。

为此，我们提出了一种PoLar算法，该算法在推理时针对冻结的预训练层预测执行程序。预测的执行程序指定了如何选择性地跳过或循环应用预训练层，并执行一次以产生最终输出。这种设计具有几个优点。首先，它通过消除昂贵的每输入搜索需求，使得层程序推理在计算上可行。其次，通过在统一的执行框架内共同支持层跳过和循环，PoLar严格推广了先前仅限于单一形式执行控制的动态深度方法。第三，它实现了完全冻结模型中的灵活测试时计算缩放，允许推理适应输入难度，同时保持模型通用性。

我们使用多个预训练LLM在一系列数学推理基准上评估了PoLar。我们的结果表明，PoLar在精度上持续优于标准推理和先前的动态深度方法，同时平均执行更少的层。此外，增加候选执行程序的数量实现了强大的测试时计算缩放，并且在分布内数据上学到的执行程序能有效泛化到跨不同领域的分布外基准。

## 2 大型语言模型中的动态推理作为层程序 \(PoLar\)

预训练LLM的推理实现为一个固定深度、固定顺序的前向传播：每个输入都通过执行相同的Transformer层序列来处理。然而，LLM的输入在难度上差异巨大。有些输入只需极少的推理就能正确回答，而另一些则需要复杂、多步的计算。这种差异引出了一个基本问题：标准的前向传播对于不同输入的正确推理是否足够？

一种可能性是，这种固定的计算对于所有情况确实足够。另一种可能性是，正确的预测需要根据输入变化计算。在这项工作中，我们研究后一种可能性。这种变化可以发生在token空间，通过更长、更明确的思维链，或者发生在模型的隐藏状态内部，这是一种我们称之为*潜在推理*的计算形式。

猜想：推理即层程序将预训练LLM中的层定义为函数，对于每个输入，可能存在多种不同的层程序执行方式（超越标准前向传播）来产生正确的预测。

推理即执行一个*程序*。在这种观点下，推理是一个逐步的过程，选择和组合预训练模块。执行过程可能在长度和顺序上都因输入而异，而每个模块仍然是一个固定的、预训练的函数。

考虑一个具有DD个Transformer层的预训练LLM，其中每一层定义了一个固定的计算函数

fi:RT×d→RT×d,i∈{0,...,D−1}\.\\textstyle f\_\{i\}:\\mathbb\{R\}^\{T\\times d\}\\rightarrow\\mathbb\{R\}^\{T\\times d\},\quad i\\in\\\{0,\\ldots,D\-1\\\}\.一个程序被定义为一个有限的层索引序列

π=\(i1,i2,...,iK\),ik∈{0,...,D−1\},\\textstyle\\pi=\(i\_\{1\},i\_\{2\},\\ldots,i\_\{K\}\),\quad i\_\{k\}\\in\\\{0,\\ldots,D\-1\\\},该序列诱导出组合计算

Fπ=fiK∘⋯∘fi1\.\\textstyle F\_\{\\pi\}=f\_\{i\_\{K\}\}\\circ\\cdots\\circ f\_\{i\_\{1\}\}\.执行一个程序就是将这个组合应用于输入并产生一个预测。如果一个程序能对给定输入产生正确预测，则被认为是*有效*的。

参见图注 图2:用于程序预测的顺序MCTS（左）vs. 端到端PoLar网络（右）。(a) 在执行程序空间中通过选择、扩展、模拟和反向传播的顺序迭代进行MCTS。每个节点代表一个部分或完整的执行程序，跳过/循环操作迭代地扩展搜索树。这种显式且彻底的搜索成本高昂且不实用。(b) 我们的PoLar训练了一个端到端、轻量级的预测网络，该网络直接生成一个程序表示，该表示由 (i) 一个将层分割成模块的二元掩码zseg\(x\)\\mathbf\{z\}^\{seg\}\(x\)，以及 (ii) 一个操作标签向量zop\(x\)\\mathbf\{z\}^\{op\}\(x\)组成，该向量对每个模块应用*跳过*、*保留*或*循环*中的一种操作。我们的方法在实践中具有可扩展性，且不需要顺序搜索。参见图注 图3:在DART\-Math五个难度级别上，不同执行深度预算下MCTS发现程序的准确率。我们将原始前向传播（橙色）与90–115%深度预算的程序（蓝色）进行比较。阴影区域表示在最高预算（115%）下达到的最大增益。搜索有效的执行程序。我们使用MCTS探索执行程序的空间。执行程序是预训练Transformer层上的变长序列，允许跳过和重复。这个空间巨大、离散且高度非凸，使得穷举搜索不可行。MCTS为优先考虑有希望的部分程序提供了一种原则性方法，使我们能够验证有效程序的存在并分析其结构特性。我们严格将MCTS用作诊断工具，而非实际的推理时方法；实现细节见附录B。所有实验均在DART\-Math\(Tonget al\.,2024\)上进行，这是一个结构化的数学推理基准，包含五个难度级别（DM\-1至DM\-5）。我们评估了四个预训练Transformer模型：LLaMA\-3\.2\-3B\-Instruct、Qwen1\.5\-MoE\-A2\.7B\-Chat、Qwen2\.5\-3B\-Instruct和Qwen3\-8B。

下面总结我们的实证发现。

发现1仅层循环的效果优于仅层跳过，但结合这两种互补操作可以产生最佳的层程序。

如表1所示，允许层循环（Loop）的执行程序在所有评估的模型和难度级别上，始终优于仅允许层跳过（Skip）的程序。此外，将跳过与循环结合（Skip&Loop）产生的增益远大于单独使用任一操作，在表1报告的每个设置中都达到了最高准确率。这些结果表明，循环提供了一种比单独跳过更强的推理改进机制，而这两种操作在结合使用时发挥着互补作用。

参见图注 图4:潜在执行程序通常存在更短的有效解。我们将标准前向传播深度与MCTS发现的有效程序的深度进行比较，针对最初正确（C→\rightarrowC）和最初错误（W→\rightarrowC）的输入。条形图报告了总执行深度占完整模型深度的百分比，带阴影叠加的部分表示有效深度（唯一层的数量）。参见图注 图5:(a)通过层段循环实现测试时计算缩放。允许通过段循环进行更多潜在执行步骤，导致跨模型发现有效执行程序的概率单调增加。 (b)对于更难的输入，对循环和跳过的需求增加。对于大多数模型，依赖于层循环或跳过才能解决的输入比例随难度增加而增加，但LLaMA\-3\.2\-3B\-Instruct除外，其偏差可由表1中的有效难度解释。参见图注 图6:准确率 vs. 总层执行次数。对于每个模型，我们报告有效执行程序的平均准确率如何随着总层执行次数（占基础模型深度的百分比）变化。在模型和难度级别中，准确率随执行层数增加而增加，揭示了深度缩放的一致效应。参见图注 图7:有效执行程序的结构偏差。有效程序主要依赖于连续的层段作为模块 (a)，并且每个模块最多需要一次循环 (b)。表1:在不同搜索空间中DART\-Math（难度1–5）的程序准确率（%）：Base（标准前向传播）、Skip（层跳过）、Loop（层循环）和Skip&Loop（跳过+循环）。Gain报告了Skip&Loop相对于Base的绝对增益。**粗体**表示每行最佳结果，下划线表示第二佳结果。指标BaseSkipLoopSkip&LoopGainLLaMA\-3\.2\-3B\-InstructDM\-137\.945\.754\.984\.7\+46\.8DM\-228\.134\.846\.872\.3\+44\.2DM\-323\.229\.738\.065\.2\+42\.0DM\-422\.828\.235\.257\.0\+34\.2DM\-527\.131\.739\.059\.1\+32\.0Qwen1\.5\-MoE\-A2\.7B\-ChatDM\-137\.442\.757\.773\.2\+35\.8DM\-226\.432\.643\.759\.7\+33\.3DM\-320\.023\.834\.450\.4\+30\.4DM\-413\.916\.425\.138\.2\+24\.3DM\-510\.913\.320\.432\.4\+21\.5Qwen2\.5\-3B\-InstructDM\-125\.447\.060\.287\.4\+62\.0DM\-211\.233\.644\.376\.5\+65\.3DM\-34\.325\.135\.565\.0\+60\.7DM\-42\.015\.822\.251\.2\+49\.2DM\-51\.213\.218\.144\.5\+43\.3Qwen3\-8BDM\-140\.766\.068\.591\.3\+50\.6DM\-229\.350\.657\.482\.2\+52\.9DM\-318\.536\.040\.967\.1\+48\.6DM\-410\.423\.829\.253\.6\+43\.2DM\-58\.820\.023\.845\.7\+36\.9

发现2 (奥卡姆剃刀) 大多数有效的执行程序通常比标准前向传播更短。

如图3所示，该图报告了在总层执行次数有明确预算下MCTS获得的最佳准确率，即使当整体计算被限制为显著短于标准前向传播时，许多输入仍然可解。与此趋势一致，图4显示，在我们发现的模型中，经常存在比标准推理需要更少层应用的有效执行程序。特别是，在那些已经通过标准推理正确解决的输入（C→\rightarrowC）中，75.5%存在更短的有效程序。即使对于最初错误解决的输入（W→\rightarrowC），36.2%也存在能够修正模型预测的更短程序。这些结果表明，标准推理通常过度计算，而正确推理通常可以通过显著更少