[論文レビュー] A Fast Post-Training Pruning Framework for Transformers

Woosuk Kwon, Sehoon Kim+4 more|arXiv (Cornell University)|Mar 29, 2022

Advanced Neural Network Applications被引用数 38

ひとこと要約

再訓練なしで大幅なFLOPsおよびレイテンシ削減を実現するTransformer用の剪定フレームワーク。Fisherベースのマスク探索、再配置、マスク調整を用いて、精度を保ちながらヘッドとFFNフィルターを剪定する。

ABSTRACT

Pruning is an effective way to reduce the huge inference cost of Transformer models. However, prior work on pruning Transformers requires retraining the models. This can add high training cost and high complexity to model deployment, making it difficult to use in many practical situations. To address this, we propose a fast post-training pruning framework for Transformers that does not require any retraining. Given a resource constraint and a sample dataset, our framework automatically prunes the Transformer model using structured sparsity methods. To retain high accuracy without retraining, we introduce three novel techniques: (i) a lightweight mask search algorithm that finds which heads and filters to prune based on the Fisher information; (ii) mask rearrangement that complements the search algorithm; and (iii) mask tuning that reconstructs the output activations for each layer. We apply our method to BERT-base and DistilBERT, and we evaluate its effectiveness on GLUE and SQuAD benchmarks. Our framework achieves up to 2.0x reduction in FLOPs and 1.56x speedup in inference latency, while maintaining < 1% loss in accuracy. Importantly, our framework prunes Transformers in less than 3 minutes on a single GPU, which is over two orders of magnitude faster than existing pruning approaches that retrain the models.

研究の動機と目的

• FLOPs/レイテンシの制約下での実用的な、再訓練不要なTransformerモデルの圧縮を推進する。
• 情報理論に導かれた、ヘッドとFFNフィルターを剪定するためのマスク探索、再配置、調整からなる3段階の剪定パイプラインを開発する。
• 厳密な予算内で精度を維持しつつ、データの小さなサンプル上での迅速な剪定を実現し、全面的な再訓練を回避する。

提案手法

• 剪定を binary masks を用いたMHAヘッドとFFNフィルター上の制約付きマスク最適化として表現する。
• Fisher情報ベースのマスク探索を用いて、FLOPs/レイテンシ予算の下で剪定するヘッド/フィルターを選択する。
• ブロック対角Fisher近似によって層内相互作用を捉えるためのマスク再配置段を適用する。
• 層ごとの線形最小二乗問題を解くことで層の活性化を再構成するマスク調整段を実行する。
• レイテンシ制約へ拡張するため、レイテンシをピースワイズ線形モデルで近似し、それに応じて探索を適応させる。
• BERT_BASEとDistilBERTで剪定を実証し、GLUEとSQuADで最小限の精度低下とともに評価する。

実験結果

リサーチクエスチョン

• RQ1 再訓練なしでFLOPs/レイテンシの制約を満たすようにTransformerモデルを剪定できるか？
• RQ2 与えられたリソース制約の下で精度低下を最小化するように、剪定するヘッドとFFNフィルターをどのように特定できるか？
• RQ3 マスク探索、再配置、調整を含むポストトレーニング剪定パイプラインは、再訓練ベースの剪定手法と比べて効率と精度のトレードオフで優れるか？
• RQ4 一般的なベンチマークで、再訓練なしの剪定を使用した場合、実機で実現可能な実用的なスピードアップはどれくらいか？

主な発見

• 本フレームワークはFLOPsを最大で2.0×、推論のレイテンシを最大で1.56×削減し、精度低下は1%未満で収める。
• 剪定は単一GPUで3分未満で完了し、再訓練ベースの剪定手法よりも2桁以上高速である。
• GLUEとSQuADで、BERT_BASEとDistilBERTを1%の精度損失で剪定すると、FLOPsの大幅な削減を達成する（複数タスクでBERT_BASEは元のFLOPsの60–70%、DistilBERTは約50%）。
• NVIDIA V100での待機時間の実験は、1%の精度制約下でバッチサイズ256で平均1.47×〜1.56×の速度向上を示す。
• 提案されたFisherベースのマスク探索、再配置、調整段階のそれぞれが精度の回復または維持に寄与し、特にマスク調整が精度回復において重要な役割を果たす。
• 既存の構造的剪定手法と比較して、再訓練不要アプローチは、FLOPsと精度のトレードオフで同等またはより良い結果を達成し、剪定コストを大幅に低減している（エンドツーエンドの剪定時間は1分未満）。