LayerRoute: 一种通过LoRA保持质量的自适应层跳过方法，用于高效LLM推理

# LayerRoute：一种自适应层跳过与LoRA保留质量的高效LLM推理方法
来源：https://arxiv.org/html/2609.13682

###### 摘要

我们提出LayerRoute，一种参数高效的自适应Transformer层跳过方法，结合了通过直通估计器训练的逐层硬门控路由与联合LoRA微调。LayerRoute在Qwen2.5-0.5B-Instruct的每个Transformer块中增加了轻量级的逐层路由器（约21.5K参数）和LoRA适配器（秩8，约1.08M参数），并在门控正则化的语言建模目标下联合训练两者。在10次独立随机种子训练中，LayerRoute均收敛到相同的跳过模式结构——一组连续的9个中间层（8-16层）在所有10次训练中均具备跳过资格——并在每次训练中均带来实际验证的时钟加速（1.02×–1.06×，平均1.04×）。在所有测试配置中，模型质量得到保持或提升：联合LoRA适配在所有10次训练中均比未修改的主干模型带来困惑度改善（平均Δ=-1.16和-1.11，对应两个评估数据集）。我们进一步验证路由器执行了真正的、非平凡的逐输入计算：在具备跳过资格的层中，门控决策改变了87-100%的留出样本的实际跳过/执行结果，证实了真正的输入依赖路由而非固定剪枝模式。LayerRoute在单个A100上训练时间低于7分钟，除路由决策外几乎不增加额外开销。我们在此工作中完整报告了可复现性方法，包括对路由器逐输入决策决定因素的系统性诊断研究。

## 1引言

现代LLM推理对每个输入token和输入序列应用统一计算，无论该输入实际需要多少计算量。这是效率提升的自然切入点：如果Transformer层中相当大一部分对给定输入的最终预测贡献甚微，跳过这些层可在质量损失极小的情况下节省实际计算量。

我们提出LayerRoute，一种轻量级适配器，它学习按输入跳过Transformer块，并与LoRA微调联合训练。我们的设计结合了三个组成部分：（1）逐层路由器，通过直通估计器产生硬二元跳过决策，消除了训练/推理时跳过应用的不一致；（2）注意力投影上的LoRA适配器，与路由器联合训练，使适配与路由能够协同演进；（3）门控正则化项和偏置中间层初始化，两者对于逃离退化的全开放平衡并达到真正的跳过行为都是必要的。

我们的核心贡献不仅在于架构，更在于对其能和不能声称的方面进行严格说明。基于路由的条件计算方法通常在单一训练运行上进行评估，这可能导致报告的效果量不具备泛化性。我们改为在10次独立随机种子训练中评估LayerRoute，并仅报告在所有训练中可可靠复现的结果：稳定的跳过模式结构、真实且一致的时钟加速，以及通过LoRA实现的一致质量保持。我们还直接验证（而非假设）路由器的决策是非平凡的且真正依赖于输入，并报告了关于路由器逐输入决策实际跟踪何种信号的系统性研究，包括我们认为有必要记录而非忽略的几项阴性结果。

我们的贡献包括：
1. 一种逐层硬门控跳过连接架构，采用直通估计，与LoRA适配器联合训练，在单个A100上训练时间低于7分钟。
2. 一项10次随机种子的可复现性研究，确定了LayerRoute的哪些特性是稳定的（跳过模式结构、时钟加速、质量保持）以及哪些不是（第7节）。
3. 直接的经验证据，证明训练后的路由器执行了真正的、实质性的逐输入计算，将其与固定剪枝模式区分开来。
4. 透明地说明了我们诊断路由器行为的方法论，旨在作为评估其他条件计算方法的严谨性模板。

## 2相关工作

早期退出与层跳过。DeeBERT（Xin等人，2020）和PABEE（Zhou等人，2020）基于中间置信度从Transformer编码器中实现早期退出；BranchyNet（Teerapittayanon等人，2016）将相同思想应用于卷积网络。CALM（Schuster等人，2022）及其早期版本（Schuster等人，2021）将逐层置信度阈值校准到token级早期退出的目标风险水平。SkipBERT（Tang等人，2023）学习通过输入依赖路由跳过整个层。深度自适应Transformer（Elbayad等人，2020）和PonderNet（Banino等人，2021）扩展了自适应计算时间（Graves，2016）以学习按样本的停止策略。LayerDrop（Fan等人，2020）通过结构化dropout训练Transformer对层移除具有鲁棒性，从而在没有学习路由器的情况下实现事后深度缩减。混合深度（MoD）（Raposo等人，2024）应用逐token路由来控制Transformer深度上的计算分配。LayerRoute与MoD共享跳过连接机制，但在训练策略上有所不同：我们使用LoRA联合微调而非从头训练，并且在序列级而非token级进行路由。事后深度剪枝分析，如ShortGPT（Men等人，2024）和Gromov等人（2024）证实了后续Transformer层通常对输出质量的贡献不成比例地小；这启发了我们的偏置中间层初始化（第3.4节），并与前馈层作为具有不同逐层重要性的键值记忆的可解释性发现相关联（Geva等人，2021）。

离散与条件计算。端到端训练硬的、离散的决策是一个长期挑战；我们使用的直通估计器（Bengio等人，2013）是一种解决方案，而连续松弛方法如Gumbel-softmax（Jang等人，2017）和Concrete分布（Maddison等人，2017）提供了我们在此未探索的替代方案。大规模混合专家路由，如Switch Transformer（Fedus等人，2022）和GShard（Lepikhin等人，2021），共享条件计算的更广泛目标，但在并行专家之间路由而非跳过顺序深度。

参数高效微调。LoRA（Hu等人，2022）将权重更新分解为低秩矩阵，以最少参数实现高效适配；QLoRA（Dettmers等人，2023）将其扩展到量化主干。更早的参数高效方法包括适配器模块（Houlsby等人，2019）、前缀微调（Li和Liang，2021）和提示微调（Lester等人，2021）。我们将LoRA应用于注意力投影，同时训练逐层路由器，允许路由策略与适配权重协同演进。

LLM的高效推理。推测解码和自推测解码（Leviathan等人，2023；Stern等人，2018）通过草拟多个token并在批处理中验证来加速自回归生成。PagedAttention（Kwon等人，2023）和FlashAttention（Dao等人，2022）通过内存管理和内核设计而非架构变更来提高吞吐量。训练后量化（Dettmers等人，2022；Frantar等人，2023）和结构化剪枝（Ma等人，2023；Sun等人，2024）直接减小模型规模。LayerRoute与所有这些方法互补——它减少每个输入的活动层数，而不是优化注意力内核、内存、数值精度或参数量。

智能体LLM系统。我们的训练数据构建受到交替使用结构化工具调用与开放式推理的智能体工作流的启发（Yao等人，2023；Schick等人，2023）；我们使用这种异构数据混合来训练LayerRoute，但根据我们的可复现性发现（第7节），并不声称路由器的决策专门基于此任务类型区别进行条件化。

## 3方法

### 3.1架构

LayerRoute在冻结的预训练Transformer（Qwen2.5-0.5B-Instruct，24层，隐藏维度896）上为每个Transformer块增加了两个组件：逐层路由器和LoRA适配器。

逐层路由器。每个路由器 \(r_i\)（\(i=0,...,23\)）是一个轻量级线性层：
\[ s_i = w_i^\top \bar{h}_i + b_i, \quad \sigma_i = \sigma(s_i), \quad g_i \in \{0,1\} \]
其中 \(\bar{h}_i \in \mathbb{R}^d\) 是进入块 \(i\) 的平均池化隐藏状态，\(w_i \in \mathbb{R}^d\)，\(\sigma(\cdot)\) 是sigmoid函数。硬门控为 \(g_i = \mathbb{1}[\sigma_i > 0.5]\)。每个路由器有 \(d+1=897\) 个参数；24个路由器的集合总共有21,528个参数。

门控跳过连接。通过块 \(i\) 的前向传播为：
\[ h_{i+1} = g_i \cdot \text{Block}_i(h_i) + (1-g_i) \cdot h_i \]
当 \(g_i=1\) 时，块正常运行（LoRA适配器激活）。当 \(g_i=0\) 时，隐藏状态不变地通过。

LoRA适配器。对于每个块中的每个注意力投影 \(W \in \{W_Q, W_K, W_V, W_O\}\)，我们添加低秩适配器 \(W' = W + \frac{\alpha}{r} BA\)，秩 \(r=8\)，\(\alpha=16\)，\(B\) 初始化为零。LoRA总参数：\(4 \times 24 \times 2 \times 896 \times 8 = 1,081,344\)。

### 3.2直通估计器

硬阈值 \(g_i = \mathbb{1}[\sigma_i > 0.5]\) 是不可微的。我们应用直通估计器（Bengio等人，2013）：
\[ \hat{g}_i = \underbrace{\mathbb{1}[\sigma_i > 0.5]}_{\text{forward}} - \underbrace{\sigma_i}_{\text{stop-grad}} + \underbrace{\sigma_i}_{\text{backward}} \]
前向传播使用硬门控；梯度通过 \(\sigma_i\) 传递，如同它是连续的，在训练和推理时使用相同的硬 \(\{0,1\}\) 决策。

推理时计算实现。在训练期间，无论 \(g_i\) 如何，每个块的前向计算都必须无条件运行，因为直通估计器的反向传播需要块的真实前向值来正确连接梯度，即使当 \(g_i=0\) 时。在推理时，此约束不适用：当 \(g_i=0\) 时，我们完全跳过块的计算，而不是计算后丢弃结果。这种区别——计算后丢弃与真正不计算——对模型输出没有影响（在任一实现中，丢弃的计算都不会影响结果），但决定了报告的加速是理论性的还是测量的；本文中我们仅报告时钟测量结果（第5节）。

### 3.3训练目标

所有可训练参数通过以下方式联合优化：
\[ \mathcal{L} = \mathcal{L}_{LM} + \lambda \cdot \frac{1}{L} \sum_{i=0}^{L-1} \sigma(s_i) \]
其中 \(\mathcal{L}_{LM}\) 是标准的自回归交叉熵损失，第二项是权重为 \(\lambda=1.0\) 的门控正则化，惩罚均匀高的软门控值以防止坍缩到全开放平衡。

### 3.4偏置初始化

均匀初始化将所有门控置于 \(\sigma(0)=0.5\) 附近，造成一种对称性阻碍区分：门控在开始跳过之前不会分化，但在分化之前不会跳过。我们通过层位置相关的偏置初始化打破这种对称性：
\[ b_i = \begin{cases} +1.0 & i \in \{0\text{--}7, 17\text{--}23\} \\ -1.0 & i \in \{8\text{--}16\} \end{cases} \]
中间层从阈值以下开始，在第一步就跳过，立即将其贡献暴露给语言模型损失的梯度。

## 4实验设置

基础模型。Qwen2.5-0.5B-Instruct（Yang等人，2024）：24个Transformer块，隐藏维度896，分组查询注意力（14个查询头 / 2个KV头），SwiGLU FFN，词汇表151,936。

训练数据。一个混合的智能体数据集，涵盖工具调用来源（Hermes Function Calling v1，Glaive Function Calling v2）和推理/规划来源（GSM8K，Turing Open Reasoning），总计10,749个训练样本 / 1,194个验证样本。

训练。AdamW（\(\beta_1=0.9\)，\(\beta_2=0.999\)，权重衰减0.01），学习率 \(2 \times 10^{-4}\) 带余弦退火，批次大小4，梯度累积4，梯度裁剪1.0，3000步。硬件：单个A100 40GB。训练时间：382秒（≈6.4分钟）。

可复现性协议。我们使用10个独立选择的随机种子训练LayerRoute，每个种子设定权重初始化、dropout和批次顺序（已验证正确到达训练流程中的每个随机性来源）。留出评估样本的选择在每次运行中单独设种子，匹配训练种子，使得仅从种子即可完全复现训练和评估。

评估。对于每个种子，我们在100个留出样本（每个评估来源50个，一个代表工具调用风格的短结构化生成，一个代表较长形式的推理）上进行评估，使用：时钟推理延迟（测量值，非理论值，启用真实推理时计算跳过）、门控模型的困惑度。