Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PreLoRA: Hybrid Pre-training of Vision Transformers with Full Training and Low-Rank Adapters

Krishu K Thapa, Reet Barik, Krishna Teja Chitty-Venkata, Murali Emani, Venkatram Vishwanath

Abstract

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. As training progresses, these changes stabilize, suggesting that the resulting updates may be amenable to approximation using low intrinsic-rank matrices. 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%.

PreLoRA: Hybrid Pre-training of Vision Transformers with Full Training and Low-Rank Adapters

Abstract

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. As training progresses, these changes stabilize, suggesting that the resulting updates may be amenable to approximation using low intrinsic-rank matrices. 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%.

Paper Structure

This paper contains 15 sections, 10 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. PreLoRA
  4. Partial Convergence Test
  5. Rank Assignment Algorithm
  6. Warmup (Full model + LoRA)
  7. Experimental Results
  8. Experimental Setup
  9. Results
  10. Accuracy-Efficiency Trade-offs
  11. Effect of warmup window size ($w$):
  12. Time, compute, and memory:
  13. Integration with ReLoRA:
  14. Conclusion
  15. Acknowledgment

Figures (10)

  • Figure 2: PreLoRA overall training workflow
  • Figure 3: ViT-Large: weight norm of Query module for all layers across training epochs.
  • Figure 4: Ablation Study. Comparison of different PreLoRA settings to the full baseline model in terms of accuracy (a,c,d) and time (b). Displays epoch 100 and later for clarity
  • Figure 5: Ablation Study. Comparison of different LoRA warmup windows: (a) Cross-Entropy loss, and (b) training speedup.
  • Figure 6: Effect of Weight norms (Query) on warmup windows for (a) the Base model and (b) LoRA parameters.
  • ...and 5 more figures