Efficient Training of Large Vision Models via Advanced Automated Progressive Learning

TL;DR

This work tackles the high computational cost of training large vision models by introducing Advanced AutoProg, a framework that automates progressive learning for both pre-training (ViTs) and fine-tuning (diffusion models). It combines a robust manual baseline (MoGrow) with AutoProg-One, which uses a one-shot elastic supernet to efficiently search growth schedules, and AutoProg-Zero, which applies zero-shot metrics and Unique Stage Identifiers to enable lossless unfreezing during fine-tuning. Empirical results show ViT pre-training can be accelerated up to $1.85\times$ and diffusion-model fine-tuning up to $2.86\times$, without sacrificing accuracy, and sometimes even improving it; AutoProg-Zero also demonstrates substantial speedups in text-to-image and DreamBooth tasks with improved generative quality. The framework’s use of zero-shot proxies (NTK-based trainability via condition number, and ZiCo-gradient metrics) and an elastic supernet for rapid sub-network evaluation yields scalable, broadly applicable improvements for efficient training of large vision models.

Abstract

The rapid advancements in Large Vision Models (LVMs), such as Vision Transformers (ViTs) and diffusion models, have led to an increasing demand for computational resources, resulting in substantial financial and environmental costs. This growing challenge highlights the necessity of developing efficient training methods for LVMs. Progressive learning, a training strategy in which model capacity gradually increases during training, has shown potential in addressing these challenges. In this paper, we present an advanced automated progressive learning (AutoProg) framework for efficient training of LVMs. We begin by focusing on the pre-training of LVMs, using ViTs as a case study, and propose AutoProg-One, an AutoProg scheme featuring momentum growth (MoGrow) and a one-shot growth schedule search. Beyond pre-training, we extend our approach to tackle transfer learning and fine-tuning of LVMs. We expand the scope of AutoProg to cover a wider range of LVMs, including diffusion models. First, we introduce AutoProg-Zero, by enhancing the AutoProg framework with a novel zero-shot unfreezing schedule search, eliminating the need for one-shot supernet training. Second, we introduce a novel Unique Stage Identifier (SID) scheme to bridge the gap during network growth. These innovations, integrated with the core principles of AutoProg, offer a comprehensive solution for efficient training across various LVM scenarios. Extensive experiments show that AutoProg accelerates ViT pre-training by up to 1.85x on ImageNet and accelerates fine-tuning of diffusion models by up to 2.86x, with comparable or even higher performance. This work provides a robust and scalable approach to efficient training of LVMs, with potential applications in a wide range of vision tasks. Code: https://github.com/changlin31/AutoProg-Zero

Figures (10)

• Figure 1: Accuracy of ViTs (DeiT touvron2020deit, VOLO Yuan2021VOLOVO) during training. Smaller ViTs converge faster in terms of runtime. Models in the legend are sorted in ascending order of model size.
• Figure 2: Generated images on multiple datasets with AutoProg-Zero fine-tuning, which only takes 0.39× training time of full fine-tuning. Resolution is 256×256.
• Figure 3: Variants of the growth operator $\bm\zeta$. $\bm\omega^\textit{A}$ and $\bm\omega^\textit{B}$ denote the parameters of two Transformer blocks in the original small network $\bm\psi_{k-1}$. (a) RandInit randomly initializes newly added layers; (b) Stacking duplicates the original layers and directly stacks the duplicated ones on top of them; (c) Interpolation interpolate new layers of $\bm\psi_\ell$ in between original ones and copy the weights from their nearest neighbor in $\bm\psi_{k-1}$. (d) Our proposed MoGrow is build upon Interpolation, by coping parameters $\widetilde{\bm\omega}$ from the momentum updated ensemble of $\bm\psi_{k-1}$.
• Figure 4: Pipeline of the $k$-th stage of AutoProg-One. In the beginning of the stage, the last optimal sub-network $\bm\psi^\star_{k-1}$ first grows to the Elastic Supernet $\Phi_k$ by $\widehat{\bm\omega} = \bm\zeta(\bm\omega^\star)$; then, we search for the optimal sub-network $\bm\psi^\star_{k}$ after supernet training; finally, the sub-network is trained in the remaining epochs of this stage. The whole process of AutoProg-One is summarized in \ref{['alg:autoprog']}.
• Figure 5: Pipeline of the $k$-th stage of AutoProg-Zero. In the beginning of the stage, different candidate learnable sub-networks unfreeze; then, we search for the optimal learnable sub-network $\bm\psi^\star_{k}$ according to zero-shot metrics; finally, the sub-network unfreezes and is trained in this stage. The whole process of AutoProg-Zero is summarized in \ref{['alg:autoprogzero']}.

Theorems & Definitions (1)

• Definition 3.1