[論文レビュー] StrassenNets: Deep Learning with a Multiplication Budget

Michael Tschannen, Aran Khanna+1 more|arXiv (Cornell University)|Dec 11, 2017

Advanced Neural Network Applications参考文献 40被引用数 11

ひとこと要約

この論文では、深層ニューラルネットワーク内の行列乗算を、和積ネットワーク（SPN）に基づく学習可能で乗算回数が少ない算術回路に置き換えるフレームワーク、StrassenNetsを提案する。これにより、乗算回数に厳密な予算を設ける中でエンドツーエンドの学習が可能になる。共同最適化と知識蒸留を用いて、完全精度モデルの精度を維持したまま乗算回数を99.5%以上削減し、2×2行列乗算を8回ではなく7回の乗算で行うストラッセンのアルゴリズムを再発見する。

ABSTRACT

A large fraction of the arithmetic operations required to evaluate deep neural networks (DNNs) consists of matrix multiplications, in both convolution and fully connected layers. We perform end-to-end learning of low-cost approximations of matrix multiplications in DNN layers by casting matrix multiplications as 2-layer sum-product networks (SPNs) (arithmetic circuits) and learning their (ternary) edge weights from data. The SPNs disentangle multiplication and addition operations and enable us to impose a budget on the number of multiplication operations. Combining our method with knowledge distillation and applying it to image classification DNNs (trained on ImageNet) and language modeling DNNs (using LSTMs), we obtain a first-of-a-kind reduction in number of multiplications (over 99.5%) while maintaining the predictive performance of the full-precision models. Finally, we demonstrate that the proposed framework is able to rediscover Strassen's matrix multiplication algorithm, learning to multiply $2 imes 2$ matrices using only 7 multiplications instead of 8.

研究の動機と目的

• リソース制限のあるハードウェアへの効率的なデプロイを目的とした、深層ニューラルネットワーク（DNN）における乗算回数の削減。
• 乗算回数の予算を課すことにより、算術的複雑さとモデル精度の間の細かい制御を可能にすること。
• ストラッセンやウィノグラードのような手作業で設計されたアルゴリズムに依存せず、データからエンドツーエンドで近似行列乗算変換を学習すること。
• 乗算回数を大幅に削減しても、完全精度モデルの予測性能を維持すること。
• 本フレームワークが、ストラッセンの行列乗算のような既知の高速アルゴリズムを再発見できることを示すこと。

提案手法

• DNN層内の行列乗算を2層構造の和積ネットワーク（SPN）として定式化し、乗算と加算の演算を分離する。
• SPNのエッジ重みを{−1, 0, 1}に属する3値行列として表現することで、非ゼロ重みの数を制御することで明示的に乗算回数を制御可能にする。
• データ駆動型学習を用いて、SPN重みとネットワークパラメータを同時に最適化するエンドツーエンドの確率的勾配降下法を実行する。
• 完全精度の教師モデルから、精度の低下を最小限に抑えた圧縮された学生モデルへの知識蒸留を適用する。
• 畳み込み層および再帰層（ResNetsやLSTMベースの言語モデルを含む）に対しても、畳み込みおよび全結合層をSPNベースの変換に再定式化することでフレームワークを適用する。
• SPNの隠れユニット数（r）をパラメータ化し、乗算予算を制御する。rは出力チャネル数または隠れユニット数に設定される。

実験結果

リサーチクエスチョン

• RQ1 手作業で設計されたアルゴリズム（ストラッセンやウィノグラードなど）と比較して、学習可能で微分可能なフレームワークがDNNにおける乗算回数をより効果的に削減できるか。
• RQ2 SPNベースの変換のエンドツーエンド学習により、乗算回数を99%以上削減した場合でも、モデル精度をどの程度維持できるか。
• RQ3 知識蒸留により、乗算回数を大幅に削減したモデルの予測性能を回復できるか。
• RQ4 本フレームワークは、8回の乗算ではなく7回の乗算で2×2行列乗算を実行する既知の高速アルゴリズム（ストラッセンの2×2行列乗算）を再発見できるか。
• RQ5 乗算回数の削減と精度の観点から、既存の量子化および圧縮技術と比較して、本手法はどのように差別化されるか。

主な発見

• ImageNetにおけるResNet-18では、StrassenNetsが乗算回数を99.63%削減したが、完全精度モデルと比較してトップ1精度が2.0%低下するにとどまった。
• 言語モデリングのタスクでは、PTBデータセットで乗算回数を99.69%削減し、困惑度（perplexity）は3.3%上昇した。
• 知識蒸留を適用したことで、画像分類および言語モデリングの両タスクで、完全精度モデルの予測性能を維持したまま、乗算回数を99.5%以上削減した。
• 蒸留済みのStrassenNetsは、言語モデル設定においてr ≥ n_outの場合、教師モデルを上回る性能を示した。これは、知識蒸留が圧縮モデルの性能を向上させることを示している。
• フレームワークは、8回ではなく7回の乗算で2×2行列乗算を実行するストラッセンのアルゴリズムを成功裏に再発見した。
• 類似した乗算回数の比較的単純なTWN量子化ベースラインと比較して、本手法は7.4%低い困惑度を達成した。これは、性能と効率のトレードオフにおいて優れた効率性を示している。