[論文レビュー] Stabilizing Equilibrium Models by Jacobian Regularization

Shaojie Bai, Vladlen Koltun+1 more|arXiv (Cornell University)|Jun 28, 2021

Model Reduction and Neural Networks被引用数 6

ひとこと要約

この論文は、深層均衡ネットワーク（DEQ）の安定性を向上させるためにヤコビアン正則化を導入し、層のヤコビアンの固有値半径を明示的に制約することで、前方および後方の収束を改善する。この手法により、DEQはメモリを一定に保ちながら、ResNet-101と同等の速度と精度を達成でき、implicit-depthモデルがarchitectural flexibilityを損なわずに、explicitネットワークと同等の効率性を達成した初めての例である。

ABSTRACT

Deep equilibrium networks (DEQs) are a new class of models that eschews traditional depth in favor of finding the fixed point of a single nonlinear layer. These models have been shown to achieve performance competitive with the state-of-the-art deep networks while using significantly less memory. Yet they are also slower, brittle to architectural choices, and introduce potential instability to the model. In this paper, we propose a regularization scheme for DEQ models that explicitly regularizes the Jacobian of the fixed-point update equations to stabilize the learning of equilibrium models. We show that this regularization adds only minimal computational cost, significantly stabilizes the fixed-point convergence in both forward and backward passes, and scales well to high-dimensional, realistic domains (e.g., WikiText-103 language modeling and ImageNet classification). Using this method, we demonstrate, for the first time, an implicit-depth model that runs with approximately the same speed and level of performance as popular conventional deep networks such as ResNet-101, while still maintaining the constant memory footprint and architectural simplicity of DEQs. Code is available at https://github.com/locuslab/deq .

研究の動機と目的

• 学習および推論中の深層均衡ネットワーク（DEQ）の不安定性と収束の遅さに対処する。
• ハイパーパramータの変化に敏感で、アーキテクチャの変更に脆いDEQの問題を克服する。
• implicit-depthモデルにおける前方（固定点の解法）および後方（勾配計算）のダイナミクスを安定化する。
• DEQがメモリのO(1)を維持しながら、ResNet-101のようなexplicitな深層ネットワークと同等の速度と性能を達成できるようにする。
• 計算的に軽量で、言語モデリングやImageNet分類のような高次元タスクにスケーラブルな正則化手法を開発する。

提案手法

• 層のヤコビアン行列の固有値半径をペナルティ化するヤコビアン正則化項を導入し、収縮的ダイナミクスを促進する。
• トレーニング中にパワー法を用いて固有値半径を推定し、その値が1未満に保たれるようにペナルティを適用することで、安定な固定点収束を保証する。
• アーキテクチャの制約を必要としない微分可能でパラメータに依存しない正則化項として、損失関数に統合する。
• 前方および後方の両方のパスに正則化を適用し、両方のダイナミクスにおける安定性を向上させる。
• 正則化の強さを制御するハイパーパramータγを用い、収束速度とモデル容量のバランスを取る。
• 本手法が、アーキテクチャの変更なしに、WikiText-103 や ImageNet といった大規模タスクにスケーリング可能であることを示す。

実験結果

リサーチクエスチョン

• RQ1 ヤコビアン正則化は、DEQの前方および後方ダイナミクスを安定化させ、高NFE（内部反復回数）に依存しなくてもよいか？
• RQ2 従来の正則化（例：重み減衰）と比較して、ヤコビアン正則化はDEQの学習をどれほど安定化させるか？
• RQ3 ヤコビアン正則化は、DEQをResNet-101のようなexplicitな深層ネットワークと同等の速度にまで加速できるか？
• RQ4 正則化はDEQの一般化性能とアーキテクチャの変更に対するロバストネスを向上させるか？
• RQ5 大規模ベンチマークにおいて、正則化強度γとモデル性能との最適なトレードオフは何か？

主な発見

• ヤコビアン正則化により、収束に必要な平均的な固定点反復回数（NFE）が減少し、正則化なしのDEQに比べて2倍から3倍の高速化が達成された。
• WikiText-103では、正則化付きDEQはわずか5 NFEで92.7%のテスト精度を達成したが、正則化なしバージョンは収束するまでに30 NFEを要した。
• ImageNetでは、正則化付きDEQはResNet-101と同等の精度を維持しながら、O(1)のメモリを保ち、推論速度も同等の水準に達した。
• 本手法により、従来は収束を破壊する原因となっていたアーキテクチャの変更に対してもDEQがロバストになり、脆さが軽減された。
• 重み減衰はDEQの安定化に効果がなく、一方でヤコビアン正則化は学習の安定性と収束性を顕著に改善した。
• 最適な正則化強度γは、モデル性能と収束速度のバランスを取るもので、CIFAR-10ではγ=0.6が最良の結果をもたらした。