PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination 论文解读

一、论文基本信息论文题目PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination作者Saurabh Goyal、Anamitra Roy Choudhury、Saurabh M. Raje、Venkatesan T. Chakaravarthy、Yogish Sabharwal、Ashish Verma发表会议ICML 2020官方代码IBM/PoWER-BERTGitHub 仓库说明这是该论文的实现并标注论文发表于 ICML 2020。(GitHub)这篇论文的核心目标是不剪 BERT 的参数不删 layer不删 attention head而是逐层删除冗余 token 对应的 word-vector从而减少后续 Transformer 层的计算量。论文摘要中明确说PoWER-BERT 通过利用中间 Transformer block 输出中的word-vector redundancy在推理阶段逐步删除冗余 word-vectors在 GLUE benchmark 上相比 BERT-base 最多获得4.5× 推理时间减少精度损失小于1%应用到 ALBERT 上时最多获得6.8× 推理时间减少。(Proceedings of Machine Learning Research)二、论文要解决的问题BERT 推理慢的一个重要原因是每一层都要处理完整输入序列的所有 token 表示。例如输入长度是 128BERT-base 有 12 层那么每一层都要处理 128 个 token 向量总共相当于处理12 层 × 128 个 word-vectors传统压缩方法通常从参数角度入手例如剪权重。剪 attention heads。删 Transformer layers。知识蒸馏成浅层模型。量化。但是 PoWER-BERT 认为BERT 里还有另一种冗余word-vector 层面的冗余。也就是说随着 self-attention 一层层传播信息会在不同 token 表示之间扩散。到中后层时很多 token 向量已经携带相似信息不一定都需要继续向后计算。论文明确说由于 self-attention 机制存在信息扩散word-vectors 在经过多层 Transformer 后会变得冗余因此可以随着层数加深逐步删除一部分 word-vectors。(Proceedings of Machine Learning Research)所以这篇论文要解决的问题是能不能在推理过程中逐层减少 token 数量而不是让每一层都处理完整序列三、核心思想PoWER-BERT 的核心思想可以概括为前几层保留较多 token 表示后几层逐步删除不重要 token 表示只让少量重要 word-vectors 继续向后传播。这里的 word-vector 指的是某一层 Transformer block 输出的某个 token 对应的 hidden vector。它不是词向量表里的 embedding 参数而是每一层中间产生的 token representation。例如输入序列长度是 128第一层可能保留 80 个 word-vectors第二层保留 73 个第三层保留 70 个越往后保留越少。论文给出的示意配置类似80 → 73 → 70 → 50 → 50 → 40 → 33 → 27 → 20 → 15 → 13 → 3这就是progressive word-vector elimination。(Proceedings of Machine Learning Research)它的直觉是浅层还需要处理比较完整的词级信息。深层已经完成了一定信息聚合不必让所有 token 都继续参与计算。最后分类任务通常只需要 [CLS] 表示因此很多普通 token 可以提前停止计算。论文也特别说明[CLS] token 永远不会被删除因为最终预测仍然由 [CLS] 表示产生。(Proceedings of Machine Learning Research)四、它剪的是什么PoWER-BERT 剪的是中间层的 token-level word-vectors。它不剪模型权重。attention heads。FFN neurons。Transformer layers。hidden dimensions。所以它和你前面看的很多剪枝方法不一样。Head pruning是删除完整 attention head。LayerDrop / Poor Man’s BERT是删除 Transformer layer。Movement Pruning是删除单个权重。CoFi是删除 MHA / FFN 子层、heads、FFN dimensions 等结构。PoWER-BERT则是删除中间层 token 表示。论文也强调PoWER-BERT 和很多基于参数删除的方法正交因为它保留所有模型参数只删除冗余 word-vectors因此它可以叠加到其他压缩模型上比如 ALBERT。(Proceedings of Machine Learning Research)所以它更准确的定位是动态 token pruning / progressive token elimination。不是传统参数剪枝。五、为什么 word-vector 可以删PoWER-BERT 的核心假设是self-attention 会让信息在 token 之间扩散。在 BERT 的 self-attention 中每个 token 表示都会从其他 token 表示中聚合信息。经过多层之后一个 token 向量不再只代表它自己而可能已经携带了句子中其他词的信息。因此某些 token 即使后续不再参与计算它们的信息也可能已经被其他保留下来的 token 表示吸收了。例如情感分类中输入句子里有很多停用词、标点、功能词。它们在前几层可能有用但到后几层时真正决定分类的可能只是少数情感词或核心短语。论文给出的 SST-2 案例中早期层会删除 stop words 和标点后续层继续保留携带句子情感信息的词。(Proceedings of Machine Learning Research)所以 PoWER-BERT 的基本判断是不是每个 token 都需要走完整个 BERT。六、如何判断哪个 word-vector 重要PoWER-BERT 用的是attention-based significance score。直观地说如果一个 token 向量经常被其他 token 注意到说明它对其他 token 的表示更新很重要。如果一个 token 向量几乎没人关注说明它对后续信息传播贡献较小可以优先删除。具体做法是对于某个 attention head统计所有 query 位置对某个 token 的 attention 总量再把所有 heads 的结果加起来得到这个 token 的总体 significance score。论文明确说一个 word-vector 的 significance score 是它在所有 heads 中对其他 word-vectors 施加的 attention 总量分数越高说明该 word-vector 对最终预测影响越大。(Proceedings of Machine Learning Research)换成更容易理解的话PoWER-BERT 不是看某个 token 自己关注了谁而是看有多少其他 token 需要它。如果很多 token 都从它那里读取信息它就重要。如果几乎没有 token 读取它它就冗余。七、Word-vector Selection每一层具体删哪些 token假设第 j 层决定只保留若干个 word-vectors那么 PoWER-BERT 会先计算当前层所有 word-vectors 的 significance score。按分数从高到低排序。保留 top-k 个 word-vectors。删除剩余低分 word-vectors。论文把这个操作放在 self-attention 和 FFN 之间通过一个extract layer完成。这个 extract layer 会根据 significance score 保留 top word-vectors。(Proceedings of Machine Learning Research)为什么放在 self-attention 后面因为 significance score 依赖当前层 self-attention 矩阵。也就是说模型先通过 self-attention 判断当前 token 之间的信息依赖再根据这个依赖关系决定哪些 token 可以继续向后传。所以它不是静态删词而是每个输入样本都有自己的 token 保留结果。这点非常重要。比如同样是情感分类不同句子里的关键词不同因此应该保留的 token 也不同。PoWER-BERT 的 attention-based selection 是输入自适应的。八、Retention Configuration每一层保留多少 token前面解决的是给定保留数量后保留哪些 token。但还有一个问题每一层到底应该保留多少 token这就是论文中的retention configuration。它是一个单调递减序列l1, l2, ..., l12其中 lj 表示第 j 个 Transformer block 保留多少 word-vectors。论文明确说PoWER-BERT 的 retention configuration 是一个单调递减序列用来指定每个 Transformer block 保留的 word-vector 数量。(Proceedings of Machine Learning Research)为什么必须单调递减因为 token 一旦删除后面层就不会再恢复。所以下一层能处理的 token 数量只能小于或等于上一层。为了学习这个配置论文引入了soft-extract layer。soft-extract layer 不直接删除 token而是给排序后不同位置的 word-vector 乘上一个可学习的 retention parameter。高排名 token 通常应该保留得更多低排名 token 可能被压小。论文说明soft-extract layer 会保留所有 word-vectors但根据它们在 significance 排序中的位置用可学习参数控制保留程度。(Proceedings of Machine Learning Research)训练时模型损失由两部分组成任务损失保证精度。retention regularizer鼓励保留更少 word-vectors尤其鼓励后层减少计算量。论文中通过超参数 λ 控制精度和推理时间之间的 trade-offλ 越大模型越倾向于删除更多 word-vectors。(Proceedings of Machine Learning Research)九、训练流程PoWER-BERT 的训练流程分三步。第一步Fine-tuning。先从预训练 BERT 出发在目标任务上正常 fine-tune得到普通任务模型。第二步Configuration Search。在 fine-tuned BERT 中插入 soft-extract layers并修改 loss使模型同时考虑任务精度和 word-vector retention cost。训练后得到每层应该保留多少 word-vectors也就是 retention configuration。第三步Re-training。把 soft-extract layer 替换成真正的 extract layer。推理时根据 attention-based significance score 动态选择 top word-vectors然后重新训练模型以恢复精度。论文明确说三个训练步骤在实验中总共只需要2–3 epochs最终推理使用 re-trained PoWER-BERT 模型。(Proceedings of Machine Learning Research)这个流程的本质是先学“每层保留多少”。再用 attention score 动态决定“每个样本保留哪些”。十、它和 Early Exit 有什么区别PoWER-BERT 不是 early exit。Early exit 方法比如 DeeBERT是让整个输入样本在某一层提前退出。如果模型在第 6 层就足够自信那么后面 6 层都不跑了。PoWER-BERT 不一样样本不会整体提前退出。模型仍然可以跑完整 12 层。但越来越少的 token 会继续向后计算。所以两者区别是Early Exit 是 sample-level depth reduction。PoWER-BERT 是 token-level sequence reduction。举例Early exit这个句子第 6 层就结束。PoWER-BERT这个句子仍然走到第 12 层但后面层只处理少数核心 token。这使 PoWER-BERT 特别适合长序列因为 self-attention 和 FFN 的计算都和 token 数量有关。Token 数量越早减少后续层节省越多。十一、它和 Head Pruning 有什么区别Head pruning 删除的是 attention 内部的多个 head。PoWER-BERT 删除的是 token 表示。两者完全不同。Head pruning 减少每个 token 的 attention 分支。PoWER-BERT 减少后续层需要处理的 token 数量。从计算角度看PoWER-BERT 的收益可能更直接。因为一个 token 被删掉后它不再参与后续层的 self-attention也不再经过后续 FFN。论文也比较了 Head-Prune发现 Head-Prune 在推理时间和精度 trade-off 上不如 PoWER-BERT。论文解释和实验显示PoWER-BERT 在多个 GLUE 任务上相较 prior methods 有更好的 Pareto trade-off。(Proceedings of Machine Learning Research)十二、实验设置论文主要在分类和回归任务上做实验包括 GLUE 任务以及 IMDB 和 RACE。论文表格列出的任务包括CoLARTEQQPMRPCSST-2MNLI-mMNLI-mmQNLISTS-BIMDBRACE并为每个任务设置不同最大输入长度例如 CoLA 和 SST-2 使用 64QQP / MRPC / MNLI / QNLI 使用 128RTE 使用 256IMDB 和 RACE 使用 512。(Proceedings of Machine Learning Research)实现方面论文代码基于 Keras推理时间实验在K80 GPU上进行batch size 大多数任务为 128RACE 为 32并对 100 次运行取平均。(Proceedings of Machine Learning Research)这一点要注意论文中的加速结果和具体硬件、框架、batch size、序列长度有关不能直接等同于所有部署环境下的速度。十三、主要实验结果13.1 相比 BERT-base最多 4.5× 推理加速论文最核心结果是在精度损失小于 1% 的约束下PoWER-BERT 相比 BERT-base 在所有数据集上至少获得 2.0× 推理时间减少最高达到 4.5×。论文在 RTE 例子中给出具体分析RTE 输入长度为 256BERT-base 12 层需要处理 12 × 256 3072 个 word-vectors而 PoWER-BERT 的一个 retention configuration 是153, 125, 111, 105, 85, 80, 72, 48, 35, 27, 22, 5总共只处理 868 个 word-vectors因此计算量显著减少。(Proceedings of Machine Learning Research)这说明 PoWER-BERT 的加速不是靠减少参数而是靠减少每层实际处理的 token 数量。13.2 相比 DistilBERT、BERT-PKD、Head-Prunetrade-off 更好论文把 PoWER-BERT 与 DistilBERT、BERT-PKD 和 Head-Prune 做了比较。结果显示PoWER-BERT 在 accuracy–inference time Pareto curve 上明显更优。论文报告在给定推理时间下PoWER-BERT 在 CoLA 上最高可带来 16% accuracy gain在 RTE 上最高带来 6% accuracy gain在给定准确率下PoWER-BERT 在 CoLA 上最高可带来 2.7× 推理时间收益在 RTE 上最高带来 2.0× 推理时间收益。(Proceedings of Machine Learning Research)这个结果说明细粒度 token elimination 比粗粒度删 Transformer block 更灵活。DistilBERT / BERT-PKD 删除的是整层因此压缩粒度比较粗。PoWER-BERT 可以在不同层保留不同数量 token也可以对不同输入动态选择不同 token所以速度–精度折中更细。13.3 应用于 ALBERT 也有效论文还把 PoWER-BERT 应用于 ALBERT。ALBERT 本身已经通过参数共享和 factorized embedding 做了压缩但 PoWER-BERT 仍然可以进一步加速它。论文结果显示PoWER-BERT 在 ALBERT 上多数 GLUE 数据集能获得约 2× 加速QQP 上最高达到6.8×精度损失小于 1%。(Proceedings of Machine Learning Research)这说明PoWER-BERT 和参数压缩方法是互补的。即使模型参数已经很少只要中间 token 表示存在冗余仍然可以通过 word-vector elimination 加速。13.4 Attention-based selection 优于静态选择论文还比较了几种 word-vector selection 方法。Head-WS保留序列开头的 word-vectors。Rand-WS随机保留 word-vectors。Attn-WS基于 attention significance score 动态保留 word-vectors。在 SST-2 上同样使用固定 retention configuration 时Attn-WS 的准确率明显更高。论文表 4 中整个数据集上 Head-WS 是 85.4%Rand-WS 是 85.7%Attn-WS 是 88.3%对长度大于 16 的输入Attn-WS 也明显优于静态方法。(Proceedings of Machine Learning Research)这个结果说明动态选择哪些 token 保留比简单保留前几个 token 或随机保留更可靠。十四、方法优点第一它发现了 BERT 中另一种冗余token-level word-vector redundancy。之前很多方法关注参数冗余PoWER-BERT 关注中间表示冗余这个角度很有启发。第二它保留模型参数不改变权重结构。这使它可以叠加到其他压缩模型上例如 ALBERT。论文也明确强调它和参数删除方法是正交的。(Proceedings of Machine Learning Research)第三它是输入自适应的。不同句子保留的 token 不同因此比静态删词更灵活。第四它逐层减少 token 数量速度收益直接。一个 token 一旦被删后续所有层都不用再处理它节省的是连续多层计算。第五速度–精度 trade-off 好。论文结果显示它在多项任务上优于 DistilBERT、BERT-PKD 和 Head-Prune 的推理时间–精度折中。(Proceedings of Machine Learning Research)十五、方法局限第一它主要适合分类 / 回归任务。论文实验集中在 GLUE、IMDB、RACE 这类分类或选择任务。对于生成任务、序列标注任务、抽取式问答等需要保留 token-level 输出的任务直接删除 token 会更复杂。第二[CLS] 分类假设很重要。PoWER-BERT 依赖最终 [CLS] 表示做预测所以普通 token 后期可以被删。如果任务需要每个 token 的最终表示例如 NER 或 token classification就不能简单删掉普通 token。第三动态 token 删除会带来实现复杂度。虽然计算量减少了但不同输入、不同层的 token 数量不同实际部署时需要支持动态 shape、gather/scatter、mask 更新等操作。不同推理框架下速度收益可能不一样。第四它不减少模型参数量。PoWER-BERT 保留所有 BERT 参数因此模型文件大小不会明显变小。它主要优化推理时间而不是存储压缩。第五剪掉 token 的决策依赖 attention score。attention score 是一个合理 proxy但不一定总能完全代表 token 对最终预测的真实贡献。某些 token 可能 attention 不高但对局部语义或对抗样本很关键。十六、和后续 Token Pruning 方法的关系PoWER-BERT 可以看作 NLP 中比较早的dynamic token pruning方法之一。它的思想和后来 ViT / VLM 里的 token pruning 非常接近先让模型在浅层处理完整 token 集。然后根据重要性分数删除冗余 token。后续层只处理保留下来的 token。区别是PoWER-BERT 的 token 是文本 word-piece token。ViT token pruning 的 token 是图像 patch token。VLM token pruning 的 token 通常是视觉 token 或多模态 token。这个思路后来在 Transformer 高效推理中非常常见计算瓶颈不一定只来自参数数量也来自序列长度。只要能减少中后层序列长度就能显著减少推理计算。十七、整体评价PoWER-BERT 的核心贡献是把 BERT 推理加速从“删参数”转向“删中间 token 表示”。它的逻辑非常清楚self-attention 会扩散信息。随着层数加深很多 token 表示变得冗余。如果某个 token 很少被其他 token 关注它可以提前停止计算。逐层减少 token 数量可以显著降低推理时间。从模型压缩角度看它不是传统剪枝从推理加速角度看它非常重要因为它直接作用于 Transformer 的一个核心成本来源sequence length。它尤其适合这样的场景输入序列较长。任务最终只需要句级表示。允许动态推理路径。目标是降低推理时间而不是减少模型文件大小。十八、一句话总结《PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination》提出一种逐层删除冗余 token 表示的 BERT 推理加速方法它不剪参数、不删 layer、不删 head而是根据 self-attention 计算每个 word-vector 的重要性在每层保留高分 token、删除低分 token使后续 Transformer 层处理更短的序列实验表明在 GLUE 等任务上 PoWER-BERT 相比 BERT-base 最多获得 4.5× 推理时间减少精度损失小于 1%并且还能叠加到 ALBERT 上进一步加速。