[論文レビュー] 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% の向上）を達成し、$N$ 台のデバイスに対して $Θ(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 圧縮手法が、各新しいハードウェア制約に対して微調整または再圧縮を必要としているスケーラビリティと再圧縮のオーバーヘッドを解消すること。
• BERT のニューラルアーキテクチャサーチにおいて、純粋な活用（例：DynaBERT）や純粋な探索（例：NAS-BERT）の限界を克服すること。
• 多様なリソース制限のあるデバイスにわたって「一度圧縮して、どこでもデプロイ」を可能にする統合圧縮パイプラインを開発すること。
• 大規模で構造化された探索空間において、アーキテクチャの探索と活用のバランスを取ることで、デプロイ時の性能劣化を最小限に抑えること。

提案手法

• 幅、深さ、ヘッドの変異を含む、ほぼすべての BERT の構成をカバーする巨大な $10^{13}$ 個のアーキテクチャ探索空間を構築する。
• 探索と活用のバランスを取るための、確率的自然勾配最適化である「探索・活用確率分布のためのステノックス自然勾配最適化（EE-SNG）」を導入する。
• プロセスを、全モデルのトレーニングフェーズ（「スーパーモード」）と、ターゲットデバイス向けの軽量な分布ベースの適応フェーズに分離する。
• 性能フィードバックに基づき、トレーニング中に動的に探索と活用のバランスを制御するための学習可能コントローラーを用いる。
• 微調整を一切行わず、スーパーベルトモデルから直接重みを継承することで、デプロイコストを $\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 デプロイでどれほど優れているか？

主な発見

• YOCO-BERT は、最先端の圧縮手法と比較して、GLUE ベンチマークで平均 2.1% から 4.5% の精度向上を達成した。
• 0.75x の圧縮比の RTE タスクにおいて、探索のみのベースラインと比較して、YOCO-BERT は 9.8 パcentage points の優位性を示した。
• 提案された EE-SNG 最適化は、図 3 に示すように、ランダム探索や遺伝的探索よりも高速に収束し、より高い精度に到達した。
• 探索後にサブアーキテクチャを微調整しても、顕著な向上は得られず、場合によっては性能が劣化し、RTE で最大 5.4% の精度低下が観測された。
• N 個の異なるデバイスをサポートするためのトレーニング複雑度は、1 回のスーパーモデルトレーニングで十分であるため、$\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).