POA: Pre-training Once for Models of All Sizes
Authors: name: Yingying Zhang; name: Xin Guo; name: Jiangwei Lao; name: Lei Yu; name: Lixiang Ru; name: Jian Wang; name: Guo Ye; name: Huimei He; name: Jingdong Chen; name: Ming Yang
Published: 2024-08-02 06:13:29

Abstract: 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.