[論文レビュー] Towards Compact Neural Networks via End-to-End Training: A Bayesian Tensor Approach with Automatic Rank Determination

Cole Hawkins, Xing Liu+1 more|arXiv (Cornell University)|Oct 17, 2020

Tensor decomposition and applications被引用数 4

ひとこと要約

本稿では、学習中に最適なテンソルランクを自動的に決定するエンドツーエンドのベイジアンテンソル化学習フレームワークを提案する。これにより、スクラッチから超コンパactなニューラルネットワークが構築可能になる。変分推論と低ランクテンソル分解（CP、Tucker、TT）を組み合わせることで、モデルパラメータを最大26,000倍まで削減（例：4.25Bから1.6×10⁵まで）しつつ、精度損失を最小限に抑え、エッジデバイスでのエネルギー効率の高い学習を可能にする。

ABSTRACT

While post-training model compression can greatly reduce the inference cost of a deep neural network, uncompressed training still consumes a huge amount of hardware resources, run-time and energy. It is highly desirable to directly train a compact neural network from scratch with low memory and low computational cost. Low-rank tensor decomposition is one of the most effective approaches to reduce the memory and computing requirements of large-size neural networks. However, directly training a low-rank tensorized neural network is a very challenging task because it is hard to determine a proper tensor rank {\it a priori}, which controls the model complexity and compression ratio in the training process. This paper presents a novel end-to-end framework for low-rank tensorized training of neural networks. We first develop a flexible Bayesian model that can handle various low-rank tensor formats (e.g., CP, Tucker, tensor train and tensor-train matrix) that compress neural network parameters in training. This model can automatically determine the tensor ranks inside a nonlinear forward model, which is beyond the capability of existing Bayesian tensor methods. We further develop a scalable stochastic variational inference solver to estimate the posterior density of large-scale problems in training. Our work provides the first general-purpose rank-adaptive framework for end-to-end tensorized training. Our numerical results on various neural network architectures show orders-of-magnitude parameter reduction and little accuracy loss (or even better accuracy) in the training process. Specifically, on a very large deep learning recommendation system with over $4.2 imes 10^9$ model parameters, our method can reduce the variables to only $1.6 imes 10^5$ automatically in the training process (i.e., by $2.6 imes 10^4$ times) while achieving almost the same accuracy.

研究の動機と目的

• 大規模で過剰パラメータ化された深層ニューラルネットワークの学習にかかる高い計算コストとエネルギー消費を低減すること。
• 低ランク因子分解における最適なテンソルランクを手動で選択するという課題を克服すること。これは、モデルの精度と効率性において極めて重要である。
• フルプレシジョンの事前学習を排除するため、スクラッチからコンパクトなニューラルネットワークを直接学習可能にする。
• 複数のテンソル形式（CP、Tucker、テンソルトレイン）をサポートするスケーラブルでエンドツーエンドのフレームワークを構築し、自動ランク適応を実現すること。
• 大規模モデルにおいて、高い圧縮比と低い推論コストを実現しながら、モデルの精度を維持または向上させること。

提案手法

• 低ランクテンソル分解（CP、Tucker、テンソル・トレイン、TTM）の上にベイジアン確率的モデルを定式化し、学習中に最適なランクを自動的に決定可能にする。
• テンソル要因とランクの事後分布を効率的に最適化するためのカスタマイズされた確率的変分推論アルゴリズムを開発する。
• スパarsity誘導型ランク適応を可能にするために、自動関連性決定（ARD）を組み込んだ階層的事前分布構造を導入する。
• 高次元かつ構造的パrameter空間における事後分布近似を改善するために、射影型カーネル化されたスティン距離を用いる。
• 完全な埋め込みテーブルを明示的に構築しないで済む微分可能学習パイプラインを設計し、大規模モデルにおけるメモリオーバーヘッドを低減する。
• 大規模なニューラルネットワーク推薦モデル（DLRM）のエンドツーエンド学習と統合することで、数十亿パラメータのモデルでもスケーラビリティを実証する。

実験結果

リサーチクエスチョン

• RQ1 事前ハイパーパramータチューニングなしに、エンドツーエンド学習中にベイジアンテンソル化学習フレームワークが最適なテンソルランクを自動的に決定できるか？
• RQ2 自動ランク決定は、大規模なディープラーニングモデルにおいて、モデルパラメータをどれほど削減しつつ、予測精度を維持または向上させられるか？
• RQ3 エンドツーエンドテンソル化学習は、学習後圧縮と比較して、メモリ使用量、エネルギー効率、推論効率の面でどの程度優れているか？
• RQ4 提案手法により、FPGAのようなリソース制限のあるプラットフォームでも、実用的なオンデバイス学習が可能になるか？
• RQ5 実世界のベンチマークにおいて、CP、Tucker、TTなどの異なるテンソル形式が、圧縮比、精度、学習効率に与える影響は何か？

主な発見

• 提案手法により、大規模なDLRMにおいてモデルパラメータを4.25Bから1.6×10⁵まで削減し、26,000倍の圧縮比を達成した。精度損失は最小限であった。
• DLRMベンチマークにおいて、ARD-LUバージョンはわずか227万パラメータで78.67%の精度を達成し、パラメータ数が少ないにもかかわらず、固定ランクベースラインを上回った。
• FPGA上で、非テンソル化学習の組み込みCPUと比較して、最大123倍のエネルギー効率向上と59倍の高速化を達成した。
• ARDによる自動ランク決定により、固定ランク手法よりも優れた一般化性能が得られ、ARD-HCはTT形式でたった163,976の最終パラメータで78.73%の精度を達成した。
• フレームワークは、完全な埋め込みテーブルを明示的に構築せず、学習後圧縮の計算コストを回避する形で、超コンパクトなモデルを正常に学習した。
• ベイジアンソルバーにより、自動ランク適応を伴う低ランクテンソル化により、数十億パラメータのモデルの安定的かつスケーラブルな学習が可能になった。