[论文解读] Stabilizing Equilibrium Models by Jacobian Regularization

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

Model Reduction and Neural Networks被引用 6

一句话总结

本文引入雅可比正则化以稳定深度平衡网络（DEQs），显式约束层雅可比矩阵的谱半径，从而提升前向与反向传播的收敛性。该方法使DEQs在保持恒定内存消耗的同时，实现与ResNet-101相当的速度与精度，首次使隐式深度模型在不牺牲架构灵活性的前提下，达到显式网络的效率水平。

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 .

研究动机与目标

• 解决深度平衡网络（DEQs）在训练与推理过程中存在的不稳定性与收敛缓慢问题。
• 克服DEQs对架构变更的脆弱性，以及在不同超参数设置下泛化能力差的问题。
• 稳定隐式深度模型中的前向（固定点求解）与反向（梯度计算）动力学。
• 使DEQs在保持O(1)内存消耗的同时，实现与ResNet-101等显式深度网络相当的速度与性能。
• 开发一种计算开销轻量、可扩展至高维任务（如语言建模与ImageNet分类）的正则化方法。

提出的方法

• 引入雅可比正则化项，惩罚层雅可比矩阵的谱半径，以促进压缩性动力学。
• 使用幂法在训练过程中估计谱半径，并施加惩罚使其低于1，确保稳定固定点收敛。
• 将正则化项整合至损失函数中，作为可微、与参数无关的正则化器，无需架构约束。
• 在前向与反向传播中均应用正则化，提升两种动力学的稳定性。
• 使用超参数γ控制正则化强度，平衡收敛速度与模型容量。
• 证明该方法可扩展至大规模任务（如WikiText-103与ImageNet），且无需架构修改。

实验结果

研究问题

• RQ1 雅可比正则化能否稳定DEQs的前向与反向动力学，从而减少对高NFE的依赖？
• RQ2 与传统正则化（如权重衰减）相比，雅可比正则化在稳定DEQ训练方面表现如何？
• RQ3 雅可比正则化在多大程度上可加速DEQs，使其达到ResNet-101等显式深度网络的速度？
• RQ4 该正则化是否提升了DEQs的泛化能力与对架构变更的鲁棒性？
• RQ5 在大规模基准测试中，正则化强度γ与模型性能之间的最优权衡是什么？

主要发现

• 雅可比正则化降低了达到收敛所需的平均固定点迭代次数（NFE），使正则化DEQs相比未正则化版本提速2至3倍。
• 在WikiText-103上，正则化DEQ仅用5次NFE即达到92.7%的测试准确率，而未正则化版本需30次NFE才能收敛。
• 在ImageNet上，正则化DEQ在保持O(1)内存的同时，达到与ResNet-101相当的准确率，并实现相近的推理速度。
• 该方法使DEQs对以往会导致收敛失败的架构修改具备鲁棒性，显著降低脆弱性。
• 权重衰减在稳定DEQs方面效果甚微，而雅可比正则化显著提升了训练稳定性与收敛性。
• 最优正则化强度γ在模型性能与收敛速度之间实现平衡，CIFAR-10上γ=0.6时表现最佳，且未过度约束模型。