[论文解读] 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），一种新型可逆架构，通过在ResNet前向传播中引入动量项，实现精确反演并显著降低显存占用。与以往可逆网络不同，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} $。
• 通过允许从$ x_{n+1} $和$ v_{n+1} $精确恢复$ x_n $来保证可逆性，从而在无需存储中间激活值的情况下计算梯度。
• 分析Momentum ResNets在连续极限下作为二阶常微分方程（ODE）的性质。
• 理论分析表明，当残差函数为线性时，Momentum ResNets可表示任意线性映射（至多一个标量因子），而标准ResNets则不能。
• 在CIFAR-10、CIFAR-100、ImageNet以及用于学习优化的Learned-ISTA（LISTA）框架上进行性能的实证评估。
• 提供一个PyTorch工具包，可仅通过少量微调将标准ResNets转换为Momentum ResNets。

实验结果

研究问题

• 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工具包以实现便捷集成。