[论文解读] Confident Adaptive Language Modeling

Tal Schuster, Adam Fisch+6 more|arXiv (Cornell University)|Jul 14, 2022

Topic Modeling被引用 39

一句话总结

CALM 引入对 Transformer 语言模型的逐词（per-token）早停决策，校准局部退出以可证明地保留全局序列质量，在若干生成任务上实现高达约 3 倍的推理加速。

ABSTRACT

Recent advances in Transformer-based large language models (LLMs) have led to significant performance improvements across many tasks. These gains come with a drastic increase in the models' size, potentially leading to slow and costly use at inference time. In practice, however, the series of generations made by LLMs is composed of varying levels of difficulty. While certain predictions truly benefit from the models' full capacity, other continuations are more trivial and can be solved with reduced compute. In this work, we introduce Confident Adaptive Language Modeling (CALM), a framework for dynamically allocating different amounts of compute per input and generation timestep. Early exit decoding involves several challenges that we address here, such as: (1) what confidence measure to use; (2) connecting sequence-level constraints to local per-token exit decisions; and (3) attending back to missing hidden representations due to early exits in previous tokens. Through theoretical analysis and empirical experiments on three diverse text generation tasks, we demonstrate the efficacy of our framework in reducing compute -- potential speedup of up to $ imes 3$ -- while provably maintaining high performance.

研究动机与目标

• 在不牺牲输出质量的前提下，推动降低大型 Transformer 语言模型的推理成本。
• 开发一个有原则的校准框架，通过局部逐词退出来强制执行全局序列级约束。
• 识别用于早停分类器的有效逐层置信信号和训练目标。
• 展示在多种文本生成任务中的效率提升和可靠性。

提出的方法

• 提出 CALM：一个框架，利用局部逐词置信分数在自回归解码过程中决定早停。
• 用无分布假设的风险控制和 Learn-Then-Test (LTT) 校准，提出必须以高概率满足的全局一致性约束（文本化或基于风险的）。
• 推导一个在置信度超过经校准阈值时就尽早退出的停止策略，同时在自回归解码中考虑状态传播。
• 训练轻量级的逐层早停分类器，并分析不同的置信度度量（softmax-diff、隐藏状态饱和度，以及显式退出分类器）。
• 引入一个衰减阈值，以在解码步骤间平衡效率与鲁棒性（λ′(λ,t)）。

实验结果

研究问题

• RQ1 如何量化并校准逐词退出决策，以确保全局序列级约束以高概率得到满足？
• RQ2 哪些置信信号最能预测自回归 Transformer 解码中的安全早停？
• RQ3 在多样的生成任务（摘要、翻译、问答）中应用 CALM 的效率提升和性能权衡是什么？
• RQ4 状态传播如何与早停交互，以及我们是否可以通过校准和训练策略缓解潜在的性能下降？

主要发现

• CALM 在可证明的全局保证下实现显著的计算量降低，在评估和训练设置中实现约 3x 的加速。
• 基于经过校准的置信度的早停决策可以在 CNN/DM、WMT、SQuAD 等任务中保留大部分全模型性能，同时显著减少 FLOPs。
• 对被跳过的层采用状态传播技术以及衰减阈值，提升逐词退出的鲁棒性和效率。
• 不同的局部置信度度量存在权衡；基于 softmax 的置信度通常在效率提升显著的情况下表现较强，而退出分类器在高性能场景中可能更节省 FLOPs。
• Learn-Then-Test 校准框架提供了一种有原则的方法来选择退出阈值，以满足文本或基于风险的全局约束，并控制错误率。