[論文レビュー] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

Zhuohan Li, Eric Wallace+5 more|arXiv (Cornell University)|Feb 26, 2020

Topic Modeling参考文献 65被引用数 51

ひとこと要約

本論文は、非常に大規模なTransformerモデルを訓練し早期終了させることが、小規模モデルを収束させて訓練するよりも計算資源効率が高くなり得ること、そしてそのような大規模モデルはより圧縮しやすいことを示しており、量子化とプルーニング後の推論コストが同等でも精度を向上させられる可能性がある。

ABSTRACT

Since hardware resources are limited, the objective of training deep learning models is typically to maximize accuracy subject to the time and memory constraints of training and inference. We study the impact of model size in this setting, focusing on Transformer models for NLP tasks that are limited by compute: self-supervised pretraining and high-resource machine translation. We first show that even though smaller Transformer models execute faster per iteration, wider and deeper models converge in significantly fewer steps. Moreover, this acceleration in convergence typically outpaces the additional computational overhead of using larger models. Therefore, the most compute-efficient training strategy is to counterintuitively train extremely large models but stop after a small number of iterations. This leads to an apparent trade-off between the training efficiency of large Transformer models and the inference efficiency of small Transformer models. However, we show that large models are more robust to compression techniques such as quantization and pruning than small models. Consequently, one can get the best of both worlds: heavily compressed, large models achieve higher accuracy than lightly compressed, small models.

研究の動機と目的

• 固定されたハードウェアと時間制約の下で、Transformerモデルのサイズが計算効率に与える影響を調査する。
• より大きなモデルを短時間訓練した場合が、収束まで訓練した小さなモデルを上回るかどうかを評価する。
• 量子化やプルーニングなどの圧縮技術に対する大規模モデルの頑健性を検討する。
• より大きなモデルが訓練速度と推論効率のトレードオフをより良くするのはいつ、なぜかを分析する。

提案手法

• RoBERTa風の MLM事前学習とWMT14 En→Fr翻訳について、Transformerの幅と深さを体系的に変化させる。
• モデルサイズ間で、ウォールクロック時間と収束までの勾配ステップを測定する。
• 事前学習済みの大規模モデルに対して、後処理の圧縮（量子化とプルーニング）を適用し、精度とメモリの関係を評価する。
• さまざまなビット幅に量子化し、異なるスパースさのレベルまで反復的にプルーニングして下流タスクへの影響を評価する。
• モデルサイズごとの圧縮影響を比較し、大規模モデルの頑健性を判断する。

実験結果

リサーチクエスチョン

• RQ1 MLM事前学習とMTタスク全体で、Transformerサイズを大きくすることは、ウォールクロック時間と勾配ステップの点で収束を加速するか？
• RQ2 下流タスクでファインタuningした際、大きなモデルは小さなモデルよりも量子化とプルーニングで圧縮しやすいか？
• RQ3 固定された推論予算の下で、より大きなモデルの訓練とその後の圧縮の間の精度トレードオフは？
• RQ4 データセットサイズと過学習の可能性は、大規模モデルと小規模モデルの利点にどのように影響するか？
• RQ5 なぜ大規模モデルはより良い圧縮性を示し、実践でどのように活用できるか？

主な発見

• 大規模モデルは、検証誤差をより少ない勾配更新回数で収束させ、オーバーヘッドを考慮した後のウォールクロック時間も速い。
• MLMとMTでは、より大きなモデルがウォールクロック時間あたりの収束を速く提供し、大きさのオーバーヘッドを上回る利得を得られる。
• 大規模モデルは、小規模モデルよりも量子化とプルーニングの両方に対して頑健で、圧縮後に同等のメモリ予算で高い精度を達成する。
• 大規模モデルでのプルーニングと量子化の組み合わせは、パレート最適な精度-効率のトレードオフを生み、特に中程度のスパラリティ付近でのプルーニングと低ビット量子化の組み合わせが強い。
• 最適なモデルサイズはデータセットサイズに関連しており、大規模データ設定では過適合の懸念が少ない場合に大規模モデルが良く機能する。