[论文解读] Magic Pyramid: Accelerating Inference with Early Exiting and Token Pruning

Xuanli He, Iman Keivanloo+5 more|arXiv (Cornell University)|Oct 30, 2021

Topic Modeling参考文献 22被引用 5

一句话总结

Magic Pyramid (MP) 通过协同结合标记剪枝和提前退出，加速 BERT 推理，在宽度和深度上均减少计算量。其在精度损失低于 0.5% 的情况下实现最高 11.95 倍的加速，相较于最先进方法，速度提升最高达 2.13 倍，且在不同输入长度下均保持稳健性能。

ABSTRACT

Pre-training and then fine-tuning large language models is commonly used to achieve state-of-the-art performance in natural language processing (NLP) tasks. However, most pre-trained models suffer from low inference speed. Deploying such large models to applications with latency constraints is challenging. In this work, we focus on accelerating the inference via conditional computations. To achieve this, we propose a novel idea, Magic Pyramid (MP), to reduce both width-wise and depth-wise computation via token pruning and early exiting for Transformer-based models, particularly BERT. The former manages to save the computation via removing non-salient tokens, while the latter can fulfill the computation reduction by terminating the inference early before reaching the final layer, if the exiting condition is met. Our empirical studies demonstrate that compared to previous state of arts, MP is not only able to achieve a speed-adjustable inference but also to surpass token pruning and early exiting by reducing up to 70% giga floating point operations (GFLOPs) with less than 0.5% accuracy drop. Token pruning and early exiting express distinctive preferences to sequences with different lengths. However, MP is capable of achieving an average of 8.06x speedup on two popular text classification tasks, regardless of the sizes of the inputs.

研究动机与目标

• 为解决大型预训练模型（如 BERT）在实时生产环境约束下的高推理延迟问题。
• 克服现有方法的局限性——标记剪枝（在长序列上有效）和提前退出（在短序列上有效）——在序列长度谱系两端表现不佳的问题。
• 开发一种统一框架，协同结合标记剪枝和提前退出，实现在可变输入长度下一致且可调节速度的推理。
• 在实现显著计算量减少的同时保持高精度，尤其适用于低延迟、高吞吐量的部署环境。

提出的方法

• MP 将标记剪枝和提前退出整合到单一分层推理流水线中，其中剪枝减少宽度方向的计算，提前退出减少深度方向的计算。
• 模型采用可学习的阈值机制，基于注意力分数的不确定性，识别并移除非显著标记。
• 在每个 Transformer 块处附加一个子分类器，当置信度超过动态阈值 τ 时可实现提前预测，该阈值在训练过程中进行调优。
• 该方法采用温度缩放和损失加权（λ）来平衡剪枝与提前退出联合优化过程中的精度与速度权衡。
• 训练过程联合优化精度、GFLOPs 和提前退出概率，使用包含知识蒸馏和不确定性正则化的多任务损失函数。
• 该框架端到端应用于基于 BERT 的模型，推理过程根据输入长度动态调整，通过剪枝标记和在满足条件时提前退出。

实验结果

研究问题

• RQ1 标记剪枝与提前退出能否有效结合，实现在不同输入长度下的持续推理加速？
• RQ2 宽度方向（剪枝）与深度方向（提前退出）计算量减少的协同效应是否优于单一方法？
• RQ3 该联合方法能否在实现显著计算节省的同时保持高精度，尤其在低延迟生产环境中？
• RQ4 所提方法在不同序列长度和 NLP 任务上的性能表现如何？

主要发现

• 在 AG News 和 Yelp 数据集上，MP 实现最高 11.95 倍加速，显著优于 FastBERT（仅提前退出）和 LTP（仅标记剪枝）在所有序列长度组中的表现。
• 在两个文本分类任务上，MP 平均实现 8.06 倍加速，无论输入长度如何，均表现出一致的性能，证明其在多样化输入下的稳定性。
• 与 BERT 相比，MP 将 GFLOPs 最多减少 70%，精度损失低于 0.5%，在效率-精度权衡上优于 LTP 和 FastBERT。
• 对于长序列（>70 个标记），MP 在 Yelp 上实现 8.25 倍加速，在 AG News 上实现 11.95 倍加速，显著优于 FastBERT 的 6.18 倍和 8.84 倍。
• 当 τ = 0.8 时，MP 在 AG News 上实现 11.95 倍加速，在 Yelp 上实现 10.10 倍加速，表明在激进的提前退出阈值下仍具备强大可扩展性。
• MP 在精度上保持或略有提升（例如 AG News 上为 94.3%），同时将 GFLOPs 降低至 1.8（实现 4.95 倍加速），优于 FastBERT 的 2.3 GFLOPs（3.97 倍加速）。