POA : Pré-entraînement unique pour des modèles de toutes tailles

POA: Pre-training Once for Models of All Sizes

August 2, 2024

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

cs.AI

Résumé

Le pré-entraînement auto-supervisé à grande échelle a ouvert la voie à un modèle de base capable de gérer de nombreuses tâches visuelles différentes. La plupart des méthodologies de pré-entraînement entraînent un seul modèle d'une certaine taille à la fois. Cependant, diverses contraintes de calcul ou de stockage dans des scénarios réels nécessitent des efforts considérables pour développer une série de modèles de différentes tailles à déployer. Ainsi, dans cette étude, nous proposons un nouveau cadre de formation auto-supervisé à trois branches, appelé POA (Pre-training Once for All), pour résoudre ce problème susmentionné. Notre approche introduit une branche étudiante élastique innovante dans un paradigme moderne d'auto-distillation. À chaque étape de pré-entraînement, nous échantillonnons aléatoirement un sous-réseau de l'étudiant original pour former l'étudiant élastique et entraînons toutes les branches de manière auto-distillante. Une fois pré-entraîné, POA permet l'extraction de modèles pré-entraînés de tailles diverses pour les tâches en aval. Remarquablement, l'étudiant élastique facilite le pré-entraînement simultané de plusieurs modèles de différentes tailles, qui agit également comme un ensemble supplémentaire de modèles de diverses tailles pour améliorer l'apprentissage de la représentation. Des expériences approfondies, incluant les k-plus proches voisins, l'évaluation par sondage linéaire et des évaluations sur plusieurs tâches en aval, démontrent l'efficacité et les avantages de notre POA. Il atteint des performances de pointe en utilisant les architectures ViT, Swin Transformer et ResNet, produisant environ une centaine de modèles de différentes tailles grâce à une seule session de pré-entraînement. Le code est disponible à l'adresse : 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.