[論文レビュー] Symmetric block-low-rank layers for fully reversible multilevel neural networks

Bas Peters, Eldad Haber+1 more|arXiv (Cornell University)|Dec 14, 2019

Sparse and Compressive Sensing Techniques参考文献 21被引用数 4

ひとこと要約

本稿では、ビデオなどの大規模3Dデータ向けに、可逆的粗化におけるチャネル増加に起因するメモリ爆発を克服するため、完全に再構成可能なマルチレベルニューラルネットワークを実現する対称的ブロック低ランク（BLR）層を提案する。畳み込みカーネルを対称的BLR構造に分解することにより、パラメータのストレージを著しく削減しながらも、可逆性を維持し、フルビデオ入力におけるエンドツーエンド学習を可能にする。これは従来の可逆ネットワークでは不可能であった。

ABSTRACT

Factors that limit the size of the input and output of a neural network include memory requirements for the network states/activations to compute gradients, as well as memory for the convolutional kernels or other weights. The memory restriction is especially limiting for applications where we want to learn how to map volumetric data to the desired output, such as video-to-video. Recently developed fully reversible neural networks enable gradient computations using storage of the network states for a couple of layers only. While this saves a tremendous amount of memory, it is the convolutional kernels that take up most memory if fully reversible networks contain multiple invertible pooling/coarsening layers. Invertible coarsening operators such as the orthogonal wavelet transform cause the number of channels to grow explosively. We address this issue by combining fully reversible networks with layers that contain the convolutional kernels in a compressed form directly. Specifically, we introduce a layer that has a symmetric block-low-rank structure. In spirit, this layer is similar to bottleneck and squeeze-and-expand structures. We contribute symmetry by construction, and a combination of notation and flattening of tensors allows us to interpret these network structures in linear algebraic fashion as a block-low-rank matrix in factorized form and observe various properties. A video segmentation example shows that we can train a network to segment the entire video in one go, which would not be possible, in terms of memory requirements, using non-reversible networks and previously proposed reversible networks.

研究の動機と目的

• 3Dデータ処理における可逆的粗化の過程で生じるチャネル増加に起因する、完全可逆ニューラルネットワークにおけるメモリ爆発を解消すること。
• 可逆性を維持しつつ、畳み込みカーネルのストレージと計算量を削減する、パラメータ効率の良いレイヤー構造を開発すること。
• 従来の可逆ネットワークではメモリ制限のため不可能であった、フル3Dビデオ入力に対する完全可逆ネットワークのエンドツーエンド学習を可能にすること。
• 訓練中にカーネルランクとメモリ使用量を明示的に制御可能な対称的BLRレイヤーを提案すること。
• 訓練中に段階的にレイヤーのランクを増加させる、学習可能なアダプティブランク戦略を設計し、事前にメモリを増やさずに表現力の向上を実現すること。

提案手法

• 畳み込みカーネル行列を $ \mathbf{K}^\top \mathbf{K} $ に因子分解する対称的ブロック低ランク（BLR）レイヤーを導入し、構成上正定値半定性と対称性を保証する。
• 低ランク構造を持つブロック行列を用いた線形変換としてレイヤーを表現し、$ \mathbf{K}^\top \mathbf{K} $ 因子分解により、効率的なストレージと計算を実現する。
• テンソルの平坦化とブロック行列表記を用い、線形代数および自動微分と互換性のある形で、レイヤーを構造的線形演算子として解釈する。
• 完全可逆マルチレベルネットワークにBLRレイヤーを統合し、粗化およびアップサンプリング段階においても可逆性を維持する。
• 低ランクレイヤーから開始し、訓練中に段階的に畳み込みカーネルを追加するアダプティブランク訓練戦略を実装し、メモリと性能のバランスを取る。
• BLR構造を活用して、非線形活性化関数の下でも、対称性や低ランク構造といった数学的性質を観測・維持する。

実験結果

リサーチクエスチョン

• RQ1 対称的ブロック低ランクレイヤーは、可逆性や表現力に妥協を来さずに、完全可逆マルチレベルネットワークにおける畳み込みカーネルのメモリ容量を効果的に削減できるか？
• RQ2 1x1畳み込み、ブロック巡回カーネル、または squeeze-and-expand ボトルネックといった既存のパラメータ効率設計と比較して、BLRレイヤー構造は、メモリ使用量と学習効率においてどのように異なるか？
• RQ3 BLRレイヤーを備えた完全可逆ネットワークは、複数回の粗化ステップに起因するカーネルメモリ爆発のため、従来の可逆ネットワークでは不可能であったフル3Dビデオ入力に対してエンドツーエンド学習が可能か？
• RQ4 アダプティブランク訓練戦略は、訓練中を通じて低メモリ使用量を維持しつつ、高品質なモデルへの収束を可能にするか？
• RQ5 非線形活性化関数の下で、BLRレイヤーのどの数学的性質が維持され、安定した学習を支援するか？

主な発見

• 対称的BLRレイヤーにより、畳み込みカーネルのメモリ使用量が $ O(n_{\text{chan}}^2) $ から $ O(m n_{\text{chan}}) $ に削減され、$ m $ がブロックランクである。これにより、従来の可逆ネットワークが直面する2次関数的増加を回避する。
• 本手法により、従来の可逆ネットワークではメモリ制限のため不可能であった、フル3Dビデオ入力に対する完全可逆ネットワークの学習が可能になった。
• BLRレイヤー構造により、カスタム行列ベクトル積の実装やコード変更を要せず、既存の可逆ネットワークフレームワークへの直接統合が可能である。
• 数値実験により、アダプティブランク訓練戦略が、高メモリモデルと同等のセグメンテーション性能を達成することが確認された。
• BLRレイヤーの対称的かつ正定値半定性の構造により、安定性が保証され、非線形変換下でのレイヤー特性の理論的分析が可能になった。