LayerRoute: 通过 LoRA 微调实现智能体语言模型的输入条件自适应层跳过 LayerRoute: Input-Conditioned Adaptive Layer Skipping via LoRA Fine-Tuning for Agentic Language Models
LoRA 联合训练实现按输入类型自适应跳过层
前置知识
LoRA (Low-Rank Adaptation)
LoRA 是一种参数高效的微调方法，它将预训练模型的权重更新分解为两个低秩矩阵的乘积 $BA$，其中 $A \in \mathbb{R}^{r \times d_{in}}$ 和 $B \in \mathbb{R}^{d_{out} \times r}$，$r$ 是秩（通常远小于隐藏层维度）。在推理时，可以通过 $W' = W + \frac{\alpha}{r}BA$ 将低秩适配器合并到原始权重中，不增加推理开销。这种方法使得在仅增加极少参数的情况下就能有效适配新任务。
本文使用 LoRA 在保持主干模型冻结的情况下进行适配训练，这是 LayerRoute 能够快速训练（6.4 分钟）且仅使用 1.10M 参数（骨干网络的 0.22%）的关键基础。
Straight-Through Estimator (STE)
直通估计器是一种用于处理离散操作（如阈值化、量化）的反向传播技术。对于不可导的硬阈值操作 $g = 1[\sigma(s) > 0.5]$，STE 在前向传播中使用离散值（0 或 1），但在反向传播时让梯度直接通过连续的 sigmoid 输出 $\sigma(s)$，即 $\frac{\partial g}{\partial s} = 1$（当在阈值附近时）。这种技巧避免了训练和推理之间的不一致性。
LayerRoute 使用硬二进制门控来真正跳过层计算，STE 是使这种离散决策能够端到端训练的核心技术，确保了训练得到的路由策略与推理时的行为一致。
智能体语言模型
智能体语言模型是指能够自主执行多步骤工作流的 AI 系统，包括调用外部工具（如 API、数据库查询）、进行规划和推理、处理检索到的文档等。这类系统的特点是输入类型异质性强，包括结构化的工具调用（如 `search_database(table='products', filter='id=100023')`）和非结构化的规划步骤（如 '制定用户留存下降的优先级策略'）。
本文针对的正是智能体系统中这种计算需求异质性的问题，不同步骤类型对 Transformer 深度的需求差异巨大，这是 LayerRoute 优化效率的动机来源。
研究动机
现有的智能体语言模型推理系统存在严重的计算资源浪费问题。这类系统交替处理两种本质完全不同的步骤类型：结构化的工具调用和开放式的规划推理。工具调用通常是短序列，下一个 token 具有高度可预测性，困惑度低；而规划步骤则需要复杂的多步推理，输出空间广阔，困惑度高。然而，当前的推理系统对这两种类型采用完全相同的计算策略，将每个输入都路由通过所有 Transformer 层，消耗相同的计算资源。这种处理方式是低效的，因为对于高置信度的工具调用，残差流在早期就达到了稳定状态，深层网络对最终预测的贡献微乎其微，实际上是在浪费浮点运算（FLOPs）。例如，对于简单的函数调用，模型可能只需要前几层就足够生成正确的参数，但仍然会运行全部 24 层，造成了不必要的延迟和能耗。
本文的目标是本文的具体目标是开发一个参数高效的机制，使语言模型能够根据输入的具体特征动态调整其计算深度，特别是在智能体工作流中区分工具调用和规划步骤。目标是学习一个路由策略，对于深度处理冗余的输入选择性地跳过 Transformer 块，从而减少推理延迟和 FLOPs 消耗。至关重要的是，这种优化必须在保证模型质量的前提下实现，理想情况下还能通过适配提升性能。另一个目标是使用极少量的可训练参数和极短的训练时间来实现这一点，使其成为一个实用的、即插即用的解决方案，能够应用于现有的模型如 Qwen2.5-0.5B-Instruct，而无需昂贵的从头训练。
与已有工作不同的是，本文的独特切入点在于专门针对智能体工作流中步骤类型异质性导致的计算效率问题，这与通用语言建模的优化有明显区别。与 Mixture of Depths (MoD) 不同，MoD 在 token 级别进行路由且需要从头训练，LayerRoute 在序列级别进行路由，基于步骤类型进行决策，这使得它能够对预训练模型进行高效的后训练适配。现有的早期退出方法通常依赖于中间置信度分数，但不同时适配模型权重。LayerRoute 引入了一种新颖的组合：使用直通估计器（STE）的硬门控跳过连接与 LoRA 微调的联合训练。这使得路由策略能够与适配权重协同演化，直接通过语言建模目标发现层的重要性，而无需显式的任务类型标签，这是与监督路由方法的一个显著区别。
核心方法
LayerRoute 提出了一个轻量级的适配器架构，它为冻结的预训练 Transformer 骨干网络（具体为 Qwen2.5-0.5B-Instruct）的每个块增加了两个可训练组件：一个每层路由器和 LoRA 适配器。直觉是不同的输入需要不同的深度；结构化的工具调用可能只需要浅层处理，而复杂的规划则需要完整的深度。该方法的工作原理是训练一个小型的线性路由器来检查进入每个块的均值池化隐藏状态，并输出一个二进制门。如果门关闭（0），则该块被完全跳过，隐藏状态原样向前传递，节省 FLOPs。如果门打开（1），则该块在 LoRA 适配器的帮助下执行，以微调其行为。训练过程使用标准的语言建模损失加上一个正则化项来联合优化这些路由器和 LoRA 权重，该正则化项强制路由器识别真正可跳过的层，有效地教导模型在适当的地方变得'懒惰'。
核心创新在于使用统一的语言建模目标结合门控正则化对 LoRA 适配器和硬门控路由器进行联合训练。与可能在显式标签上训练路由器或使用软概率的现有方法不同，LayerRoute 使用 STE 通过硬二进制决策（$g_i \in \{0, 1\}$）传递梯度，确保模型学习到它在推理时将使用的确切行为。另一个关键创新是偏置初始化策略：中间层被初始化为倾向于跳过，而早期和晚期层倾向于执行。这打破了对称性问题（即门可能永远无法区分），并立即将中间层的效用暴露给梯度信号。这使得系统能够纯粹通过观察 LM 损失，自主发现中间层（8-16）对于工具调用是多余的，但对于规划是必不可少的，而无需任何显式的监督标签来定义'工具调用'或'规划'。
方法步骤详情
首先，架构增强了 24 个 Transformer 块中的每一个。一个路由器 $r_i$ 从均值池化输入 $\bar{h}_i$ 计算分数 $s_i = w_i^T \bar{h}_i + b_i$。该分数通过 sigmoid 函数 $\sigma(s_i)$ 并阈值化以产生硬门 $g_i = 1[\sigma(s_i) > 0.5]$。同时，秩为 8 的 LoRA 适配器被添加到 Q/K/V/O 投影中（$W' = W + \frac{\alpha}{r}BA$）。其次，在前向传播过程中，块的输出计算为 $h_{i+1} = g_i \cdot \text{Block}_i(h_i) + (1-g_i) \cdot h_i$。如果 $g_i=0$，则跳过该块。第三，使用组合损失函数 $\mathcal{L} = \mathcal{L}_{\text{LM}} + \lambda \sum \sigma(s_i)$ 训练模型，其中 $\lambda=1.0$ 惩罚保持门打开。离散门操作的梯度使用直通估计器（STE）近似。最后，路由器使用偏置值初始化（层 0-7 和 17-23 为 $+1.0$，层 8-16 为 $-1.0$），以鼓励在训练开始时立即进行区分。
技术新颖性
这项工作的技术新颖性源于几个解决动态推理特定挑战的设计选择。首先，带有硬门的 STE 的使用消除了软门控方法中常见的训练/推理差异，确保模型为确切的推理行为进行优化。其次，偏置初始化是层跳过中冷启动问题的一个新颖解决方案，强制中间层'证明其价值'或立即被跳过，这加速了向两簇结构的收敛。第三，该方法是参数高效的，仅使用 1.10M 参数（骨干网络的 0.22%），使其高度适合部署。最后，通过 LM 目标隐式学习任务特定路由（工具 vs 规划），而无需显式数据集标签，是一个重要的方法论贡献，证明了当与正确的架构偏差耦合时，结构效率如何可以从标准的预测训练信号中涌现。
实验结果
实验证明 LayerRoute 成功学会了区分工具调用和规划步骤。在 100 个样本的保留集上，该方法实现了 12.91% 的跳过差异，意味着工具调用跳过了 15.25% 的层，而规划步骤仅跳过了 2.34%。这转化为工具调用 15.2% 的 FLOPs 减少，验证了结构化调用需要较少深度的假设。令人惊讶的是，门控模型实际上比全层基线改善了困惑度，工具调用上的 delta 为 -1.29，规划步骤为 -1.30。这归因于 LoRA 适配，它补偿了跳过的质量损失。消融研究证实了门控正则化的必要性（LayerRoute-NoReg 实现了 0.1% 的跳过）和偏置初始化（LayerRoute-UniformInit 将差异反转为 -5.4%）。学习到的门结构收敛到一个稳定的模式，其中层 8-16 形成工具调用的跳过块，这与关于 Transformer 层功能角色的可解释性发现一致。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 跳过差异 | 工具调用跳过% - 规划跳过% | 12.91% | LayerRoute-UniformInit: -5.4% | 优于均匀初始化 18.31 个百分点 |
| 工具调用 FLOPs 减少 | FLOPs 降低百分比 | 15.2% | Full model: 0.0% | 绝对降低 15.2 个百分点 |
| 工具调用困惑度改进 | PPLΔ (PPL(gated) - PPL(full)) | -1.29 | LayerRoute-BCE: +0.06 | 优于 BCE 方法 1.35 |
| 规划步骤困惑度改进 | PPLΔ | -1.30 | LayerRoute-UniformInit: -1.28 | 略优于均匀初始化 |
局限与改进
作者承认了几个局限性。首先，路由有些粗糙；层 8-16 作为一个单一的跳过块，没有内部区分，这意味着系统无法识别该簇内的哪些特定层是最或最不冗余的。其次，FLOPs 的减少是适度的（15.2%），对于规划繁重的工作负载，益处可以忽略不计（仅 2.34% 跳过）。第三，评估仅限于 Qwen2.5-0.5B 模型；这种行为可能无法推广到更大的模型（3B、7B）或层重要性可能不同的其他架构。第四，数据集不平衡，仅来自 Turing Open Reasoning 数据集的 50 个样本，可能限制了模型对多样化规划模式的暴露。此外，该方法依赖于序列级路由，这对于在相同序列中包含工具调用和推理的混合内容输入可能不如 token 级路由灵活。
独立分析的弱点
一个明显的弱点是缺乏簇内路由的粒度。该方法将所有中间层（8-16）视为一个要跳过的单一单元。这阻止了它发现，例如，即使层 11-15 是冗余的，层 10 可能也是关键的。这限制了潜在的 FLOPs 节省，因为路由器无法在更细的粒度上进行优化。另一个弱点是对 LoRA 适配以保持质量的依赖。虽然 LoRA 在这些实验中改善了困惑度，但它增加了部署的复杂性；如果下游任务需要确切的原始权重（例如，为了严格的安全对齐），这种修改可能是不可取的。此外，规划步骤上的低跳过率（2.34%）表明，虽然该方法对工具有效，但它还不是所有智能体任务的通用解决方案。改进可能涉及分层路由（首先决定'跳过块'，然后决定'哪些内部层'）或动态正则化权重，根据预测的困惑度变化。
未来方向
作者提出了几个未来方向。中间簇内的更细粒度路由是首要任务，可能需要特定层的蒸馏目标或更高秩的适配器来捕获单个层的贡献。扩展到更大的模型尺寸（3B、7B+）对于验证两簇结构（早期/晚期必需，中间可跳过）是否适用于更深的模型至关重要。与其他效率技术如投机解码的集成可以提供乘法增益——LayerRoute 减少验证步骤的成本，而投机解码处理草稿生成。另一个方向是逐 token 路由，其中同一序列中的不同 token 可能被路由到不同的深度，尽管这会增加实现的复杂性，并可能失去序列级跳过的效率增益。最后，探索对验证损失更敏感的替代正则化策略可能有助于发现更激进的跳过时间表，而不会牺牲质量。
复现评估
该论文通过详细的训练设置提供了高可复现性：使用 Qwen2.5-0.5B-Instruct，AdamW 优化器，$\beta_1=0.9, \beta_2=0.999$，学习率 $2 \times 10^{-4}$ 并带有余弦退火，并在 A100 40GB 上精确训练 3000 步（大约 6.4 分钟）。数据集是标准的公共基准（Hermes Function Calling、Glaive、GSM8K、Turing），使得数据获取变得简单明了。架构规范精确：秩 8 LoRA，$\alpha=16$，路由器维度 896，以及特定的初始化偏置值。虽然提供的文本中没有明确链接代码，但给出了精确的参数计数和损失公式（$\mathcal{L} = \mathcal{L}_{\text{LM}} + \lambda \sum \sigma(s_i)$）。短训练时间（不到 7 分钟）和低参数数量（1.1M）使得这个实验在配备适度 VRAM 的标准消费级硬件上极其容易复制，假设基础模型可用。
论文图表
该图展示了工具调用（如 search_products(...)）和规划输入（如 Summarize quarterly sales trends...）的典型门控模式。每张图有 24 个位置代表 24 个 Transformer 层，用 1 和 0 表示层是运行还是跳过。从图中可以看出，工具调用的门控模式中 0 的数量明显多于规划输入，具体来说工具调用在中间层有更多的跳过（如 29.2% 和 33.3% 的跳过率），而规划输入的跳过率较低（12.5%）。图中的模式还显示，不同输入的跳过模式有所变化，证实了输入条件路由的有效性。
这张图对理解论文至关重要，因为它直观地展示了 LayerRoute 实际学到了什么——即如何根据输入类型（工具调用 vs 规划）产生不同的层跳过模式。它证明了模型确实学会了区分这两种任务类型，并且在工具调用上更加激进地跳过中间层。