PTM: Pré-treinamento Uma Vez para Modelos de Todos os Tamanhos

POA: Pre-training Once for Models of All Sizes

August 2, 2024

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

cs.AI

Resumo

O pré-treinamento auto-supervisionado em larga escala abriu caminho para um modelo base lidar com diversas tarefas de visão diferentes. A maioria das metodologias de pré-treinamento treina um único modelo de determinado tamanho de cada vez. No entanto, várias restrições de computação ou armazenamento em cenários do mundo real exigem esforços substanciais para desenvolver uma série de modelos com diferentes tamanhos para implantação. Assim, neste estudo, propomos um novo framework de treinamento auto-supervisionado de três ramos, denominado POA (Pré-treinamento Uma Vez para Todos), para lidar com o problema mencionado anteriormente. Nossa abordagem introduz um ramo estudante elástico inovador em um paradigma moderno de auto-dissipação. Em cada etapa de pré-treinamento, amostramos aleatoriamente uma sub-rede do estudante original para formar o estudante elástico e treinamos todos os ramos de forma auto-dissipativa. Uma vez pré-treinado, o POA permite a extração de modelos pré-treinados de diversos tamanhos para tarefas posteriores. Notavelmente, o estudante elástico facilita o pré-treinamento simultâneo de múltiplos modelos com diferentes tamanhos, que também atua como um conjunto adicional de modelos de vários tamanhos para aprimorar o aprendizado de representação. Experimentos extensivos, incluindo vizinhos mais próximos, avaliação de sondagem linear e avaliações em diversas tarefas posteriores, demonstram a eficácia e as vantagens de nosso POA. Ele alcança desempenho de ponta usando espinha dorsal ViT, Transformer Swin e ResNet, produzindo cerca de cem modelos com diferentes tamanhos por meio de uma única sessão de pré-treinamento. O código está disponível em: 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.