[論文レビュー] PLATON: Pruning Large Transformer Models with Upper Confidence Bound of Weight Importance

Qingru Zhang, Simiao Zuo+5 more|arXiv (Cornell University)|Jun 25, 2022

Topic Modeling被引用数 11

ひとこと要約

PLATONは、感度の上側信頼区間（UCB）を用いて重みの重要度推定に不確実性を組み込むことで、学習安定性と性能を向上させる構造的プルーニング手法を提案する。時間的変動による不確実性の定量化と指数的平滑化された感度を組み合わせることで、重要度スコアのばらつきを低減し、より信頼性の高いプルーニングを実現する。特に高スパarsity領域で顕著で、SQuADにおいて90%のスパarsityでF1スコアが5%以上向上する。

ABSTRACT

Large Transformer-based models have exhibited superior performance in various natural language processing and computer vision tasks. However, these models contain enormous amounts of parameters, which restrict their deployment to real-world applications. To reduce the model size, researchers prune these models based on the weights' importance scores. However, such scores are usually estimated on mini-batches during training, which incurs large variability/uncertainty due to mini-batch sampling and complicated training dynamics. As a result, some crucial weights could be pruned by commonly used pruning methods because of such uncertainty, which makes training unstable and hurts generalization. To resolve this issue, we propose PLATON, which captures the uncertainty of importance scores by upper confidence bound (UCB) of importance estimation. In particular, for the weights with low importance scores but high uncertainty, PLATON tends to retain them and explores their capacity. We conduct extensive experiments with several Transformer-based models on natural language understanding, question answering and image classification to validate the effectiveness of PLATON. Results demonstrate that PLATON manifests notable improvement under different sparsity levels. Our code is publicly available at https://github.com/QingruZhang/PLATON.

研究の動機と目的

• ミニバッチ上で推定される重要度スコアの高い分散が原因で生じる反復的プルーニングにおける学習不安定性を解消すること。
• 訓練中に感度推定がノイズを含むため、重要な重みが誤ってプルーニングされるリスクを低減すること。
• 大規模なTransformerモデルにおける高スパarsityプルーニング領域における一般化性能と安定性を向上させること。
• 不確実性の高い低重要度重みの探索と、高重要度重みの活用のバランスを取る手法を開発すること。
• グループレベルの感度スコアにおける不確実性を定量化することで、より信頼性の高い構造的プルーニングを可能にすること。

提案手法

• PLATONは、重み勾配と重み値の内積の絶対値として感度を計算する：$ I = |\theta^\top \nabla_\theta \mathcal{L}(\theta) | $。
• 感度を時間的に平滑化するために指数移動平均（EMA）を適用する：$ \overline{I}^{(t+1)} = \beta_1 \overline{I}^{(t)} + (1 - \beta_1) I^{(t)} $。
• 不確実性は平滑化された感度からの絶対偏差として定量化する：$ U^{(t)} = |I^{(t)} - \overline{I}^{(t)}| $。
• 不確実性に対しても2番目のEMAを適用する：$ \overline{U}^{(t+1)} = \beta_2 \overline{U}^{(t)} + (1 - \beta_2) U^{(t)} $。
• 最終的な重要度スコアは、平滑化された感度と不確実性の積として得る：$ S^{(t)} = \overline{I}^{(t)} \cdot \overline{U}^{(t)} $。
• 重要度スコアが最も低い重みまたはグループをゼロ化することでプルーニングを実行し、上位-r%の選択戦略を用いる。

実験結果

リサーチクエスチョン

• RQ1 大規模なTransformerモデルにおける重要度スコア推定の不確実性を効果的にモデル化することで、プルーニングの安定性を向上させることができるか？
• RQ2 上側信頼区間（UCB）による不確実性の組み込みが、反復的プルーニングにおける感度スコアのばらつきを低減するか？
• RQ3 UCBベースの重要度スコアは、高スパarsityプルーニング環境において、標準的な感度またはマグニチュードベースの指標を上回る性能を示すか？
• RQ4 構造的プルーニング（例：列全体をプルーニング）において、本手法はどのように性能を発揮するか？
• RQ5 従来の反復的プルーニングベースラインと比較して、本手法は学習不安定性をどの程度低減し、一般化性能を向上させるか？

主な発見

• DeBERTaV3-baseを90%のスパarsityにプルーニングした際、SQuAD v1.1ベンチマークでF1スコアが5.0%以上向上し、先行手法を著しく上回った。
• SQuAD v1.1で微調整されたBERT-baseでは、残存パラメータが50%の状態でEMが78.5%、F1が87.2%を維持し、ベースライン手法を上回った。
• 図1（下段）に示すように、PLATONにおける重要度スコアのばらつきは、標準的な感度ベースのプルーニングと比較して著しく低減されており、学習の安定性が向上している。
• 既存の反復的プルーニング手法（例：movement pruning）と比較して、PLATONの訓練プロセスはより安定しており、ハイパーパrameterの選択に敏感ではない。
• 構造的プルーニングにおいて、SQuAD v1.1で残存パラメータが40%の状態でEMが78.0%、F1が86.9%を達成し、非構造的ベースラインを上回った。
• 高い不確実性と低い感度を示す重みを保持することで、本手法は重要な重みの過剰な早期プルーニングを効果的に防止し、モデル容量のより良い探索を可能にした。