[論文レビュー] Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations

Steffen Schotthöfer, Emanuele Zangrando+3 more|arXiv (Cornell University)|May 26, 2022

Model Reduction and Neural Networks被引用数 5

ひとこと要約

本論文は、行列微分方程式を用いてニューラルネットワーク内の低ランク重み行列を動的に維持する、新しいトレーニングアルゴリズムを提案する。これにより、メモリと計算量を大幅に削減しつつ、効率的なトレーニングと推論が可能になる。この手法は、フルランクネットワークと同等の性能を達成する一方で、トレーニング中にランクを自動で適応させる。これは、初期化に依存しない『低ランクウィンニングチケット』——つまり、初期化に依存せずにトレーニング中に直接発見可能な高性能な部分ネットワーク——の存在を示している。

ABSTRACT

Neural networks have achieved tremendous success in a large variety of applications. However, their memory footprint and computational demand can render them impractical in application settings with limited hardware or energy resources. In this work, we propose a novel algorithm to find efficient low-rank subnetworks. Remarkably, these subnetworks are determined and adapted already during the training phase and the overall time and memory resources required by both training and evaluating them are significantly reduced. The main idea is to restrict the weight matrices to a low-rank manifold and to update the low-rank factors rather than the full matrix during training. To derive training updates that are restricted to the prescribed manifold, we employ techniques from dynamic model order reduction for matrix differential equations. This allows us to provide approximation, stability, and descent guarantees. Moreover, our method automatically and dynamically adapts the ranks during training to achieve the desired approximation accuracy. The efficiency of the proposed method is demonstrated through a variety of numerical experiments on fully-connected and convolutional networks.

研究の動機と目的

• リソース制約のある環境における大規模ニューラルネットワークのトレーニングおよびデプロイの高コストなメモリと計算量を低減すること。
• 固定ランクを必要とするか、トレーニング後の圧縮を要する従来の低ランクおよびプルーニング手法の限界を克服すること。
• 最適化中に自然に効率的な低ランク部分ネットワーク（低ランクウィンニングチケット）を発見・維持するトレーニング手法を開発すること。
• 特に小さな特異値に近い領域においても、低ランク多様体上での数値的安定性と収束保証を確保すること。
• トレーニング中にランクを自動的かつ動的に適応させることで、ランクの手動ハイパーパramータチューニングの必要性を排除すること。

提案手法

• 重み行列を $ W = U S V^ op $ の形で表現することで、トレーニング中に低ランク要因 $ U $, $ S $, $ V $ を更新することにより、重み行列を低ランク多様体に制限する。
• トレーニングプロセスを、低ランク多様体上での連続時間勾配フローとして定式化し、行列ODEにおける動的低ランク近似（DLRA）の技術を用いる。
• 低ランク構造をトレーニング全体にわたって維持する低ランク数値積分法を採用し、特に小さな特異値に近い場合でも安定性を確保する。
• 逆伝播の勾配は行列連鎖則を用いて計算され、随伴法を用いて $ U $, $ S $, $ V $ の効率的な更新が導出される。
• 畳み込み層に対しては、カーネルテンソルを行列にリシェイプし、フラット化されたカーネルに低ランク分解を適用することで、畳み込み構造を保持する。
• 低ランク多様体上での最適化における近似精度、降下行動、安定性に関する理論的保証を提供する。

実験結果

リサーチクエスチョン

• RQ1 後処理によるプルーニングや初期化に依存する探索を経ずに、トレーニングプロセス中に低ランク部分ネットワークを効率的に発見・トレーニングできるか？
• RQ2 特異値が小さい場合でも、低ランク構造を維持しながら数値的安定性と収束性を確保できるか？
• RQ3 提案手法は、必要な精度を達成しつつ最小限のパラメータ数で、トレーニング中に重み行列のランクを自動的に適応的に変更できるか？
• RQ4 得られる低ランク部分ネットワーク——すなわち『低ランクウィンニングチケット』——は、ランダム初期化からでもフルランクモデルと同等の性能を達成できるか？
• RQ5 標準的な低ランクトレーニングやスパarsityに基づくプルーニング手法と比較して、本手法は効率性と精度の面で優れているか？

主な発見

• 提案手法は、完全ランクネットワークと同等のテスト精度を、全結合および畳み込みアーキテクチャで達成しており、パラメータ数を大幅に削減しても成立する。
• 本手法はトレーニング中に重み行列のランクを動的に適応させ、ランクの手動ハイパーパramータチューニングの必要性を排除する。
• 数値実験の結果、トレーニングおよび推論の両方において、フルランクモデルと比較して顕著なメモリ使用量と計算コストの削減が確認された。
• 近似品質、降下行動、安定性に関する理論的保証が提供されており、特に小さな特異値が存在する場合にも有効である。
• 『低ランクウィンニングチケット』——つまり、初期化に依存せずトレーニング中に直接発見可能な高性能な低ランク部分ネットワーク——の存在が示唆された。
• 適切な多様体制約最適化を用いることで、安定性と精度を維持する点で、単純な低ランクトレーニング手法（例：$ U $ と $ V $ に対する交互SGD）を上回る性能を示した。