[论文解读] You Only Compress Once: Towards Effective and Elastic BERT Compression via Exploit-Explore Stochastic Nature Gradient

Shaokun Zhang, Xiawu Zheng+8 more|arXiv (Cornell University)|Jun 4, 2021

Topic Modeling参考文献 33被引用 11

一句话总结

YOCO-BERT 提出了一种新颖的一次性压缩框架，用于 BERT 模型，使其能够在无需微调的情况下弹性部署于各种硬件约束下。通过在 $10^{13}$ 个架构的搜索空间中采用具有探索-利用平衡的随机自然梯度优化，该方法实现了最先进水平的准确率提升（GLUE 平均提升 2.1–4.5%），并通过直接从超 BERT 模型继承权重，将 $N$ 个设备的训练成本降低至 $/mathcal{O}(1)$。

ABSTRACT

Despite superior performance on various natural language processing tasks, pre-trained models such as BERT are challenged by deploying on resource-constraint devices. Most existing model compression approaches require re-compression or fine-tuning across diverse constraints to accommodate various hardware deployments. This practically limits the further application of model compression. Moreover, the ineffective training and searching process of existing elastic compression paradigms[4,27] prevents the direct migration to BERT compression. Motivated by the necessity of efficient inference across various constraints on BERT, we propose a novel approach, YOCO-BERT, to achieve compress once and deploy everywhere. Specifically, we first construct a huge search space with 10^13 architectures, which covers nearly all configurations in BERT model. Then, we propose a novel stochastic nature gradient optimization method to guide the generation of optimal candidate architecture which could keep a balanced trade-off between explorations and exploitation. When a certain resource constraint is given, a lightweight distribution optimization approach is utilized to obtain the optimal network for target deployment without fine-tuning. Compared with state-of-the-art algorithms, YOCO-BERT provides more compact models, yet achieving 2.1%-4.5% average accuracy improvement on the GLUE benchmark. Besides, YOCO-BERT is also more effective, e.g.,the training complexity is O(1)for N different devices. Code is availablehttps://github.com/MAC-AutoML/YOCO-BERT.

研究动机与目标

• 解决现有 BERT 压缩方法在面对新硬件约束时需微调或重新压缩所带来的可扩展性与重压缩开销问题。
• 克服纯利用（如 DynaBERT）或纯探索（如 NAS-BERT）在 BERT 神经架构搜索中的局限性。
• 开发统一的压缩流水线，实现跨多种资源受限设备的“一次压缩，处处部署”。
• 通过在大规模结构化搜索空间中平衡架构探索与利用，最小化部署过程中的性能下降。

提出的方法

• 构建一个覆盖几乎所有 BERT 配置（包括宽度、深度和头数变化）的 $10^{13}$ 个架构的超大规模搜索空间。
• 提出一种探索-利用随机自然梯度（EE-SNG）优化方法，联合训练权重以及利用/探索概率分布的参数。
• 将流程解耦为两个阶段：首先在“超模式”下训练完整模型，然后在目标设备上通过轻量级分布适应阶段进行部署。
• 使用可学习控制器，根据性能反馈动态平衡训练过程中的探索与利用。
• 直接从超 BERT 模型继承目标子架构的模型权重，无需微调，将部署成本降低至 $\mathcal{O}(1)$。
• 将 NAS 过程形式化为使用随机自然梯度的连续似然最大化问题，实现可微分且高效的搜索。

Figure 1: The proposed YOCO-BERT. We separate YOCO-BERT into two stages. Given expanded training space, we first propose a EE-SNG to optimize the probability distribution $P_{\theta}$ and super-BERT. In the second stage, searching algorithm is supervised by the optimized $P_{\theta}$ to obtain the o

实验结果

研究问题

• RQ1 一个统一的 BERT 压缩框架是否能在无需重新压缩或微调的情况下，实现对多样化硬件约束的高性能表现？
• RQ2 在 BERT 压缩中，如何有效建模并优化探索-利用平衡？
• RQ3 结合大规模结构化搜索空间与可微分优化，是否能生成优于随机搜索或进化搜索的子架构？
• RQ4 与微调后的子网络相比，直接从超 BERT 模型继承权重在压缩 BERT 部署中能多大程度上表现更优？

主要发现

• 与最先进压缩方法相比，YOCO-BERT 在 GLUE 基准上实现了 2.1% 至 4.5% 的平均准确率提升。
• 在压缩比为 0.75x 的 RTE 任务中，YOCO-BERT 比仅探索的基线方法高出 9.8 个百分点。
• 如图 3 所示，所提出的 EE-SNG 优化方法收敛更快且准确率更高，优于随机搜索与进化搜索。
• 在搜索后对子架构进行微调并未带来显著性能增益，反而可能损害性能，在 RTE 任务上准确率最高下降达 5.4%。
• 支持 $N$ 种不同设备的训练复杂度被降低至 $\mathcal{O}(1)$，因为仅需一次超模型训练。
• 该方法能直接从优化后的概率分布中识别出高性能子架构，消除了昂贵重训练的需求。

Figure 2: Comparison of parameters, FLOPs between YOCO-BERT and previous works under different datasets. We use accuracy as the metric for the three datasets. Since YOCO-BERT does not use external data, Dyanbert’s results are obtained without using data augmentation (DA).