[論文レビュー] projUNN: efficient method for training deep networks with unitary matrices

Bobak T. Kiani, Randall Balestriero+2 more|arXiv (Cornell University)|Mar 10, 2022

Domain Adaptation and Few-Shot Learning被引用数 11

ひとこと要約

本稿では、射影または輸送を用いてユニタリティを維持するランク-$k$更新を用いて、ユニタリ行列を用いた深層ネットワークの学習を効率化する方法projUNNを提案する。$O(kN^2)$の近似的最適な学習複雑度を達成し、最先端のユニタリRNNを上回るか同等の性能を発揮するとともに、直交畳み込みネットワークの安定的かつスケーラブルな学習を可能にする。

ABSTRACT

In learning with recurrent or very deep feed-forward networks, employing unitary matrices in each layer can be very effective at maintaining long-range stability. However, restricting network parameters to be unitary typically comes at the cost of expensive parameterizations or increased training runtime. We propose instead an efficient method based on rank-$k$ updates -- or their rank-$k$ approximation -- that maintains performance at a nearly optimal training runtime. We introduce two variants of this method, named Direct (projUNN-D) and Tangent (projUNN-T) projected Unitary Neural Networks, that can parameterize full $N$-dimensional unitary or orthogonal matrices with a training runtime scaling as $O(kN^2)$. Our method either projects low-rank gradients onto the closest unitary matrix (projUNN-T) or transports unitary matrices in the direction of the low-rank gradient (projUNN-D). Even in the fastest setting ($k=1$), projUNN is able to train a model's unitary parameters to reach comparable performances against baseline implementations. In recurrent neural network settings, projUNN closely matches or exceeds benchmarked results from prior unitary neural networks. Finally, we preliminarily explore projUNN in training orthogonal convolutional neural networks, which are currently unable to outperform state of the art models but can potentially enhance stability and robustness at large depth.

研究の動機と目的

• ユニタリ行列を用いた深層または再帰的ネットワークの学習における高い計算コストを軽減すること。これは、勾配ステップごとに通常$O(N^3)$の演算を要する。
• 特に大きな$N$に対して、$O(N^3)$の演算を必要としない低ランク更新を用いて、$N\times N$ユニタリ行列の完全なパラメータ化を維持しつつ、学習実行時間を短縮すること。
• SVDや行列逆行列計算の高コストなステップを回避することで、ユニタリティを保ちながらネットワークの安定性とスケーラビリティを向上させること。
• 特に深層アーキテクチャにおいて、ユニタリ／直交行列の適用範囲を畳み込みネットワークへ拡張すること。

提案手法

• projUNNはランク-$k$更新を用いて勾配を近似し、ユニタリ行列の更新に要する計算コストを$O(N^3)$から$O(kN^2)$に削減する。
• Directバージョン（projUNN-D）は、低ランク更新部分空間の固有値分解を用いて、更新された行列を最も近いユニタリ行列に射影する。
• Tangentバージョン（projUNN-T）は、指数写像更新を用いて、ユニタリ行列を低ランク勾配の方向に輸送する。
• 両バージョンとも、全ユニタリ多様体をパラメータ化し、各勾配更新後にユニタリティを強制することで、数値的安定性を確保する。
• 畳み込み層では、FFTを用いてフーリエ領域で直交畳み込みを実行し、ヘルミート対称性を維持することで、効率的な計算を実現する。
• ユニタリ／直交パラメータ用に特別な最適化アルゴリズムを用い、通常のパラメータの1/10〜1/20程度の学習率を設定する。

実験結果

リサーチクエスチョン

• RQ1 低ランク勾配近似を用いることで、性能を損なわずにユニタリ行列の深層ネットワーク学習を効率的に行えるか？
• RQ2 提案手法はユニタリ多様体の完全なパラメータ化を維持しつつ、近似的最適な学習複雑度を達成するか？
• RQ3 長期間のシーケンスタスクにおいて、既存のユニタリRNNと比較して、精度と学習速度の点でどのように差がつくか？
• RQ4 projUNNは直交畳み込みネットワークへ効果的に拡張可能であり、特に深層アーキテクチャにおける安定性の向上に寄与するか？
• RQ5 ハール乱数または単位行列初期化と比較して、どのような初期化手法がprojUNNの収束性と安定性を向上させるか？

主な発見

• projUNNは勾配ステップごとに$O(kN^2)$の学習複雑度を達成し、従来のリー代数ベース手法の$O(N^3)$コストを顕著に低減する。
• $k=1$であっても、benchmarckシーケンス学習タスクにおいて、最先端のユニタリRNNと同等またはそれ以上の性能を発揮する。
• Tangentバージョン（projUNN-T）は、多くの学習ステップにわたり、Directバージョン（projUNN-D）よりも優れた数値的安定性を示す。
• projUNNは直交畳み込みネットワークの効率的学習を可能にし、非常に深いモデルにおける安定性の向上への道筋を提供する。
• 実験的に、$2\times2$歪対称ブロックを用いた初期化（Henaffスタイル）が、単位行列やハール乱数初期化を上回る性能を発揮することが示された。
• $O(N)$ステップごとに間欠的かつ完全な射影を実施することで、ユニタリティを維持しつつ、平均的な$O(kN^2)$実行時間に影響を与えない。