[Paper Review] Automated Progressive Learning for Efficient Training of Vision Transformers

Changlin Li, Bohan Zhuang+4 more|arXiv (Cornell University)|Mar 28, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes AutoProg, an automated progressive learning framework that accelerates Vision Transformer (ViT) training by dynamically and efficiently growing model capacity during training. By formulating growth schedule optimization as a sub-network architecture search and leveraging one-shot performance estimation via an elastic supernet with parameter recycling, AutoProg achieves up to 85.1% training speedup on ImageNet with no accuracy drop across DeiT and VOLO models.

ABSTRACT

Recent advances in vision Transformers (ViTs) have come with a voracious appetite for computing power, high-lighting the urgent need to develop efficient training methods for ViTs. Progressive learning, a training scheme where the model capacity grows progressively during training, has started showing its ability in efficient training. In this paper, we take a practical step towards efficient training of ViTs by customizing and automating progressive learning. First, we develop a strong manual baseline for progressive learning of ViTs, by introducing momentum growth (MoGrow) to bridge the gap brought by model growth. Then, we propose automated progressive learning (AutoProg), an efficient training scheme that aims to achieve lossless acceleration by automatically increasing the training overload on-the-fly; this is achieved by adaptively deciding whether, where and how much should the model grow during progressive learning. Specifically, we first relax the optimization of the growth schedule to sub-network architecture optimization problem, then propose one-shot estimation of the sub-network performance via an elastic supernet. The searching overhead is reduced to minimal by recycling the parameters of the supernet. Extensive experiments of efficient training on ImageNet with two representative ViT models, DeiT and VOLO, demonstrate that AutoProg can accelerate ViTs training by up to 85.1% with no performance drop. Code: https://github.com/changlin31/AutoProg

Motivation & Objective

• To address the high computational and environmental cost of training large Vision Transformers (ViTs), which scale poorly with model size and data.
• To improve training efficiency by rethinking the role of all parameters in every training step, especially in early training stages.
• To develop a generalizable, automated progressive learning scheme for ViTs that avoids manual hyperparameter tuning.
• To achieve lossless training acceleration by dynamically adapting model growth based on real-time performance feedback.
• To reduce search overhead in progressive learning through parameter recycling and one-shot supernet estimation.

Proposed method

• Proposes a growth operator called Momentum Growth (MoGrow) that uses an effective momentum update scheme to bridge performance gaps during model capacity expansion.
• Reformulates the growth schedule optimization as a sub-network architecture search problem to enable automated, adaptive growth decisions.
• Introduces an elastic supernet that enables one-shot estimation of sub-network performance, drastically reducing search cost.
• Recycles parameters from the supernet across training stages to minimize computational overhead and maintain efficiency.
• Employs adaptive regularization (AdaReg) that scales augmentation intensity and dropout probability with model size, improving stability and performance.
• Uses a two-stage training process: first train a sub-network, then grow it to full capacity with optimized growth parameters.

Experimental results

Research questions

• RQ1 Can progressive learning be automated to achieve consistent training acceleration across diverse ViT architectures without performance degradation?
• RQ2 How can the growth schedule be optimized efficiently without expensive retraining or hyperparameter tuning?
• RQ3 To what extent can one-shot supernet estimation and parameter recycling reduce the search overhead in progressive learning?
• RQ4 Does momentum-based weight initialization (MoGrow) outperform standard methods like interpolation or random initialization during model growth?
• RQ5 How does adaptive regularization affect convergence and accuracy in progressive training of ViTs?

Key findings

• AutoProg achieves up to 85.1% training speedup on VOLO-D1 with no drop in ImageNet top-1 accuracy, demonstrating lossless acceleration.
• The proposed MoGrow growth operator improves final accuracy by 0.2% over baseline methods and boosts supernet performance by 2.73%.
• Weight recycling increases speedup by 12.3% on VOLO-D1, achieving 65.6% speedup, with no accuracy degradation.
• Adaptive regularization is critical for VOLO-D1, where its absence causes a 1.2% accuracy drop, but is ineffective for DeiT-S, which fails to converge when used.
• On DeiT-S, AutoProg achieves 40.7% speedup without AdaReg, and 85.1% speedup on VOLO-D1 with AdaReg, showing strong generalization across models.
• Ablation studies confirm that MoGrow, weight recycling, and AdaReg are all essential components for optimal performance in AutoProg.