[論文レビュー] Fast Differentiable Matrix Square Root

Song Yue, Nicu Sebe+1 more|arXiv (Cornell University)|Jan 21, 2022

Advanced Vision and Imaging被引用数 5

ひとこと要約

本稿では、高速で微分可能な行列平方根の2つの手法—行列テイラー多項式（MTP）と行列パデ近似（MPA）—を提案する。これらは前向き計算に用いられ、逆誤差伝搬には反復的リャプノフ解法を組み合わせる。この手法はSVDやニュートン・シュルツ反復法と比較して顕著な高速化を達成し、視覚タスクにおいても競争力ある性能を発揮する。

ABSTRACT

Computing the matrix square root or its inverse in a differentiable manner is important in a variety of computer vision tasks. Previous methods either adopt the Singular Value Decomposition (SVD) to explicitly factorize the matrix or use the Newton-Schulz iteration (NS iteration) to derive the approximate solution. However, both methods are not computationally efficient enough in either the forward pass or in the backward pass. In this paper, we propose two more efficient variants to compute the differentiable matrix square root. For the forward propagation, one method is to use Matrix Taylor Polynomial (MTP), and the other method is to use Matrix Padé Approximants (MPA). The backward gradient is computed by iteratively solving the continuous-time Lyapunov equation using the matrix sign function. Both methods yield considerable speed-up compared with the SVD or the Newton-Schulz iteration. Experimental results on the de-correlated batch normalization and second-order vision transformer demonstrate that our methods can also achieve competitive and even slightly better performances. The code is available at \href{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}.

研究の動機と目的

• 深層学習フレームワークにおけるSVDおよびニュートン・シュルツ反復法の計算非効率性を解消すること、特にバッチ処理された行列平方根演算において。
• 固有値の差による影響により、逆誤差伝搬における固有値分解の不安定性と高コストを克服すること。
• GPUに適した、微分可能で高速かつ数値的に安定した行列平方根手法を開発すること。
• 視覚応用分野（共分散プーリングや正規化層など）における行列平方根の効率的利用を可能にすること。

提案手法

• 前向きパスにおいて、SVDや反復的手法に代わって行列テイラー多項式（MTP）および行列パデ近似（MPA）を用いて行列平方根を近似する。
• 逆誤差伝搬中に勾配を効率的に計算するため、行列の符号関数に基づく反復的リャプノフ解法を適用する。
• 逆誤差伝搬を連続時間リャプノフ方程式の解法として定式化し、明示的な固有値分解を回避する。
• GPU上で効率的に実行される行列乗算を活用し、最小限の計算オーバーヘッドで深層学習フレームワークに統合する。
• 反復的コンponentの収束安定性を向上させるために、入力行列を事前に正規化する。
• 特に大規模モデルにおいて効率を高めるために、混合精度学習との互換性を活用する。

実験結果

リサーチクエスチョン

• RQ1 MTPおよびMPAは、深層学習における行列平方根計算においてSVDの高速で微分可能な代替手段として機能するか？
• RQ2 提案されたリャプノフに基づく逆誤差伝搬ソルバーは、従来のSVDおよびニュートン・シュルツ反復法と比較して勾配計算の速度で優れているか？
• RQ3 実世界の視覚ベンチマーク（ZCAホワイトニングや2次元視覚変換器など）において、提案手法は速度と精度の両面で優れているか？
• RQ4 大規模モデルにおいて、学習時間とメモリ使用量を削減しながらも、性能を維持または向上できるか？

主な発見

• 我々のMPA-Lya手法は、So-ViT-14モデルにおいてニュートン・シュルツ反復法の3.19倍、SVD-Padéの25.63倍高速である。
• MTP-Lya手法は1回の行列平方根計算あたり2.36msという最速の前向きパスを達成し、SVDおよびNS反復法を上回っている。
• CIFAR-100では、MPA-Lyaは2.52msの推論時間で19.55% ± 0.20の検証誤差を達成し、速度と精度の両面でSVDおよびNS反復法を上回っている。
• 2次元視覚変換器では、MPA-Lyaは254エポックで収束するが、SVDベースの手法より少ない回数であり、トップ1精度は同等またはわずかに上回っている。
• MTP-Lya手法はResNet-18ではSVDおよびNSと同等の性能を達成するが、ResNet-50ではわずかな低下が見られ、これは近似誤差の蓄積による可能性がある。
• リャプノフに基づく逆誤差伝搬ソルバーは、NS反復法と比較して行列乗算の回数を削減し、より高速な逆誤差伝搬を実現している。