PreLoRA: Hybrid Pre-training of Vision Transformers with Full Training and Low-Rank Adapters
This incremental improvement addresses the problem of resource-intensive training for researchers and practitioners working with large vision models.
The paper tackles the high resource cost of training large vision transformers by proposing a hybrid pre-training method that dynamically switches from full training to low-rank adaptation based on convergence, achieving a 3x throughput improvement and 20% GPU memory reduction while preserving accuracy.
Training large models ranging from millions to billions of parameters is highly resource-intensive, requiring significant time, compute, and memory. It is observed that most of the learning (higher change in weights) takes place in the earlier stage of the training loop. These changes stabilize as training continues, enabling them to be captured by matrices of a low intrinsic rank. Therefore, we propose an approach to identify such states of partial convergence and dynamically switch from full parameter training to Low-Rank Adaptation (LoRA) on the ViT-Large model. We introduce a flexible approach that leverages user-defined hyperparameters to determine the switching point and assign a rank specific to each module layer based on its level of convergence. Experimental results show that this approach preserves model accuracy while reducing the number of trainable parameters to 10% of its original size, resulting in a 3x improvement in throughput, and a 1.5x reduction in average training time per epoch while also reducing GPU memory consumption by 20%