[論文レビュー] Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed) Neural Networks

Yequan Zhao, Xinling Yu+4 more|arXiv (Cornell University)|Aug 18, 2023

Model Reduction and Neural Networks被引用数 6

ひとこと要約

この論文は、テンソル圧縮されたゼロ次順序（ZO）最適化を用いて、バックプロパゲーションフリーなニューラルネットワークおよび物理情報付きニューラルネットワーク（PINNs）の学習フレームワークを提案する。分散低減のためのテンソル圧縮とハイブリッド ZO 勾配推定戦略を組み合わせることで、自動微分を必要とせず、メモリ効率的かつスケーラブルにエッジデバイス上で学習可能であり、MNIST では準SOTAの精度を達成し、20次元のハミルトニアン＝ジャコビ＝ベルマン方程式を効果的に解くことに成功した。

ABSTRACT

Backward propagation (BP) is widely used to compute the gradients in neural network training. However, it is hard to implement BP on edge devices due to the lack of hardware and software resources to support automatic differentiation. This has tremendously increased the design complexity and time-to-market of on-device training accelerators. This paper presents a completely BP-free framework that only requires forward propagation to train realistic neural networks. Our technical contributions are three-fold. Firstly, we present a tensor-compressed variance reduction approach to greatly improve the scalability of zeroth-order (ZO) optimization, making it feasible to handle a network size that is beyond the capability of previous ZO approaches. Secondly, we present a hybrid gradient evaluation approach to improve the efficiency of ZO training. Finally, we extend our BP-free training framework to physics-informed neural networks (PINNs) by proposing a sparse-grid approach to estimate the derivatives in the loss function without using BP. Our BP-free training only loses little accuracy on the MNIST dataset compared with standard first-order training. We also demonstrate successful results in training a PINN for solving a 20-dim Hamiltonian-Jacobi-Bellman PDE. This memory-efficient and BP-free approach may serve as a foundation for the near-future on-device training on many resource-constraint platforms (e.g., FPGA, ASIC, micro-controllers, and photonic chips).

研究の動機と目的

• FPGA や ASIC、マイコンなどのエッジプラットフォームにおけるハードウェア制限およびメモリ制限により、オンデバイスでのニューラルネットワーク学習におけるバックプロパゲーションの必要性を排除すること。
• 自動微分や勾配計算を一切行わず、実用的なニューラルネットワークおよび PINNs のエンドツーエンド学習を可能にすること。
• 高次元設定において従来の ZO 最適化の高い分散とスケーラビリティの欠如を、テンソル圧縮とハイブリッド勾配推定により克服すること。
• PINNs におけるバックプロパゲーションフリーな学習を実現するために、PDE 損失関数の導出計算をスパースグリッド・シュタイン推定器に置き換えること。

提案手法

• 高次元パラメータ勾配を低次元のテンソル部分空間に射影することで分散を著しく低減する、テンソル圧縮による分散低減技術を導入する。
• ランダムな摂動と有限差分近似を組み合わせることで、前方パスの回数を最小限に抑えつつ勾配の正確さを維持する、ハイブリッド ZO 勾配評価手法を提案する。
• PINN の損失関数における導関数を、バックプロパゲーションを用いずに計算するためのスパースグリッドに基づくシュタイン推定器を開発する。これにより、バックプロパゲーションフリーな PINN 学習が可能になる。
• すべての学習ステップにおいて前方評価のみを採用することで、中間活性化の保存やチェーンルールに基づく勾配計算の必要性を排除する。
• 前方パスのみを用いて損失評価とパラメータ更新が行えるため、標準的なフィードフォワードネットワークおよび PINNs にこのフレームワークを適用する。
• 二段階の学習プロセスを採用する：最初に符号に基づく ZO 勾配推定（Zo-SignRGE）を用いた粗い学習段階を経て、その後に ZO-CGE を用いた微調整段階で収束性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 ゼロ次順序最適化は、バックプロパゲーションを一切用いずに、実用的サイズのニューラルネットワークをスクラッチからスケーラブルに学習可能か？
• RQ2 大規模モデルに対するバックプロパゲーションフリーな学習を可能にするために、ZO 勾配の分散を十分に低減できるか？
• RQ3 ハイブリッド ZO 勾配推定戦略により、学習に必要な前方パスの回数を削減できるか？
• RQ4 PDE の導出計算をスパースグリッド推定器に置き換えることで、バックプロパゲーションを用いずに PINNs を学習可能か？
• RQ5 エッジ関連のタスクにおいて、バックプロパゲーションフリーな学習の性能は、一次順序法と比較して精度および収束性においてどうなるか？

主な発見

• テンソル圧縮された ZO アプローチは、スパarsityに基づく ZO 学習と比較して、MNIST における検証損失を顕著に低減し、FC-3 ネットワークでは TT-3 パラメータ化で 4.25e-05 の損失に到達した。
• MNIST データセットにおいて、バックプロパゲーションフリーな学習は、3層の全結合ネットワークを用いて 3.58e-04 の検証損失を達成し、自動微分を用いた一次順序法と同等の性能を示した。
• 本手法は、20次元のハミルトニアン＝ジャコビ＝ベルマン方程式を、自動微分を用いた一次順序学習とほぼ同等の精度で解く PINN の学習に成功した。
• スパースグリッド・シュタイン推定器により、バックプロパゲーションを用いずに損失関数内の導関数を正確に近似可能となり、バックプロパゲーションフリーな PINN 学習が実現された。
• ハイブリッド ZO 学習戦略により、前方パスの回数を削減しながらも勾配の品質を維持でき、リソース制限のあるプラットフォームでの効率性が向上した。
• 微調整段階でのクエリコストが高めではあるが、全体としてエッジデバイス上でエンドツーエンドのバックプロパゲーションフリーな学習が可能となり、オンデバイス AI 学習の道筋が開けた。