[論文レビュー] Only Train Once: A One-Shot Neural Network Training And Pruning Framework

Tianyi Chen, Bo Ji+8 more|arXiv (Cornell University)|Jul 15, 2021

Advanced Neural Network Applications参考文献 81被引用数 43

ひとこと要約

OTOはゼロ不変群を用いて、ワンパスで最初から完全なニューラルネットワークを訓練・圧縮します。微調整なしで剪定するための新規Half-Space Stochastic Projected Gradientを用います。

ABSTRACT

Structured pruning is a commonly used technique in deploying deep neural networks (DNNs) onto resource-constrained devices. However, the existing pruning methods are usually heuristic, task-specified, and require an extra fine-tuning procedure. To overcome these limitations, we propose a framework that compresses DNNs into slimmer architectures with competitive performances and significant FLOPs reductions by Only-Train-Once (OTO). OTO contains two keys: (i) we partition the parameters of DNNs into zero-invariant groups, enabling us to prune zero groups without affecting the output; and (ii) to promote zero groups, we then formulate a structured-sparsity optimization problem and propose a novel optimization algorithm, Half-Space Stochastic Projected Gradient (HSPG), to solve it, which outperforms the standard proximal methods on group sparsity exploration and maintains comparable convergence. To demonstrate the effectiveness of OTO, we train and compress full models simultaneously from scratch without fine-tuning for inference speedup and parameter reduction, and achieve state-of-the-art results on VGG16 for CIFAR10, ResNet50 for CIFAR10 and Bert for SQuAD and competitive result on ResNet50 for ImageNet. The source code is available at https://github.com/tianyic/only_train_once.

研究の動機と目的

• pruning ワークフローにおける多段階の微調整を要せず、効率的なモデル圧縮の必要性を動機づける。
• 0から同時に完全なモデルを訓練し、剪定できるフレームワークを導入する。
• ZIGs（zero-invariant groups）を定義し、出力を保持したまま安全に剪定できるようにする。
• 構造的スパース性を促進する新しい最適化アルゴリズムHalf-Space Stochastic Projected Gradient (HSPG)を開発する。
• CNNsとNLPモデルの標準ベンチマークで、最先端または競争力のある圧縮と精度を示す。

提案手法

• モデルパラメータをzero-invariant groups (ZIGs)に分割し、ゼロのグループを出力に影響を与えずに剪定できるようにする。
• ZIGs上での混合l1/l2ノルムによる構造的スパーシティ正規化を定式化し、グループスパース性を誘導する。
• 初期化段階（SGD）とグループスパース性段階（Half-Space Projection）を持つHalf-Space Stochastic Projected Gradient (HSPG)を用いて非滑らかな正規化問題を解く。
• 降下条件が満たされたとき全グループをゼロにするHalf-Space Projection演算子を導入し、積極的なスパーシティを可能にする。
• 追加の微調整なしにゼロZIGを除去してスリムなアーキテクチャを得ることで、全モデルを剪定する。
• 共通アーキテクチャ（Conv-BN、Residual blocks、Fully ConnectedおよびMulti-Head Attention）に対する自動的または規定されたZIG分割ガイダンスを提供する。

実験結果

リサーチクエスチョン

• RQ1 微調整なしで、ワンラン訓練から圧縮され高性能なモデルを得ることは可能か？
• RQ2 出力を変更せずにゼロ値グループを剪定できるように、パラメータをZIGsとしてどのように分割するか？
• RQ3 HSPGを用いた構造的スパース性アプローチは、標準的なプロキシ法よりもゼログループを促進できるか？
• RQ4 OTOを適用した場合、標準ベンチマークでのFLOPs、パラメータ数、精度の経験的利得はどの程度か？
• RQ5 このアプローチはCNNsおよびBERTのようなNLPモデルにもアーキテクチャに依存せずスケーラブルか？

主な発見

• OTOはベンチマーク全体で競争力のある精度を維持しつつ、FLOPsとパラメータ削減を大幅に実現する。
• CIFAR-10上のVGG16およびVGG16-BNは、Top-1精度を維持しつつFLOPsおよびパラメータ削減を大幅に達成。
• CIFAR-10およびImageNet上のResNet-50は、精度の小さな損失で顕著なFLOPs削減とパラメータ削減を達成。
• SQuAD上のBERTはワンショット剪定フレームワーク下で競争力のある結果を示す。
• HSPGはグループスパース性の促進において標準的なプロキシ法より優れており、収束性の特性を維持する。
• 剪定プロセスは、同じ入力に対してフルモデルと同一の出力を生成する圧縮モデルを微調整なしに得る。