[2107.00910] Learned Token Pruning for Transformers

“Learned Token Pruning for Transformers”由Sehoon Kim等人撰写。在Transformer模型应用中，其推理成本随输入序列长度呈二次增长，高效部署面临挑战。本文提出基于阈值的自适应令牌修剪方法LTP，通过可微软二值化掩码学习阈值，在不影响准确性的前提下减少计算量。实验表明，LTP在GLUE和SQuAD任务上表现优异，相比现有方法，在相同FLOPs下精度更高，能大幅降低FLOPs并提升吞吐量，对输入序列长度变化更具鲁棒性，且与量化和知识蒸馏等方法兼容。

研究背景

  • Transformer模型的应用与挑战：基于Transformer的模型如BERT和RoBERTa在NLP任务中取得优异成绩，但因其规模大、推理成本高，在实际部署中面临内存、延迟和能耗等挑战。其核心的多头自注意力机制使计算复杂度与输入序列长度呈二次关系，限制了长序列应用。

  • 现有修剪方法的局限：修剪是减少神经网络规模和计算量的常用技术，分为非结构化和结构化修剪。非结构化修剪在硬件上难以高效部署，结构化修剪虽更实用，但现有令牌修剪方法存在不足。部分方法为数据集固定单一修剪配置，易导致欠修剪或过修剪；一些方法虽考虑输入序列长度调整配置，但缺乏对内容的适应性或训练成本高，且多数依赖昂贵的top - k操作。

LTP方法介绍

  • 阈值令牌修剪：定义令牌重要性分数为各头注意力概率均值，采用阈值策略修剪分数低于阈值的令牌，通过二元掩码确定保留或修剪，此操作简单高效，避免了top - k计算。

  • 可学习阈值机制：为解决阈值确定问题，采用软修剪方案，用sigmoid操作的可微软掩码替代非可微掩码，并乘以层输出激活实现类似硬修剪效果，使阈值可通过反向传播学习。同时添加L1正则化项惩罚未修剪令牌，推动阈值学习，调整修剪比例。

实验设置与结果

  • 实验设置：在RoBERTa_base上实现LTP，使用HuggingFace库，在GLUE和SQuAD 2.0任务上测试，以分类准确率、F1分数等指标评估。训练分软修剪和硬修剪阶段，与SpAtten和LAT对比，使用PyTorch 1.8，在特定CPU和GPU环境进行推理速度实验，避免输入序列额外填充。

  • 性能评估：LTP在多数任务中，在1%精度下降内实现平均1.96倍（最高2.10倍）的FLOPs降低，在相同FLOPs下比SpAtten精度高约2%，比LAT在多数任务中高约2.5%。在SQuAD 2.0上也有类似良好表现，且在不同子集上F1分数变化符合预期，需合理设置平衡精度。

  • 鲁棒性分析：针对序列长度变化，在QNLI和QQP任务上训练和评估LTP与LAT，结果显示LTP在不同长度序列上性能更优，如在QNLI和QQP的长序列评估中，LTP比LAT精度分别高16.44%和9.20%。

  • 消融研究：手动设置阈值实验表明，学习阈值方法在相同FLOPs下性能优于手动设置，证明阈值学习的有效性。

  • 硬件吞吐量测量：在NVIDIA V100 GPU和Intel Haswell CPU上测量吞吐量，通过合理处理修剪后序列，LTP在不同批次大小下，在QNLI和QQP任务上相比基线实现高达约1.9倍和2.0倍的提升，且速度提升随批次增大而增加。

  • 与其他方法结合：将LTP与量化和知识蒸馏结合，对模型参数量化并以RoBERTa_base为教师模型蒸馏知识，在QQP和SST - 2数据集上实现高达10倍的BOPs降低且精度下降小于2%，证明兼容性。