（2024|ICML PMLR，强化学习，动态决策，网络跳过，质量和计算的权衡）可切换决策：动态神经生成网络

Switchable Decision: Dynamic Neural Generation Networks

公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）

目录

0. 摘要

3. 方法

3.1 构建离散决策空间

3.2 强化学习代理

3.3. 约束优化

4. 实验

7. 结论

---

0. 摘要

自回归生成模型在许多不同的自然语言处理任务中取得了竞争性的性能，如摘要、问答和分类。然而，它们也因为推理速度较慢而闻名，这使得它们在实时应用中部署变得具有挑战性。我们提出了一种可切换的决策，通过动态分配计算资源来加速推理过程。通过自动决定跳过何处以及如何平衡质量和计算成本，我们的动态神经生成网络强制执行高效的推理路径，并确定优化的权衡。在问答、摘要和分类基准上进行的实验表明，我们的方法在推理过程中受益于更少的计算成本，同时保持相同的准确性。广泛的实验和消融研究表明，我们的方法可以是通用的、有效的，并且对许多自然语言处理任务都是有益的。 

相似的网络跳过策略：

（2024，FLOPs 动态分配，MoD，MoDE，top-k 路由，块丢弃）在基于 Transformer 的语言模型中动态分配计算 

3. 方法

我们的可切换决策（见图1）网络专注于加速自回归语言生成模型的推理时间。具体来说，我们提出了一种可切换决策的通用方案：1）构建多功能的决策空间，2）利用依赖输入的强化学习代理，3）提出词典（lexico）优化策略。

记输入为 o = (o_0, · · · , o_n)。对于一个包含 n 个标记的序列，一个基于 Transformer 的语言生成模型 M 具有 L 层，首先嵌入标记获得矩阵 O_e ∈ R^(n×e)，其中 e 是嵌入空间的维度。然后，这些标记表示经过语言模型的编码器和解码器。为了加速推理时间，同时保持类似的高质量，我们决定每个输入数据是否应该跳过一层。随着层数的增加，这个决策问题呈指数级增长。此外，由于决策空间的离散性质，优化变得具有挑战性。在本节中，我们概述了我们的独特设计选择，以实现我们的目标并克服优化挑战。

3.1 构建离散决策空间

我们提出使用可切换决策网络来加速推理时间，为每个示例学习最佳的（输入，推理路径）对配置。我们考虑三个搜索空间候选者，即注意力层、前馈层和第一层之后的查询（query）输入。我们现在解释下面每个搜索空间的详细信息。

注意力候选者。Transformer-based 语言模型的一个关键组成部分是注意力层。张等人（2019年）发现某些层是多余的。为了决定是否跳过某个层，我们将这些决策建模为由策略网络 q 参数化的 i.i.d. 伯努利随机变量序列。令 b_l 表示第 l 层的可切换决策，定义为

其中 x ∈ R^e 表示决策单元的输入，并且我们将第一个编码器层的输出作为 x。策略（policy）网络 g 学习保留每个层的隐藏表示的实例专属概率。为了执行跳过操作，我们从这个分布中采样，并将指示符 b^att_l 广播到注意力层的输入表示中。

前馈候选者。在同样的精神下，前馈层可能包含多余的信息。因此，我们考虑使用与注意力层相同的方法跳过这些层。我们根据指示符 b^ffn_l 决定是否跳过。策略网络的设计与注意力层相同。

Token 候选者。除了跳过层之外，跳过标记也可以是节省计算成本的替代方法。我们创建了两种标记跳过策略：➀ 跳过最后的 p% 标记，和 ➁ 均匀跳过 p% 标记。对于前者，我们将 p 设置为 10、20 和 30。对于后者，p 等于 25、33 和 50。为了决定使用哪种策略，我们优化一个由函数 h(·) 参数化的分类随机变量。h(·) 的输入与 g(·) 相同，其输出是所有六个候选决策的分布。

编码器和解码器结构。我们感兴趣的架构包含编码器和解码器。对于编码器，我们将注意力跳过和前馈跳过与标记跳过一起应用。对于解码器，由于每个标记对于最终输出都是有意义的，我们只应用注意力跳过和前馈跳过。在做出决策时，我们从策略网络的输出中采样，并将决策广播到每一层的隐藏表示中。

3.2 强化学习代理

策略（Policy）网络架构。由于我们的目标是加速推理过程，我们采用了简单的策略网络设计。我们利用一个带有层归一化和 ReLU 激活函数的单层 MLP。为了在决策上输出二项分布，我们对注意力和前馈候选者的网络输出应用 Sigmoid 激活函数。对于标记候选者的选择，我们使用 softmax 函数输出分布。

参数化。在训练过程中，我们从由策略网络参数化的决策分布中采样。层的可切换决策分布可以表示为一个 2L 维的伯努利分布，可以写成：

相似地，token 跳过决策的分布可以表示为一个分类分布，可以形式化为：

其中 a 表示跳过策略的选择，J 表示总的策略数量。我们在实践中应用了七个候选策略。

奖励。我们将奖励函数（Yang等人，2022b;a；Feng等人，2023）定义为质量和计算成本之间的权衡。给定一个推理路径和一个数据实例，奖励可以从计算（估计的 FLOPs）中计算出来。直观地，跳过层将具有较高的奖励。我们进一步将质量称为准确性和损失，如下所示：

其中质量为 -loss，计算为估计的 FLOPs（浮点运算次数），λ 是一个系数。整体损失函数定义为奖励的期望值：

其中 π 和 η 分别在 (2) 和 (3) 中定义。

优化。为了优化我们的策略网络，我们应用策略梯度来计算 J 的梯度，并更新策略网络的参数。我们使用自临界基线来减少梯度的方差。进一步在质量和计算上应用约束优化策略。详细信息见下一节。

推理期间。与训练过程不同，我们在推理过程中不对跳过决策进行采样。相反，我们选择最大化似然函数的决策。

3.3. 约束优化

权衡是一个问题。在主网络和策略网络的联合训练中，质量和计算之间的权衡很重要。多个目标的线性组合是最广泛使用的方法。然而，组合的系数需要手动调整，并且在理论上不适用于非凸函数。在这项工作中，我们考虑在两个目标之间进行约束优化，特别强调词典（lexico）优化。我们的方程。为了优化方程（4）中质量和计算之间的权衡，我们提出使用词典优化，其中参数按以下方式迭代更新：

其中 γt ≥ 0 是自适应步长，et ∈ R^d 是要选择的更新方向，以平衡 f 的最小化和 q 的约束满足。其中一个目标（f，我们的情况是计算）相对于另一个目标（q，我们的情况是质量）不太重要。约束优化的设计标准是当约束不满足时（即 q(θt) ≥ c），重点是尽快减小 q 以满足约束；与此同时，f 作为一个次要目标，应尽量减少到不会对 q 的下降造成伤害。因此，我们应用以下更新规则来实现这个目标：

其中 ∇computation 和 ∇quality 是由评分函数估计的，λ 可以计算为

其中 ϕ(θt) 等于 q(θt) − c，c 表示最小损失。

提出的算法。我们的可切换决策（switchable decision，SD）与高效的候选空间和约束优化在算法 1 中显示。我们以单循环方式迭代更新自回归模型和策略网络。策略网络参数 θ 通过方程（6）进行更新，以平衡质量的优化和对计算的约束满足。

4. 实验

根据Lewis等人（2019）的做法，我们采用预训练的 BART 模型作为骨干，并利用提供的检查点在下游数据集上进行微调。 

7. 结论

我们的工作展示了引入动态网络可切换决策的益处。所提出的方法可以显著提高推理效率，同时仍能保持模型性能。在摘要、问答和分类基准测试中观察到了明显的 FLOPs 节省和一致的性能。我们进一步在不同设置下对所提出的可切换策略进行了详细研究，例如，与不同的架构搜索空间进行比较，提供更多根据隐藏表示进行决策的证据，并验证组件的影响。总而言之，可切换决策是有效的和通用的，有潜力被整合到现有的生成模型中，用于各种自然语言处理任务。