[論文レビュー] Top-KAST: Top-K Always Sparse Training

Siddhant M. Jayakumar, Razvan Pascanu+3 more|arXiv (Cornell University)|Jun 7, 2021

Advanced Neural Network Applications被引用数 21

ひとこと要約

Top-KAST は、活性化のための上位-K 絶対値重みと、勾配更新のためのより大きな集合を選択することで、順伝播および逆伝播の両方で一定のスパarsityを維持する完全スパーストレーニング手法である。探索損失を導入することで、劣悪なマスクへの収束を防ぐ。ImageNet および enwik8 で最先端の性能を達成し、密度の高いモデルと比較して 80% のスパarsityまで一致するが、密度の高いパラメータや勾配計算を一切行わず、効率的かつスケーラブルなトレーニングを可能にする。

ABSTRACT

Sparse neural networks are becoming increasingly important as the field seeks to improve the performance of existing models by scaling them up, while simultaneously trying to reduce power consumption and computational footprint. Unfortunately, most existing methods for inducing performant sparse models still entail the instantiation of dense parameters, or dense gradients in the backward-pass, during training. For very large models this requirement can be prohibitive. In this work we propose Top-KAST, a method that preserves constant sparsity throughout training (in both the forward and backward-passes). We demonstrate the efficacy of our approach by showing that it performs comparably to or better than previous works when training models on the established ImageNet benchmark, whilst fully maintaining sparsity. In addition to our ImageNet results, we also demonstrate our approach in the domain of language modeling where the current best performing architectures tend to have tens of billions of parameters and scaling up does not yet seem to have saturated performance. Sparse versions of these architectures can be run with significantly fewer resources, making them more widely accessible and applicable. Furthermore, in addition to being effective, our approach is straightforward and can easily be implemented in a wide range of existing machine learning frameworks with only a few additional lines of code. We therefore hope that our contribution will help enable the broader community to explore the potential held by massive models, without incurring massive computational cost.

研究の動機と目的

• トレーニング中に密度の高いパラメータや勾配計算を回避する完全スパーストレーニング手法を開発すること。
• 大規模で密度の高いモデルのトレーニングに伴う計算コストと環境への影響を、トレーニング全体を通してスパarsityを維持することで軽減すること。
• 一定のスパarsityを維持することで、ハードウェアとエネルギー要件を削減し、大規模モデルへのアクセスを民主化すること。
• 特に自然言語処理分野において、リソースが限られた環境で大規模なアーキテクチャ（例：Transformers）の有効なスパーシフィケーションを可能にすること。
• 既存のディープラーニングライブラリに簡単に統合可能なシンプルでフレームワーク互換性の高い手法を提供すること。

提案手法

• 順伝播では、推論中にスパarsityを維持するために、絶対値で上位-K のパラメータを選択する。
• 逆伝播では、勾配がより大きな集合 B に適用され、B には上位-K の集合 A が含まれるため、より効果的な重み更新が可能になる。
• 動的適応によるスパarsityマスクの促進と、劣悪な重みサブセットへの固定を防ぐために、補助的な探索損失を導入する。
• 完全に密度の高いパラメータや勾配計算を回避することで、順伝播および逆伝播の両方で一定のスパarsityを保証する。
• スパarsityレベルは、ユーザーが定義する上位-K および探索集合のサイズのハイパーパrameterによって制御され、リソース制約に応じた柔軟性を提供する。
• 最小限のコード変更で既存のディープラーニングフレームワークへの統合が容易に設計されている。

実験結果

リサーチクエスチョン

• RQ1 密度の高いパラメータや勾配を一度も生成しない完全スパーストレーニング手法は、密度の高いモデルと同等の性能を維持できるか？
• RQ2 ImageNet や enwik8 といった標準ベンチマークにおいて、Top-KAST は既存のスパース・スパースおよび密度・スパーストレーニング手法と比較してどのように性能を発揮するか？
• RQ3 Top-KAST は、性能劣化を最小限に抑えながら、言語モデリングにおいて大規模なTransformerアーキテクチャを効果的にスパース化できるか？
• RQ4 探索損失の導入により、一般化性能が向上し、劣悪なスパースサブネットワークへの収束が防げるか？
• RQ5 実際の応用において、Top-KAST はモデルサイズをどれほどスケーリング可能にし、計算およびメモリのオーバーヘッドを削減できるか？

主な発見

• Top-KAST は、複数の FLOP バジェットにおいて ImageNet ベンチマークで最先端の性能を達成し、既存のスパース・スパース手法を上回り、RigL と同等の性能を示した。
• enwik8 言語モデリングタスクでは、Top-KAST でトレーニングされた Transformer-XL が 5500万パラメータで 1.00 ビット/文字（bpc）を達成し、2億7700万パラメータの小さな密度モデル（0.99 bpc）を上回った。
• ImageNet では 80% のスパarsityまで高い性能を維持し、劣化は最小限であり、90% スパarsityを超えたあたりからのみ性能が低下し始めた。
• Top-KAST は、密度モデルと 80% スパarsityまで同等の性能を達成するスパースモデルのトレーニングを可能にし、優れた効率性とスケーラビリティを示した。
• このアプローチは、視覚および言語モデリングの両方で有効であり、モデルタイプや分野を問わず広範に適用可能である。
• 既存のフレームワークへの実装が非常に簡単で、追加コードを数行のみで実現でき、広範な採用が可能である。