[論文レビュー] TT-PINN: A Tensor-Compressed Neural PDE Solver for Edge Computing

Ziyue Liu, Xinling Yu+1 more|arXiv (Cornell University)|Jul 4, 2022

Model Reduction and Neural Networks被引用数 4

ひとこと要約

本稿では、テンソル・トレイン（TT）分解を活用して、偏微分方程式（PDE）を解く際のモデルパラメータ数とメモリ使用量を著しく削減しながらも高い精度を維持する、テンソル圧縮型物理学情報ニューラルネットワーク（TT-PINN）を提案する。全結合層内の密な重み行列を低ランクのTTコアに置き換えることで、TT-PINNは最大15倍のパラメータ削減を達成し、標準PINNと比較してはるかに少ないパラメータで優れた性能を発揮する。これにより、リソース制限のあるエッジデバイス上でも効率的な学習が可能になる。

ABSTRACT

Physics-informed neural networks (PINNs) have been increasingly employed due to their capability of modeling complex physics systems. To achieve better expressiveness, increasingly larger network sizes are required in many problems. This has caused challenges when we need to train PINNs on edge devices with limited memory, computing and energy resources. To enable training PINNs on edge devices, this paper proposes an end-to-end compressed PINN based on Tensor-Train decomposition. In solving a Helmholtz equation, our proposed model significantly outperforms the original PINNs with few parameters and achieves satisfactory prediction with up to 15$ imes$ overall parameter reduction.

研究の動機と目的

• 限られたメモリ、計算リソース、エネルギーを備えたエッジデバイス上で大規模な物理学情報ニューラルネットワーク（PINN）を学習するという課題に対処すること。
• モデルサイズとメモリフットプリントを削減することで、エッジプラットフォーム上でのPINNのエンドツーエンド学習を可能にすること。ただし、解の精度を損なわないようにすること。
• PDEを解く際のPINNに、テンソル・トレイン（TT）分解を圧縮手法として用いることの妥当性を検討すること。
• 圧縮されたPINNアーキテクチャが、より大きな非圧縮PINNと同等またはそれ以上の性能を発揮できることを示すこと。

提案手法

• TT-PINNフレームワークは、標準PINNの全結合層内の密な重み行列を、低ランクのテンソル・トレイン（TT）分解に置き換える。これにより、パラメータ数とメモリ使用量が削減される。
• 重み行列は、元の行列の圧縮表現を保持する小さなテンソルであるTTコアの系列に因子分解される。
• 勾配は自動微分を用いて計算され、勾配は直接TTコアを介して逆伝播され、パラメータが更新される。
• TTランクは所望の圧縮比に基づいて決定され、固定ランクおよび適応的ランク決定の両方をサポートする。
• アーキテクチャはシグモイド関数ではなく、Sine活性化関数とXavier初期化を用い、Adam最適化法と学習率の段階的減少を用いて学習される。
• 物理法則（PDE、境界条件、初期条件）は損失関数に直接埋め込まれており、ネットワークの物理学情報性が保持される。

実験結果

リサーチクエスチョン

• RQ1 テンソル・トレイン分解は、エッジデバイス上でのPDEに対するPINNの解の精度を損なわずに、効果的に圧縮可能か？
• RQ2 ヘルムホルツ方程式を解く際、TT圧縮PINNの性能は、同程度またはより大きなサイズの標準PINNと比べてどうか？
• RQ3 予測性能が低下しない範囲で、PINNにおけるパラメータ数とメモリ使用量の圧縮はどの程度達成可能か？
• RQ4 TT-PINNは、計算リソースとメモリが制限されたエッジデバイス上でも、高い解の忠実度を維持しながら学習を可能にするか？

主な発見

• パラメータ数が3,713のTT-PINNは、相対ℓ₂誤差3.03e-2を達成し、パラメータ数50,049の標準PINNと同等の性能を示した。
• 256×256の重み行列を100倍圧縮したTT-PINNは、平均二乗誤差2.25e-4および相対ℓ₂誤差3.03e-2を達成し、12,737パラメータの標準PINNを上回った。
• 256×256行列を40倍圧縮したTT-PINNは、相対ℓ₂誤差7.75e-3を達成し、12,737パラメータの標準PINNよりも顕著に優れた性能を示した。
• TT-PINNは、標準PINNと比較して最大15倍の全体的なパラメータ削減を達成しながら、解の精度を維持または向上させた。
• 可視化比較では、3,713パラメータのTT-PINNが、12,737パラメータの標準PINNよりもより正確な解を生成した。
• 本手法は、低ランクのTT構造が、より大きなPINNの表現力と同等の能力を保持できることを示しており、エッジデバイス上での効率的デプロイメントを可能にする。