[论文解读] Momentum^2 Teacher: Momentum Teacher with Momentum Statistics for Self-Supervised Learning

Zeming Li, Songtao Liu+1 more|arXiv (Cornell University)|Jan 19, 2021

Music and Audio Processing参考文献 53被引用 11

一句话总结

该论文提出 Momentum² Teacher，一种自监督学习方法，通过不仅对学生网络权重，还对批归一化（BN）统计量应用动量更新，从而提升教师网络的稳定性。通过引入‘动量 BN’——一种基于历史 BN 统计量的动量更新机制，该方法在仅使用小批量训练（128）的情况下，实现了 ImageNet 上 74.5% 的 SOTA（最先进）线性评估 top-1 准确率，且无需 TPU 或同步 BN。

ABSTRACT

In this paper, we present a novel approach, Momentum$^2$ Teacher, for student-teacher based self-supervised learning. The approach performs momentum update on both network weights and batch normalization (BN) statistics. The teacher's weight is a momentum update of the student, and the teacher's BN statistics is a momentum update of those in history. The Momentum$^2$ Teacher is simple and efficient. It can achieve the state of the art results (74.5\%) under ImageNet linear evaluation protocol using small-batch size(\eg, 128), without requiring large-batch training on special hardware like TPU or inefficient across GPU operation (\eg, shuffling BN, synced BN). Our implementation and pre-trained models will be given on GitHub\footnote{https://github.com/zengarden/momentum2-teacher}.

研究动机与目标

• 解决小批量自监督训练中批归一化不稳定的問題，该問題限制了学生-教师框架的性能。
• 克服现有 SOTA 方法（如 BYOL）对大批次训练或专用硬件（如 TPU）的依赖，以及同步 BN 操作效率低下的问题。
• 通过将动量更新从网络权重扩展到 BN 统计量，提升教师网络的稳定性。
• 在普通 GPU 硬件上，仅使用标准且高效的低批量训练，实现高性能的自监督表示学习。

提出的方法

• 提出‘动量 BN’——一种新型操作，通过在训练步骤间维护批量统计量（均值和方差）的运行平均值，而非依赖当前小批量的统计量。
• 像标准学生-教师框架一样，对学生的网络权重和教师的网络权重应用动量更新。
• 将动量更新扩展到教师的 BN 统计量，使用历史批量统计量的移动平均值以稳定归一化过程。
• 保持教师网络不进行反向传播，避免因 BN 统计量历史的不可微性而引发的梯度问题。
• 使用标准小批量 BN 的学生网络以保证效率，而教师网络则利用动量 BN 实现稳定且精确的特征归一化。
• 通过对比学习或自举学习目标训练学生网络，教师网络提供目标表示。

实验结果

研究问题

• RQ1 基于动量的批归一化统计量更新是否能提升学生-教师自监督学习的稳定性和性能？
• RQ2 在不依赖大批次或基于 TPU 的训练条件下，小批量训练（如 128）在自监督表示学习中能达到多大程度的 SOTA 性能？
• RQ3 与标准 BN、打乱 BN 或同步 BN 相比，动量 BN 在准确率和训练效率方面是否表现更优？
• RQ4 所提出的 Momentum² Teacher 框架是否能在不同下游任务和数据集上实现泛化？
• RQ5 与 BYOL 和 MoCo 等 SOTA 自监督模型相比，该方法在线性评估和微调基准测试中的表现如何？

主要发现

• Momentum² Teacher 方法在仅使用 128 批次大小的情况下，于 ImageNet 上的线性评估协议下实现了 SOTA 的 74.5% top-1 准确率。
• 该方法优于使用打乱 BN 的 MoCo，并且在不依赖大批次训练或 TPU 硬件的条件下，性能与 BYOL 相当。
• 在 COCO 目标检测和实例分割任务上，该方法显著优于监督预训练基线，并大幅超越 MoCo。
• 在长尾类别较多的 LVIS 实例分割任务中，该方法相比监督预训练将 mAP 提升了 1.7 个百分点，表明其在低资源类别上具有强大的泛化能力。
• 在仅使用 1% 和 10% ImageNet 标签的半监督学习设置下，该方法分别实现了 57.7% 和 62.3% 的 top-1 准确率，超越了所有先前的自监督和半监督方法。
• 该方法在检测和分割任务之间表现出强大的迁移能力，尤其在长训练周期下，对密集目标检测器（如 RetinaNet 和 FCOS）表现尤为突出。