POA: Pre-addestramento Unico per Modelli di Tutte le Dimensioni

POA: Pre-training Once for Models of All Sizes

August 2, 2024

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

cs.AI

Abstract

Il pre-training su larga scala in modalità self-supervised ha aperto la strada affinché un singolo modello di base possa gestire molteplici task visivi. La maggior parte delle metodologie di pre-training addestra un unico modello di una determinata dimensione alla volta. Tuttavia, i vari vincoli computazionali o di archiviazione negli scenari reali richiedono notevoli sforzi per sviluppare una serie di modelli di dimensioni diverse da distribuire. Pertanto, in questo studio, proponiamo un nuovo framework di training self-supervised a tre rami, denominato POA (Pre-training Once for All), per affrontare il problema sopra menzionato. Il nostro approccio introduce un innovativo ramo studente elastico in un moderno paradigma di self-distillation. Ad ogni passo di pre-training, campioniamo casualmente una sotto-rete dallo studente originale per formare lo studente elastico e addestriamo tutti i rami in modalità self-distilling. Una volta pre-addestrato, POA consente l'estrazione di modelli pre-addestrati di dimensioni diverse per task downstream. In modo notevole, lo studente elastico facilita il pre-training simultaneo di più modelli con dimensioni diverse, che funge anche come un ulteriore ensemble di modelli di varie dimensioni per migliorare l'apprendimento delle rappresentazioni. Esperimenti estensivi, tra cui k-nearest neighbors, valutazione lineare e test su più task downstream, dimostrano l'efficacia e i vantaggi del nostro POA. Raggiunge prestazioni all'avanguardia utilizzando backbones come ViT, Swin Transformer e ResNet, producendo circa un centinaio di modelli di dimensioni diverse attraverso una singola sessione di pre-training. Il codice è disponibile all'indirizzo: 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.