POA: Pre-entrenamiento Único para Modelos de Todos los Tamaños

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

Resumen

El pre-entrenamiento auto-supervisado a gran escala ha allanado el camino para que un modelo base maneje muchas tareas de visión diferentes. La mayoría de las metodologías de pre-entrenamiento entrenan un único modelo de cierto tamaño a la vez. Sin embargo, diversas limitaciones de cómputo o almacenamiento en escenarios del mundo real requieren esfuerzos sustanciales para desarrollar una serie de modelos con diferentes tamaños para implementar. Por lo tanto, en este estudio, proponemos un nuevo marco de entrenamiento auto-supervisado de tres ramas, denominado POA (Pre-entrenamiento Una Vez para Todos), para abordar este problema mencionado anteriormente. Nuestro enfoque introduce una innovadora rama estudiante elástica en un paradigma moderno de auto-difusión. En cada paso de pre-entrenamiento, muestreamos aleatoriamente una subred del estudiante original para formar el estudiante elástico y entrenamos todas las ramas de manera auto-difundida. Una vez pre-entrenado, POA permite la extracción de modelos pre-entrenados de diversos tamaños para tareas posteriores. Notablemente, el estudiante elástico facilita el pre-entrenamiento simultáneo de múltiples modelos con diferentes tamaños, que también actúa como un conjunto adicional de modelos de varios tamaños para mejorar el aprendizaje de representaciones. Experimentos extensos, incluyendo vecinos más cercanos, evaluación de sondas lineales y evaluaciones en múltiples tareas posteriores, demuestran la efectividad y ventajas de nuestro POA. Logra un rendimiento de vanguardia utilizando las arquitecturas ViT, Swin Transformer y ResNet, generando alrededor de cien modelos con diferentes tamaños a través de una sola sesión de pre-entrenamiento. El código está disponible en: 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.