[論文レビュー] Momentum Residual Neural Networks

Michael E. Sander, Pierre Ablin+2 more|arXiv (Cornell University)|Feb 15, 2021

Model Reduction and Neural Networks参考文献 48被引用数 12

ひとこと要約

本稿では、勾配逆伝播における活性化の保存に起因するメモリコストを低減するために、Momentum ResNets（モーメンタム・リーダル・ニューラルネットワーク）を提案する。これは、前向き伝搬においてモーメンタム項を導入することで、正確な逆算が可能となり、メモリ使用量を著しく削減する、新しい可逆アーキテクチャである。従来の可逆ネットワークとは異なり、Momentum ResNetsは標準のResNetsの即時置換として使用可能であり、CIFAR-10、CIFAR-100、ImageNetで同等の精度を達成するとともに、既存の可逆モデルが失敗する学習最適化設定でも正常に収束する。

ABSTRACT

The training of deep residual neural networks (ResNets) with backpropagation has a memory cost that increases linearly with respect to the depth of the network. A way to circumvent this issue is to use reversible architectures. In this paper, we propose to change the forward rule of a ResNet by adding a momentum term. The resulting networks, momentum residual neural networks (Momentum ResNets), are invertible. Unlike previous invertible architectures, they can be used as a drop-in replacement for any existing ResNet block. We show that Momentum ResNets can be interpreted in the infinitesimal step size regime as second-order ordinary differential equations (ODEs) and exactly characterize how adding momentum progressively increases the representation capabilities of Momentum ResNets. Our analysis reveals that Momentum ResNets can learn any linear mapping up to a multiplicative factor, while ResNets cannot. In a learning to optimize setting, where convergence to a fixed point is required, we show theoretically and empirically that our method succeeds while existing invertible architectures fail. We show on CIFAR and ImageNet that Momentum ResNets have the same accuracy as ResNets, while having a much smaller memory footprint, and show that pre-trained Momentum ResNets are promising for fine-tuning models.

研究の動機と目的

• バックプロパゲーション中に中間活性化を保存する必要があることによる、深層ResNetsの訓練における高いメモリコストを解消すること。
• 既存のResNetブロックと互換性を持つ可逆残差ネットワークアーキテクチャを開発すること。
• 特にメモリ制限のある環境において、メモリ効率の良い訓練およびファインチューニングを可能にすること。
• モーメンタムを強化した残差ネットワークの理論的表現能力を調査すること。
• 画像分類および学習最適化ベンチマークにおいて、手法の実証的検証を実施すること。

提案手法

• モーメンタム項を用いた速度項を導入することで、ResNetの残差ブロックを変更し、更新則を次のように定義する：$ v_{n+1} = \gamma v_n + (1-\gamma)f(x_n, \theta_n) $, $ x_{n+1} = x_n + v_{n+1} $。
• 中間活性化を正確に復元可能であるため、可逆性が保証され、中間活性化の保存を必要とせずに勾配計算が可能になる。
• Momentum ResNetsの連続極限を、2階常微分方程式（ODE）として分析する。
• 理論的分析により、線形残差関数を用いる場合、Momentum ResNetsはスカラー係数を除き任意の線形写像を表現可能であることが示され、標準のResNetsとは対照的である。
• CIFAR-10、CIFAR-100、ImageNet、および学習最適化のためのLearned-ISTA（LISTA）フレームワークを用いて、性能を実証的に評価する。
• 標準のResNetsをMomentum ResNetsに変換するためのPyTorchパッケージを提供し、最小限の再トレーニングで実現可能である。

実験結果

リサーチクエスチョン

• RQ1 ResNetの前向き伝搬に単純な修正を加えることで、性能を損なわせずにメモリ使用量を削減できる可逆アーキテクチャを構築できるか？
• RQ2 連続極限において、モーメンタムの導入が残差ネットワークの表現能力にどのように影響を与えるか？
• RQ3 既存の可逆ネットワークが失敗する学習最適化タスクにおいて、Momentum ResNetsは収束可能か？
• RQ4 Momentum ResNetsは、CIFAR や ImageNet といった標準的なビジョンベンチマークにどの程度一般化可能か？
• RQ5 事前学習済みのMomentum ResNetsは、最小限の追加トレーニングで効果的にファインチューニング可能か？

主な発見

• Momentum ResNetsは、メモリ使用量を削減しているにもかかわらず、CIFAR-10、CIFAR-100、ImageNetで標準のResNetsと同等のテスト精度を達成した。
• 活性化の正確な逆算が可能であるため、モーメンタムパラメータ$ \gamma $ を調整することで、深層ネットワークのメモリコストをほぼゼロにまで低減できる。
• 学習最適化設定（Learned-ISTA）では、Momentum ResNetsは収束し、可逆性制約により収束に失敗するRevNetベースの変種を上回る性能を示した。
• 理論的分析により、線形残差関数を用いる場合、Momentum ResNetsはスカラー係数を除き任意の線形変換を表現可能であるのに対し、標準のResNetsでは不可能であることが示された。
• 事前学習済みのMomentum ResNetsは、数エポック程度の追加トレーニングで、完全なトレーニングと同等の性能に達することができ、メモリ制限のある環境での効率的適応が可能である。
• フレームワークは、ResNetブロックの即時置換として完全に互換性があり、PyTorchパッケージを用いて簡単に統合可能である。