QuAILoRA: Quantization-Aware Initialization for LoRA

TL;DR

QuAILoRA is introduced, a quantization-aware initialization for LoRA that mitigates this negative impact of quantization errors at initialization by decreasing quantization errors at initialization, without increasing the memory-cost of fine-tuning.

Abstract

QLoRA reduces the memory-cost of fine-tuning a large language model (LLM) with LoRA by quantizing the base LLM. However, quantization introduces quantization errors that negatively impact model performance after fine-tuning. In this paper we introduce QuAILoRA, a quantization-aware initialization for LoRA that mitigates this negative impact by decreasing quantization errors at initialization. Our method spends a small amount of computational overhead to compute this quantization-aware initialization, without increasing the memory-cost of fine-tuning. We evaluate our method on several causal language modeling and downstream evaluation tasks using several different model sizes and families. We observe that almost all LLMs fined-tuned with QuAILoRA achieve better validation perplexity. When evaluated on downstream tasks, we find that QuAILoRA yields improvements proportional to the negative effect of quantization error. On average, applying QuAILoRA to 4-bit QLoRA models yields 75% of the validation perplexity decrease and 86% of the downstream task accuracy increase as doubling the quantization precision to 8-bit, without increasing GPU memory utilization during fine-tuning.

Figures (2)

• Figure 1: Effect of LoRA rank on validation perplexity after fine-tuning 4-bit models, averaged across six 4-bit LLMs and 4 causal language modeling tasks. Increasing the LoRA rank results in a continual decrease in average validation perplexity when initializing with our method, albeit with diminishing returns. In contrast, increasing the LoRA rank does not significantly affect performance when using the baseline initialization. We plot results for Pythia-70m separately (not included in the six-model average) as this was the only baseline to show a strong decrease in validation perplexity with increasing LoRA rank.
• Figure 2: Fine-tuning convergence for each of our 9 base models. We also include a six-model average convergence curve that excludes OPT-13b and OPT-30b (due to the perplexity spikes in the middle of fine-tuning) as well as Pythia-70m (which we fine-tune for 10 times as many steps as the other models).