[論文レビュー] Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models

Tri Dao, Beidi Chen+5 more|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 90被引用数 7

ひとこと要約

Pixelated Butterfly は、フラットなブロックバタフライ行列と低ランク成分を組み合わせることで、ニューラルネットワークの層をスパース化する、シンプルでハードウェアに優れたスパーストレーニング手法を提案する。この手法は、ビジョントランスフォーマー、GPT-2、MLP-Mixer などの密度型モデルと比較して、ImageNet や WikiText-103 で最大 2.5 倍の高速化を達成し、精度の低下なしに、ブロックアラインドで並列化可能な計算を可能にすることで、密度型モデルに非常に近い近似性能を維持する。

ABSTRACT

Overparameterized neural networks generalize well but are expensive to train. Ideally, one would like to reduce their computational cost while retaining their generalization benefits. Sparse model training is a simple and promising approach to achieve this, but there remain challenges as existing methods struggle with accuracy loss, slow training runtime, or difficulty in sparsifying all model components. The core problem is that searching for a sparsity mask over a discrete set of sparse matrices is difficult and expensive. To address this, our main insight is to optimize over a continuous superset of sparse matrices with a fixed structure known as products of butterfly matrices. As butterfly matrices are not hardware efficient, we propose simple variants of butterfly (block and flat) to take advantage of modern hardware. Our method (Pixelated Butterfly) uses a simple fixed sparsity pattern based on flat block butterfly and low-rank matrices to sparsify most network layers (e.g., attention, MLP). We empirically validate that Pixelated Butterfly is 3x faster than butterfly and speeds up training to achieve favorable accuracy--efficiency tradeoffs. On the ImageNet classification and WikiText-103 language modeling tasks, our sparse models train up to 2.5x faster than the dense MLP-Mixer, Vision Transformer, and GPT-2 medium with no drop in accuracy.

研究の動機と目的

• 動的スパースリティマスクや非構造的スパースリティに依存する既存のスパーストレーニング手法の非効率さと高いオーバーヘッドを是正すること。
• スパース行列演算におけるハードウェア非効率性を解消し、スパースリティパターンをブロック最適化されたメモリアクセスに一致させること。
• 特にアテンション層や MLP コンポonent に適した、多様なニューラルネットワーク層に適用可能な統一的で静的スパースリティパターンの開発。
• 特に過パラメータ化されたモデルにおいて、精度を損なわずに高いトレーニングスピードアップを達成すること。
• 実用的で広範囲への導入を可能にするために、実装を簡素化し、計算オーバーヘッドを低減すること。

提案手法

• スパース行列のスーパーセット上で連続的最適化フレームワークを提案し、スパースリティマスク探索を簡素化するため、バタフライ行列の積を用いる。
• 標準バタフライ行列の一次近似として、フラットバタフライ行列を導入し、逐次的積を和に置き換えることで、より良い並列化を実現する。
• GPU でのメモリアクセス効率を向上させるために、固定でハードウェアにアラインされたスパースリティパターンを持つブロックバタフライ行列を設計する。
• フラットブロックバタフライ行列と低ランク成分を組み合わせることで、グローバルおよびローカルな特徴表現を捉える。
• トレーニング全体を通して静的スパースリティパターンを用いることで、動的マスクの更新が必要なくなり、トレーニングのオーバーヘッドが削減される。
• この手法をトランスフォーマーおよび MLP-Mixer モデルのアテンション層と MLP 層のスパース化に適用し、効率的でスケーラブルなトレーニングを実現する。

実験結果

リサーチクエスチョン

• RQ1 静的でハードウェアにアラインされたスパースリティパターンは、スパースニューラルネットワークにおいて、高いトレーニング効率と強力なモデル精度を両立できるか？
• RQ2 異なるニューラルネットワークアーキテクチャにおいて、フラットブロックバタフライ行列の近似誤差は、密度型モデルと比べてどの程度か？
• RQ3 ブロックアラインドスパースリティパターンを用いたスパーストレーニングは、精度の低下なしに、ウォールクロックスピードで密度型トレーニングをどの程度上回れるか？
• RQ4 統一されたスパースリティパターンは、アテンション層や MLP 層など多様なニューラルネットワークコンポーネントに効果的に一般化できるか？
• RQ5 バタフライ構造と低ランク成分の組み合わせは、非構造的または標準バタフライスパースリティと比較して、より優れた近似性能と効率性を達成できるか？

主な発見

• Pixelated Butterfly は、ImageNet や WikiText-103 において、ビジョントランスフォーマー、GPT-2 メディアム、MLP-Mixer といった密度型モデルと比較して、最大 2.5 倍の高速化を達成し、精度の低下なしに実現した。
• この手法は、ハードウェアの効率的利用と並列化の向上のおかげで、標準バタフライ行列トレーニングよりも 3 倍速い。
• フラットバタフライ行列は、逐次的積を和に置き換えることで、効率的で並列化可能な行列乗算を可能にした。
• ブロックバタフライ行列はハードウェアのメモリブロックに一致しており、メモリアクセス効率を顕著に向上させ、遅延を低減した。
• フラットブロックバタフライ行列と低ランク成分の組み合わせは、近似精度とトレーニングスピードの両面で、他のスパースリティパターンを常に上回った。
• このアプローチは、過パラメータ化されたモデルにおいても強力な一般化性能を維持し、密度型ネットワークで観察されるダブルデセントの挙動を支持した。