重访局部监督学习：端到端训练的替代方案
计算机视觉与模式识别
2021-01-27 v1 机器学习
机器学习
摘要
由于需要存储中间激活以进行反向传播，深度网络的端到端（E2E）训练通常面临高 GPU 内存占用的问题。本文旨在通过重访局部监督学习来解决该问题，其中网络被拆分为梯度隔离的模块并以局部监督进行训练。我们通过实验表明，仅用 E2E 损失训练局部模块往往会使早期层中与任务相关的信息坍缩，从而损害完整模型的性能。为避免此问题，我们提出信息传播（InfoPro）损失，其鼓励局部模块尽可能保留有用信息，同时逐步丢弃与任务无关的信息。由于 InfoPro 损失以其原始形式难以计算，我们推导出一个可行的上界作为替代优化目标，从而得到一种简单而有效的算法。事实上，我们表明所提方法可归结为最小化重构损失与常规交叉熵/对比项之组合。在五个数据集（即 CIFAR、SVHN、STL-10、ImageNet 和 Cityscapes）上的大量实证结果验证了 InfoPro 能够以低于 E2E 训练 40% 的内存占用实现有竞争力的性能，同时允许在相同 GPU 内存约束下使用更高分辨率或更大批量的训练数据。我们的方法还支持异步训练局部模块以潜在加速训练。代码见：https://github.com/blackfeather-wang/InfoPro-Pytorch。
引用
@article{arxiv.2101.10832,
title = {Revisiting Locally Supervised Learning: an Alternative to End-to-end Training},
author = {Yulin Wang and Zanlin Ni and Shiji Song and Le Yang and Gao Huang},
journal= {arXiv preprint arXiv:2101.10832},
year = {2021}
}
备注
Accepted by ICLR 2021