ПОА: Предварительное обучение один раз для моделей всех размеров

POA: Pre-training Once for Models of All Sizes

August 2, 2024

Авторы: Yingying Zhang, Xin Guo, Jiangwei Lao, Lei Yu, Lixiang Ru, Jian Wang, Guo Ye, Huimei He, Jingdong Chen, Ming Yang

cs.AI

Аннотация

Масштабное предварительное обучение без учителя открыло путь для создания одной базовой модели, способной обрабатывать множество различных задач компьютерного зрения. Большинство методов предварительного обучения обучают одну модель определенного размера за один раз. Тем не менее, различные вычислительные или ограничения по хранению в реальных сценариях требуют значительных усилий для разработки серии моделей различных размеров для развертывания. Таким образом, в данном исследовании мы предлагаем новую трехветвую структуру самообучения, названную POA (Предварительное обучение один раз для всех), чтобы решить вышеупомянутую проблему. Наш подход вводит инновационную упругую ветвь студента в современную парадигму самодистилляции. На каждом шаге предварительного обучения мы случайным образом выбираем подсеть из исходного студента для формирования упругого студента и обучаем все ветви в режиме самодистилляции. После предварительного обучения POA позволяет извлекать предварительно обученные модели различных размеров для последующих задач. Замечательно, упругий студент облегчает одновременное предварительное обучение нескольких моделей различных размеров, что также выступает в качестве дополнительного ансамбля моделей различных размеров для улучшения обучения представлений. Обширные эксперименты, включая оценку ближайших соседей, линейное зондирование и оценку на нескольких последующих задачах, демонстрируют эффективность и преимущества нашего подхода POA. Он достигает передовых результатов, используя основные модели ViT, Swin Transformer и ResNet, создавая около сотни моделей различных размеров через одну сессию предварительного обучения. Код доступен по ссылке: https://github.com/Qichuzyy/POA.

English

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.