POA：已开源，蚂蚁集团提出同时预训练多种尺寸网络的自监督范式

论文提出一种新颖的POA自监督学习范式，通过弹性分支设计允许同时对多种尺寸的模型进行预训练。POA可以直接从预训练teacher生成不同尺寸的模型，并且这些模型可以直接用于下游任务而无需额外的预训练。这个优势显著提高了部署灵活性，并有助于预训练的模型在各种视觉任务中取得SOTA结果。

来源：晓飞的算法工程笔记公众号

论文: POA: Pre-training Once for Models of All Sizes

论文地址：https://arxiv.org/abs/2408.01031
论文代码：https://github.com/Qichuzyy/POA

Abstract

大规模自监督预训练为一个基础模型处理多种不同的视觉任务铺平了道路。大多数预训练方法在一次训练中训练一个特定大小的单一模型。然而，在现实世界的场景中，由于各种计算或存储限制，需要大量的工作来开发一系列不同大小的模型进行部署。因此，在这项研究中，我们提出了一种新颖的三分支自监督训练框架，称为POA（Pre-training Once for All），来解决上述问题。我们的方法在现代自蒸馏范式中引入了一种创新的弹性student分支。在每个预训练步骤中，我们随机从原始student中抽样一个子网络来形成弹性student，并以自蒸馏的方式训练所有分支。一旦预训练完成，POA允许提取不同大小的预训练模型用于下游任务。值得注意的是，弹性student促进了多个不同大小模型的同时预训练，同时也作为各种大小模型的额外集合，增强了表示学习。大量实验证明了我们的POA的有效性和优势，包括k最近邻、线性探测评估以及多个下游任务的评估。它使用ViT、Swin Transformer和ResNet骨干网络实现了最先进的性能，并通过一次预训练会话生成了大约一百个不同大小的模型。代码可在以下链接找到：https😕/github.com/Qichuzyy/POA。

Introduction

通过自监督学习在大型模型中学习可泛化的视觉表示，近年来在各种视觉任务上取得了卓越的性能。然而，当部署到现实世界的应用程序时，大型模型必须根据计算、存储、功耗等各种资源限制进行调整。例如，一个良好设计的人工智能产品通常包括一套为不同场景量身定制的模型，比如Gemini Nano、Pro和Ultra。对于一个大型预训练模型，将其部署到具有不同资源约束的多个应用场景的常见解决方案包括额外的权重修剪、知识蒸馏，甚至从头开始重新训练一个小网络，这些都需要大量的开发工作。因此，这个问题引发了一个关键问题：是否可能进行一次预训练以同时生成多个具有不同大小的模型，每个模型都提供足够好的表示。

为了解决这一挑战，论文引入了一种名为POA（Pre-training Once for All）的新型自监督学习范式。POA建立在流行的teacher-student自蒸馏框架之上，具有一个额外的创新性弹性student分支。弹性student分支通过参数共享嵌入了一系列子网络，这是基于观察到对于现代网络结构来说，较小尺寸的模型是较大尺寸模型的子网络。此外，该分支的参数与原始的或完整的studennt共享。在每个预训练步骤中，从完整student中随机抽样一部分参数，形成相应的弹性studennt。原始完整student和弹性student都被训练以模拟teacher网络的输出。teacher本身通过对student参数的指数移动平均（EMA）不断优化，包括采样的弹性student。弹性student有助于在不同参数子集上进行有效和高效的预训练，从而成功地从预训练teacher中提取出高性能子网络，用于后续的下游场景。它还作为一种训练正则化形式，通过强制teacher和各种子网络之间的输出匹配来促进稳定的训练过程。

POA代表了第一个能够同时训练多个不同尺寸模型的自监督学习方法，每个模型在不需要进一步预训练的情况下，都能获得适用于不同资源约束的高质量表示。图1显示了通过POA预训练的ViT-L模型提取的143个子网络的k最近邻（k-NN）评估结果。通过选择不同的弹性宽度和深度，预训练teacher模型可以根据可用计算资源定制的适用于下游应用的合适模型，生成足够数量的候选子网络以供选择。值得注意的是，由于在同视图蒸馏上进行了精心设计，每个子网络都经过了良好训练，并表现出优越性能。特别是，ViT-S、ViT-B和ViT-L模型创造了新的基准，与那些由现有方法预训练的模型相比取得了SOTA结果。

为了严格评估方法的有效性，使用三种广泛使用的骨干架构，即ViT、Swin Transformer和ResNet，进行了大量实验。每个骨干架构都在ImageNet-1K数据集上进行了预训练，并使用k-NN和线性探测分类评估，以及在下游密集预测任务进行评估，如目标检测和语义分割。POA在单次预训练会话中跨多种模型尺寸实现了最先进的准确性。

本文的技术贡献总结如下：

POA是第一个将无监督表示学习和一次性模型生成集成到单个预训练会话中的预训练范式，解决了社区很少探讨的一次性预训练挑战。这对实际部署非常重要，因为实际部署通常需要一套模型。
提出了一个新颖而优雅的组件，称为弹性student（Elastic Student），具有一系列弹性算子，可以使POA与包括ViT、Swin Transformer和ResNet在内的流行骨干结构兼容，具备生成各种大小模型的能力。此外，还作为模型集成来平滑训练过程并改善学到的表示。
通过对k-NN、线性探测和下游密集任务评估的彻底评估，在多个指标上展现出优于现有最先进预训练方法的性能。此外，将POA与自监督蒸馏（SEED）进行了比较，SEED是一种专为自监督学习设计的知识蒸馏方法，进一步验证了POA的有效性。

POA Self-supervised Learning Framework

论文的主要目标是通过单次自监督预训练会话来预训练多种规模的模型，受到自蒸馏技术最新进展的启发,提出了一个名为POA的新型SSL（Self-supervised Learning）框架。POA架构如图2所示，包括一个teacher模型、一个完整的student模型、一个弹性student模型以及对应的头部。teacher模型使用student模型的指数移动平均（EMA）进行更新。弹性student模型是完整student模型的派生版本，其主干网络和头部参数是共享的。

在两个方面利用蒸馏技术：完整student和弹性student都是通过使用同一图像不同视图的teacher模型进行蒸馏，而弹性student还通过使用相同视图的完整student进行学习。交叉视图蒸馏作为一种表示学习形式，如所介绍的那样。值得注意的是，除了仅使用完整student进行常规EMA更新外，弹性student在每个预训练步骤中还提供一个随机抽样的子网络，参与teacher模型的EMA优化。这个过程实际上模拟了多个子网络的集成，这在监督学习领域也被证明是有益的。同视图蒸馏是完整student和弹性student之间的标准知识蒸馏，提升了弹性student的质量。

Design of Elastic Student

弹性student是一个子网络，其参数是从完整student中提取的。在transformer主干网络的背景下，宽度指的是标记的维度，而在卷积主干网络中，宽度表示通道数。深度则定义为transformer或卷积网络中基本块的数量。给定宽度和深度的值，会产生一定的网络结构。为简单起见，论文将重点放介绍ViT的弹性设计。

ViT的基本块主要由多头自注意力（MSA）模块和多层感知器（MLP）模块组成。在每个模块之前应用层归一化（LN），并在每个模块后使用残差连接。如图3的左侧所示，弹性块是指在ViT原始基本块中调整宽度后堆叠的弹性MSA、MLP和LN。在论文的方法中，弹性student分支是通过在每个训练迭代中组装特定数量的这些弹性块来构建的。

Elastic MSA

一个原始或完整的MSA模块由三个主要组件组成，即输入投影层，包含注意力和连接的操作符，以及输出投影层。将投影层定义为( $w^{\ast}, b^{\ast}$ )，其中 $w^{\ast}$ 表示线性转换权重， $b^{\ast}$ 表示相应的偏置， $\ast$ 表示层的名称。如图3的右侧所示，给定一个标记维度 $D_{max}=N_h \cdot D_h$ ，其中 $N_h$ 是注意力头的数量， $D_h$ 是头部维度，具有长度 $T$ 的输入序列 $\in \mathbb{R}^{T \times D_{max}}$ 最初被投影以形成查询 $\in \mathbb{R}^{T \times D_h}$ 、键 $\in \mathbb{R}^{T \times D_h}$ 和值 $\in \mathbb{R}^{T \times D_h}$ 。为了生成弹性MSA，定义了M+1个弹性宽度，包括 $D_{max}$ ，间隔为 $D_h$