[论文解读] TR-BERT: Dynamic Token Reduction for Accelerating BERT Inference

Deming Ye, Yankai Lin+2 more|arXiv (Cornell University)|May 25, 2021

Topic Modeling参考文献 48被引用 7

一句话总结

TR-BERT 提出一种动态标记剪枝方法，通过强化学习学习在每一层保留哪些标记，以基于其任务相关性加速 BERT 推理。该方法在仅造成轻微准确率下降的情况下实现了 2–5 倍的加速，尤其在长序列任务中受益显著，通过选择性地传播标记减少了二次方复杂度的自注意力计算。

ABSTRACT

Existing pre-trained language models (PLMs) are often computationally expensive in inference, making them impractical in various resource-limited real-world applications. To address this issue, we propose a dynamic token reduction approach to accelerate PLMs' inference, named TR-BERT, which could flexibly adapt the layer number of each token in inference to avoid redundant calculation. Specially, TR-BERT formulates the token reduction process as a multi-step token selection problem and automatically learns the selection strategy via reinforcement learning. The experimental results on several downstream NLP tasks show that TR-BERT is able to speed up BERT by 2-5 times to satisfy various performance demands. Moreover, TR-BERT can also achieve better performance with less computation in a suite of long-text tasks since its token-level layer number adaption greatly accelerates the self-attention operation in PLMs. The source code and experiment details of this paper can be obtained from https://github.com/thunlp/TR-BERT.

研究动机与目标

• 为解决预训练语言模型（PLMs）在资源受限环境中的高推理延迟问题。
• 克服固定层剪枝的局限性，后者会统一减少模型深度，从而损害复杂任务的性能。
• 通过动态选择哪些标记在深层中传播，实现在每个标记上的自适应计算。
• 在自注意力随序列长度呈二次方增长的长文本场景中提升效率。
• 为实际部署提供灵活的按需速度-准确率权衡。

提出的方法

• TR-BERT 将标记剪枝建模为多步选择问题，其中每个标记在每一层被选择或跳过。
• 它在强化学习框架中采用策略网络，基于任务特定的重要性学习动态标记选择策略。
• 强化学习智能体接收一个结合预测置信度和所选标记数量惩罚的奖励，以平衡速度与准确率。
• 被认为不重要的标记会被早期剪枝，从而减少后续计算并避免冗余的自注意力操作。
• 该方法在标记级别运行，允许对每个输入标记的计算进行细粒度控制，不同于层级别的剪枝。
• 该框架即插即用，无需微调基础 BERT 模型，且与现有 PLMs 兼容。

实验结果

研究问题

• RQ1 动态的、基于标记级别的剪枝是否能在不牺牲模型性能的前提下提升推理速度？
• RQ2 基于任务相关性的动态标记选择与固定层或逐层剪枝相比有何差异？
• RQ3 强化学习能否在无人工标注监督的情况下有效学习标记重要性信号？
• RQ4 标记级别剪枝在长序列 NLP 任务中能多大程度提升效率？
• RQ5 该方法是否能在多种下游任务中实现灵活的速度-准确率权衡？

主要发现

• TR-BERT 在十一个 NLP 基准测试中实现了 2–5 倍的推理加速，具体取决于所需的准确率水平。
• 在问答任务中，TR-BERT 表现优于基线方法，尤其在复杂推理场景中，固定剪枝方法会失效。
• 对于长文本任务（例如 1,024 个输入标记），TR-BERT 在显著降低推理时间的同时实现了更高的性能，优于原始 BERT。
• 模型在剪枝后的序列上仍保持高预测置信度，表明其有效保留了语义上有意义的标记。
• 基于强化学习的标记选择策略能够无需人工标注即学习保留任务相关标记，如疑问词和情感指示词。
• 该方法在多种 NLP 任务中表现出鲁棒性，包括文本分类、自然语言蕴含和问答任务。