用于对比自监督学习的可瘦身网络

自监督学习在预训练大型模型方面取得了显著进展，但在小模型上表现不佳。解决该问题主流方案主要依赖知识蒸馏，其涉及两阶段过程：先训练大型教师模型，再蒸馏以提升较小模型的泛化能力。本工作中，我们引入另一种单阶段方案，无需额外教师即可获得预训练小模型，即可瘦身网络用于对比自监督学习(SlimCLR)。可瘦身网络由完整网络与若干权重共享子网络组成，可一次性预训练得到多种网络，包括低计算开销的小网络。然而，权重共享网络间的干扰导致自监督场景下严重性能退化，表现为梯度幅值失衡与梯度方向发散。前者表明一小部分参数在反向传播中产生主导梯度，而主要参数可能未充分优化。后者显示梯度方向无序、优化过程不稳定。为解决这些问题，我们引入三种技术使主要参数产生主导梯度且子网络具有一致输出，包括子网络慢启动训练、在线蒸馏以及按模型尺寸损失重加权。此外，理论结果表明单个可瘦身线性层在线性评估中次优，故线性评估时采用可切换线性探针层。我们以典型对比学习框架实例化SlimCLR，以更少参数与FLOPs取得优于既往方法的性能。代码见https://github.com/mzhaoshuai/SlimCLR。