Gradient-based Intra-attention Pruning on Pre-trained Language Models

要約

事前トレーニングされた言語モデルは優れたパフォーマンスを実現しますが、計算コストが高くなります。
枝刈りや知識の蒸留などの技術が、サイズと待ち時間を削減するために開発されてきました。
本研究では、知識の蒸留によりタスク固有の枝刈りを実行し、非常に効果的なモデルを生成する構造化枝刈り手法 GRAIN (Gradient-based Intra-attention pruning) を提案します。
各アテンションヘッドを全体としてプルーニングする一般的なアプローチとは異なり、GRAIN はアテンション内の構造を検査してプルーニングするため、構造検索スペースが大幅に拡張され、より柔軟なモデルが可能になります。
また、2 つのアプローチをより適切に組み合わせるために、枝刈りに対する蒸留の干渉を軽減する勾配分離戦略も提案します。
GLUE、SQuAD、および CoNLL 2003 での実験では、特に高スパース性領域において GRAIN が他の手法よりも顕著に優れており、$93\%\sim99\%$ のパフォーマンスを維持しながら $6\sim7\times$ の高速化を達成することが示されています。
$3\%$ のトランス重量のみが残る極度の圧縮下でも、プルーニングされたモデルは、より大きなモデルと比較して依然として競争力があります。

要約(オリジナル)

Pre-trained language models achieve superior performance but are computationally expensive. Techniques such as pruning and knowledge distillation have been developed to reduce their sizes and latencies. In this work, we propose a structured pruning method GRAIN (Gradient-based Intra-attention pruning), which performs task-specific pruning with knowledge distillation and yields highly effective models. Different from common approaches that prune each attention head as a whole, GRAIN inspects and prunes intra-attention structures, which greatly expands the structure search space and enables more flexible models. We also propose a gradient separation strategy that reduces the interference of distillation on pruning for a better combination of the two approaches. Experiments on GLUE, SQuAD, and CoNLL 2003 show that GRAIN notably outperforms other methods, especially in the high sparsity regime, and achieves $6\sim7\times$ speedups while maintaining $93\%\sim99\%$ performance. Under extreme compression where only $3\%$ transformer weights remain, the pruned model is still competitive compared to larger models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google