[论文解读] Revisiting Locally Supervised Learning: an Alternative to End-to-end Training

Yulin Wang, Zanlin Ni+3 more|arXiv (Cornell University)|Jan 26, 2021

Advanced Neural Network Applications参考文献 65被引用 24

一句话总结

该论文提出 InfoPro，一种用于局部监督深度学习的新型损失函数，通过采用非贪婪、信息论目标鼓励信息保留，防止早期层中任务相关特征的崩溃。通过推导 InfoPro 损失的可计算上界，将其表示为重建项与对比项的组合，该方法实现了与端到端训练相当的性能，且 GPU 显存使用率低于 40%，并支持异步、并行训练。

ABSTRACT

Due to the need to store the intermediate activations for back-propagation, end-to-end (E2E) training of deep networks usually suffers from high GPUs memory footprint. This paper aims to address this problem by revisiting the locally supervised learning, where a network is split into gradient-isolated modules and trained with local supervision. We experimentally show that simply training local modules with E2E loss tends to collapse task-relevant information at early layers, and hence hurts the performance of the full model. To avoid this issue, we propose an information propagation (InfoPro) loss, which encourages local modules to preserve as much useful information as possible, while progressively discard task-irrelevant information. As InfoPro loss is difficult to compute in its original form, we derive a feasible upper bound as a surrogate optimization objective, yielding a simple but effective algorithm. In fact, we show that the proposed method boils down to minimizing the combination of a reconstruction loss and a normal cross-entropy/contrastive term. Extensive empirical results on five datasets (i.e., CIFAR, SVHN, STL-10, ImageNet and Cityscapes) validate that InfoPro is capable of achieving competitive performance with less than 40% memory footprint compared to E2E training, while allowing using training data with higher-resolution or larger batch sizes under the same GPU memory constraint. Our method also enables training local modules asynchronously for potential training acceleration. Code is available at: https://github.com/blackfeather-wang/InfoPro-Pytorch.

研究动机与目标

• 解决深度神经网络端到端（E2E）训练中因存储中间激活而导致的高 GPU 显存占用问题。
• 克服因贪婪、短视训练导致的局部监督学习性能下降问题，该问题会使得早期层中有效输入信息发生崩溃。
• 开发一种训练目标，以在本地模块间保留与任务相关的信息，同时丢弃无关（干扰）成分。
• 通过解耦模块间的梯度更新，实现模型并行化和异步训练。
• 提供一种实用、内存高效的端到端训练替代方案，且不损失模型精度。

提出的方法

• 基于信息论原理，提出一种信息传播（InfoPro）损失，以在最大化输入与中间特征间互信息的同时，最小化对干扰变量的依赖。
• 推导 InfoPro 损失的可计算上界，使其在计算上可行，从而得到一个由重建损失与对比损失组合而成的代理目标。
• 使用代理 InfoPro 损失训练本地模块，以促使特征表示保留对下游层有用的信息。
• 将优化过程分解为无需在整个网络反向传播的本地目标，从而实现梯度隔离训练。
• 通过消除对后续层误差信号的依赖，支持本地模块的异步与并行训练。
• 在 PyTorch 中实现该方法，并在包括 CIFAR、SVHN、STL-10、ImageNet 和 Cityscapes 在内的多样化基准上进行验证。

实验结果

研究问题

• RQ1 为何深度神经网络中的贪婪局部监督相较于端到端训练会导致性能下降？
• RQ2 信息论损失函数能否缓解局部监督训练中早期层的信息崩溃问题？
• RQ3 如何有效近似非贪婪、信息保留目标，以适用于实际的深度学习训练？
• RQ4 在保持精度的前提下，使用 InfoPro 损失的局部监督训练能在多大程度上降低 GPU 显存使用？
• RQ5 所提出的方法能否在不损害性能的前提下支持网络模块的异步或并行训练？

主要发现

• InfoPro 训练可防止早期层中任务相关特征的崩溃，这是标准局部监督学习的关键失败模式。
• 该方法在 CIFAR-10、SVHN、STL-10、ImageNet 和 Cityscapes 上的性能与端到端训练相当，且 GPU 显存占用低于 40%。
• 在相同显存约束下，InfoPro 支持 50% 更大的批量大小或 50% 更高的输入分辨率，相比 E2E 训练。
• 由重建项与对比项组成的代理 InfoPro 损失，能有效近似理论上的信息论目标。
• 该方法支持本地模块的异步训练，通过模型并行化实现潜在加速。
• 该方法在多样化视觉基准上表现出鲁棒性与有效性，涵盖分类与语义分割任务。