トランスフォーマーにおけるロバストなN:Mスパース性トレーニングのためのプログレッシブ・グラディエント・フロー

Progressive Gradient Flow for Robust N:M Sparsity Training in Transformers

February 7, 2024

著者: Abhimanyu Rajeshkumar Bambhaniya, Amir Yazdanbakhsh, Suvinay Subramanian, Sheng-Chun Kao, Shivani Agrawal, Utku Evci, Tushar Krishna

cs.AI

要旨

N:M構造化スパース性は、比較的少ないオーバーヘッドと効率の向上により、大きな関心を集めています。さらに、この形式のスパース性は、その控えめな表現オーバーヘッドにより、メモリフットプリントを削減する点でも非常に魅力的です。N:M構造化スパース性のためのトレーニング手法の開発が進められてきましたが、それらは主に低スパース性領域（約50％以下）に焦点を当てています。しかしながら、これらのアプローチでトレーニングされたモデルの性能は、高スパース性領域（>80％）に直面すると低下する傾向があります。本研究では、既存のスパーストレーニング手法が高スパース性領域でどの程度有効であるかを調査し、これらの手法が低スパース性領域と同等のモデル品質を維持できないことを指摘します。この差異の主な要因として、勾配の大きさに誘発されるノイズのレベルが高いことを示します。この望ましくない効果を緩和するために、プルーニングされた要素への勾配の流れを段階的に制限する減衰メカニズムを採用します。私たちのアプローチは、高スパース性領域において、視覚モデルと言語モデルの品質をそれぞれ最大2％と5％向上させます。また、モデルの精度とトレーニング計算コスト（FLOPs）のトレードオフを評価します。等しいトレーニングFLOPsにおいて、従来のスパーストレーニング手法と比較して、私たちの手法は最大2％の精度向上を示し、より優れた性能を発揮します。ソースコードは以下で公開されています： https://github.com/abhibambhaniya/progressive_gradient_flow_nm_sparsity

English

N:M Structured sparsity has garnered significant interest as a result of relatively modest overhead and improved efficiency. Additionally, this form of sparsity holds considerable appeal for reducing the memory footprint owing to their modest representation overhead. There have been efforts to develop training recipes for N:M structured sparsity, they primarily focus on low-sparsity regions (sim50\%). Nonetheless, performance of models trained using these approaches tends to decline when confronted with high-sparsity regions (>80\%). In this work, we study the effectiveness of existing sparse training recipes at high-sparsity regions and argue that these methods fail to sustain the model quality on par with low-sparsity regions. We demonstrate that the significant factor contributing to this disparity is the presence of elevated levels of induced noise in the gradient magnitudes. To mitigate this undesirable effect, we employ decay mechanisms to progressively restrict the flow of gradients towards pruned elements. Our approach improves the model quality by up to 2% and 5% in vision and language models at high sparsity regime, respectively. We also evaluate the trade-off between model accuracy and training compute cost in terms of FLOPs. At iso-training FLOPs, our method yields better performance compared to conventional sparse training recipes, exhibiting an accuracy improvement of up to 2%. The source code is available at https://github.com/abhibambhaniya/progressive_gradient_flow_nm_sparsity.