type
status
date
slug
summary
tags
category
icon
password
为了解决自回归语言模型表现出的高推理延迟，之前的研究提出了一种早退框架，该框架根据生成后续 token 的复杂度为每个 token 分配自适应计算路径。
缺点：由状态复制机制或大量推出路径引起的性能下降；对退出置信度阈值的敏感性
因此，本文提出了一个快速且鲁棒的早退框架（FREE），该框架包含”浅-深模块“和同步并行解码。
该框架通过将当前 token 的解码过程与先前堆叠的早退 token 同步以实现更快的推理。
此外，由于并行解码允许我们观察浅层和深层模型的预测，因此我们提出了一种新颖的自适应阈值估计器，它利用 Beta 混合模型来确定合适的置信度阈值。
1 简介
为了减小推理延迟，早退框架根据每个 token 生成的复杂性动态分配计算路径。容易预测的 token 仅需几层计算即可产生一致预测，难度较高的 token 需要计算更多层才能准确预测。
理想情况下模型早退能够显著加速推理并且不影响生成质量。
然而早退框架中存在三个挑战：
- 尽管有可能在较早的层退出，但处理后续 token 仍然需要剩余层的键和值状态。虽然有状态复制机制通过重用早退层的隐藏状态来有效计算这些状态，但研究结果表明对于大模型和长序列输出表现不佳。
- 将所有层都设为可退出位置并不能保证更快推理，因为：
- 浅层的性能缺陷可能导致生成异常长序列
- 每层置信度测量的计算开销
- 通过早退实现目标延迟和准确性很大程度上取决于为目标任务选择合适的置信阈值，这通常需要付出巨大努力和额外计算开销
本文提出的 FREE 框架结合了“浅-深模块”和“同步并行解码”，不仅可以为更大的模型和更长的输出序列提供一致的加速和性能，而且还消除了寻找合适退出阈值计算成本高昂的问题。
具体来说，浅深模块将计算路径分为浅层模型（具有指定数量的早期层）和深层模型（包括所有层）。同步并行解码会累积连续的早期退出 token，这些 token 仅通过浅层模型，直到遇到非早退 token。将当前的非早退 token 的解码过程与之前堆叠的令牌同步。这通过利用实际注意力计算的键和值而不是通过状态复制的近似状态来防止性能下降，同时与自回归解码单个令牌相比，并行解码更高效。
此外，我们利用并行解码输出预测的事实设计了一种新颖的自适应阈值估计器。该估计器使用 Beta 混合模型 (BMM) 来捕获两个模型的置信度分数和预测对齐之间的相关性，从而确定每个数据集的正确置信度阈值。
2 相关工作
3 再评估早退框架
对现有框架重新评估的四项新发现
3.1 对大模型和长输出缺乏鲁棒性
对于较大的模型和具有较长序列的数据集，状态复制导致的生成质量下降会变得严重
3.2 退出位置导致性能下降
当仅使用 T5-large 的一小部分浅层的时候，所有 token 在同一层退出会导致性能显著下降
3.3 不可忽略的计算成本
与不使用提前退出的基线模型相比，提前退出通常会导致总解码时间意外增加。
这可以归因于测量每一层的置信度所涉及的不可忽略的计算成本，特别是由于具有大词表的 softmax 操作。
此外，虽然状态复制方法的目的是减少剩余层的 MHA 和 FFN 层的计算时间，但使用重复的隐藏状态计算键和值状态会产生额外的不可忽略的开销。
3.4 不同的最佳置信阈值
在相同性能下实现最低延迟的最佳置信阈值在不同数据集上存在显着差异
4 FREE
4.1 浅-深模块
有策略地将早期的一些层作为浅层模型，全部层作为深层模型，解决传统早退框架共同训练多个退出层带来的性能下降问题。
为了增强浅层模型的能力，用逐层知识蒸馏（KD）作为 Equation 1 的额外损失项 和 ：
Equation 1
表示深层模型中将知识提取到浅层模型的对应的层
和 是浅层和深层模型的隐藏状态
实验从最后一层、从固定的统一映射层、动态映射层蒸馏（）
虽然四种损失函数在改变退出阈值时的性能差异并不显著，但 KD-dyna 有更好的权衡
4.2 同步并行解码
4.3 自适应阈值估计
与之前利用额外校准集的方法不同，我们通过使用早期实例的信息快速调整阈值
具体来说，在并行解码过程中，我们收集样本来评估浅层模型的置信度分数与浅层模型和深层模型之间的预测对齐之间的对应关系。
