[論文レビュー] Learning Best Combination for Efficient N:M Sparsity

Yuxin Zhang, Mingbao Lin+6 more|arXiv (Cornell University)|Jun 14, 2022

Machine Learning and ELM被引用数 23

ひとこと要約

本研究はN:Mのスパーシティを組み合わせ選択問題として再定義し、Learning Best Combination (LBC) を提案する。LBCは各候補ウェイトサブセットのスコアを学習し、標準的な訓練中に低スコアの組み合わせを徐々に除去して、強い精度を保ちながら効率的なN:Mスパース性を実現する。

ABSTRACT

By forcing at most N out of M consecutive weights to be non-zero, the recent N:M network sparsity has received increasing attention for its two attractive advantages: 1) Promising performance at a high sparsity. 2) Significant speedups on NVIDIA A100 GPUs. Recent studies require an expensive pre-training phase or a heavy dense-gradient computation. In this paper, we show that the N:M learning can be naturally characterized as a combinatorial problem which searches for the best combination candidate within a finite collection. Motivated by this characteristic, we solve N:M sparsity in an efficient divide-and-conquer manner. First, we divide the weight vector into $C_{ ext{M}}^{ ext{N}}$ combination subsets of a fixed size N. Then, we conquer the combinatorial problem by assigning each combination a learnable score that is jointly optimized with its associate weights. We prove that the introduced scoring mechanism can well model the relative importance between combination subsets. And by gradually removing low-scored subsets, N:M fine-grained sparsity can be efficiently optimized during the normal training phase. Comprehensive experiments demonstrate that our learning best combination (LBC) performs consistently better than off-the-shelf N:M sparsity methods across various networks. Our project is released at \url{https://github.com/zyxxmu/LBC}.

研究の動機と目的

• GPU上でのハードウェア加速ポテンシャルを保持する効率的なN:M sparsityを動機づける。
• N:M sparsityを有限の組合せ問題としてモデル化し、その解法を分割統治的に行う。
• 最適なN/ Mのウェイト組み合わせを識別する学習可能なスコアリング機構を導入。
• プルーニング過程のための密な勾配計算を必要とせず、エンドツーエンドの訓練を可能にする。
• LBCがCNNsとViTモデル全般で既存のN:M sparsity法を上回ることを示す。

提案手法

• 各ウェイトグループ W^l_g,:（サイズ M）をすべての C(M,N) 個のNサイズの組み合わせに分割する。
• 各組み合わせ Θ^l_g,i に学習可能なスコア S^l_g,i を割り当て、ネットワークウェイトと jointly 最適化する。
• 訓練エポックを通じて低スコアの組み合わせを徐々に剪定するスケジュールを用い、最終的に最良の組み合わせを残す。
• 正向計算 Z^l = Z^{l-1} ⊗ (B^l ⊙ W^l) を用いて N:M スパース性を実現するために、W^l に二値マスク B^l を適用する。
• 各組み合わせの重要度を反映するよう、直通近似法 (STE) 勾配で S^l_g,i を更新する。
• 学習されたスコアがウェイト削除による損失変化を近似し、剪定を導くことを示す解析的正当化を提供。
• 密な勾配計算を回避しつつ高いスパース性性能を達成することで、エンドツーエンドの訓練効率を実証する。

実験結果

リサーチクエスチョン

• RQ1 重い pre-training や密な勾配計算なしに、N:M sparsityをゼロ-shotで効果的に学習できるか？
• RQ2 学習可能な組み合わせスコアは、N:M sparsityのウェイトサブセットの相対的重要度を正確に反映できるか？
• RQ3 徐々に剪定する分割統治法は、既存のN:M手法より精度と効率が向上するか？
• RQ4 標準ベンチマーク（ImageNet、COCO）上で、CNNとVision TransformerにおけるLBCの性能は、最先端のスパース性手法と比較してどうか？
• RQ5 訓練スケジュール（t_i, t_f）が訓練コストと精度のトレードオフに与える影響は何か？

主な発見

• LBCはResNet-50とDeiTで2:4, 2:8, 1:16のスパース性パターンで競争力ある、あるいは優れたTop-1精度を達成し、いくつかのベースラインより低い訓練コスト。
• ImageNetのResNet-50では、2:4スパース性で120訓練エポック、Top-1 77.2%を達成し、ASPとSR-STEを上回る。
• LBCはResNet-50で1:16スパース性でTop-1 71.8%を達成し、同程度のスパース性でSTRを精度と見かけの効率の双方で上回る。
• 物体検出およびインスタンスセグメンテーション（COCO）では、2:4および2:8パターンでSR-STEより一貫してmAP/boxおよびマスク結果を改善。
• DeiT-smallでは2:4スパース性でTop-1 80.1%を達成し、SR-STEより0.5ポイント高く、訓練FLOPsを抑える。
• 高いスパース性での非構造スパース性手法と比較して、LBCはNVIDIA Ampereハードウェア上でより良い実用的圧縮と加速を提供し、精度を維持または向上させる。