POA: Pre-training Once for Models of All Sizes

📄 arXiv: 2408.01031v1 📥 PDF

作者: Yingying Zhang, Xin Guo, Jiangwei Lao, Lei Yu, Lixiang Ru, Jian Wang, Guo Ye, Huimei He, Jingdong Chen, Ming Yang

分类: cs.CV

发布日期: 2024-08-02

备注: Accepted by ECCV2024

🔗 代码/项目: GITHUB

💡 一句话要点

提出POA：一次预训练得到各种尺寸的模型，解决部署难题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 自监督学习 预训练模型 模型压缩 知识蒸馏 弹性网络 视觉任务 模型部署

📋 核心要点

  1. 现有预训练方法通常针对特定尺寸模型，难以适应实际部署中计算资源受限的场景，需要大量工作训练不同尺寸模型。
  2. POA引入弹性学生分支，通过自蒸馏方式同时训练多个尺寸模型，实现一次预训练即可获得多种尺寸的模型。
  3. 实验表明，POA在ViT、Swin Transformer和ResNet等骨干网络上均取得了SOTA性能，并能生成上百个不同尺寸的模型。

📝 摘要（中文）

大规模自监督预训练已经为单个基础模型处理各种视觉任务铺平了道路。然而，大多数预训练方法一次只训练一个特定尺寸的模型。在实际应用中，计算或存储资源的限制往往需要开发一系列不同尺寸的模型进行部署，这需要大量的精力。因此，本研究提出了一种新颖的三分支自监督训练框架，称为POA（Pre-training Once for All），以解决上述问题。我们的方法在现代自蒸馏范式中引入了一个创新的弹性学生分支。在每个预训练步骤中，我们从原始学生网络中随机采样一个子网络，形成弹性学生网络，并以自蒸馏的方式训练所有分支。预训练完成后，POA允许提取各种尺寸的预训练模型用于下游任务。值得注意的是，弹性学生网络促进了同时预训练多个不同尺寸的模型，并且可以作为各种尺寸模型的额外集成，以增强表征学习。大量的实验，包括k近邻、线性探测评估以及在多个下游任务上的评估，证明了POA的有效性和优势。它使用ViT、Swin Transformer和ResNet骨干网络实现了最先进的性能，通过一次预训练过程生成大约一百个不同尺寸的模型。代码已开源。

🔬 方法详解

问题定义：论文旨在解决实际应用中，由于计算或存储资源限制，需要部署不同大小的模型，而传统预训练方法需要为每种尺寸单独训练模型，导致效率低下的问题。现有方法的痛点在于无法一次性获得多种尺寸的预训练模型，需要重复训练。

核心思路：论文的核心思路是引入一个“弹性学生分支”，该分支可以在训练过程中随机采样不同的子网络，从而模拟不同尺寸的模型。通过自蒸馏的方式，让弹性学生分支学习教师分支的知识，同时原始学生分支也进行学习。这样，一次预训练就可以得到多个不同尺寸的模型。

技术框架：POA框架包含三个分支：教师分支、原始学生分支和弹性学生分支。教师分支提供知识指导，原始学生分支进行常规学习，弹性学生分支则通过随机采样子网络模拟不同尺寸的模型。训练过程中，所有分支都以自蒸馏的方式进行训练，即学生分支学习教师分支的输出。预训练完成后，可以从原始学生分支和弹性学生分支中提取不同尺寸的预训练模型。

关键创新：POA的关键创新在于弹性学生分支的设计。通过随机采样子网络，弹性学生分支能够模拟不同尺寸的模型，从而实现一次预训练得到多种尺寸模型的目标。此外，弹性学生分支还可以看作是多个不同尺寸模型的集成，有助于提升表征学习能力。

关键设计：弹性学生分支的关键设计在于如何进行子网络采样。论文中采用随机采样的方式，即在每个训练步骤中，随机选择一部分网络层或神经元进行训练。损失函数方面，主要采用自蒸馏损失，即最小化学生分支和教师分支输出之间的差异。具体的网络结构选择ViT、Swin Transformer和ResNet等常见的backbone，并根据实际情况进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

POA在多个下游任务上取得了SOTA性能。例如，在ImageNet数据集上，使用ViT作为骨干网络，POA预训练的模型在top-1准确率上超过了其他自监督学习方法。此外，POA能够通过一次预训练生成上百个不同尺寸的模型，极大地提高了模型部署的灵活性。

🎯 应用场景

POA技术可广泛应用于各种需要部署不同尺寸模型的场景，例如移动设备上的图像识别、边缘计算中的目标检测、以及资源受限环境下的视觉任务。该方法能够显著降低模型训练成本，提高部署效率，并促进AI技术在更广泛领域的应用。

📄 摘要（原文）

Large-scale self-supervised pre-training has paved the way for one foundation model to handle many different vision tasks. Most pre-training methodologies train a single model of a certain size at one time. Nevertheless, various computation or storage constraints in real-world scenarios require substantial efforts to develop a series of models with different sizes to deploy. Thus, in this study, we propose a novel tri-branch self-supervised training framework, termed as POA (Pre-training Once for All), to tackle this aforementioned issue. Our approach introduces an innovative elastic student branch into a modern self-distillation paradigm. At each pre-training step, we randomly sample a sub-network from the original student to form the elastic student and train all branches in a self-distilling fashion. Once pre-trained, POA allows the extraction of pre-trained models of diverse sizes for downstream tasks. Remarkably, the elastic student facilitates the simultaneous pre-training of multiple models with different sizes, which also acts as an additional ensemble of models of various sizes to enhance representation learning. Extensive experiments, including k-nearest neighbors, linear probing evaluation and assessments on multiple downstream tasks demonstrate the effectiveness and advantages of our POA. It achieves state-of-the-art performance using ViT, Swin Transformer and ResNet backbones, producing around a hundred models with different sizes through a single pre-training session. The code is available at: https://github.com/Qichuzyy/POA.