LISA: Layerwise Importance Sampling for Memory-Efficient Large Language Model Fine-Tuning

TL;DR

This work tackles the memory bottleneck in fine-tuning very large language models by uncovering a skew in layerwise weight updates under LoRA and introducing Layerwise Importance Sampled AdamW (LISA). LISA uses per-layer importance sampling to selectively unfreeze layers during optimization, achieving memory use on par with LoRA while often outperforming both LoRA and full-parameter fine-tuning across diverse tasks and model sizes (up to 70B parameters). The approach demonstrates strong results on MT-Bench, GSM8K, PubMedQA, MMLU, AGIEval, and WinoGrande, with notable gains in instruction-following directives and memory savings enabling training on modest hardware. Theoretical convergence guarantees accompany empirical evidence, and results suggest LISA as a practical, domain-robust alternative to LoRA for memory-efficient LLM fine-tuning, with potential for future enhancements by combining width-focused and depth-focused strategies.

Abstract

The machine learning community has witnessed impressive advancements since large language models (LLMs) first appeared. Yet, their massive memory consumption has become a significant roadblock to large-scale training. For instance, a 7B model typically requires at least 60 GB of GPU memory with full parameter training, which presents challenges for researchers without access to high-resource environments. Parameter Efficient Fine-Tuning techniques such as Low-Rank Adaptation (LoRA) have been proposed to alleviate this problem. However, in most large-scale fine-tuning settings, their performance does not reach the level of full parameter training because they confine the parameter search to a low-rank subspace. Attempting to complement this deficiency, we investigate the layerwise properties of LoRA on fine-tuning tasks and observe an unexpected but consistent skewness of weight norms across different layers. Utilizing this key observation, a surprisingly simple training strategy is discovered, which outperforms both LoRA and full parameter training in a wide range of settings with memory costs as low as LoRA. We name it Layerwise Importance Sampled AdamW (LISA), a promising alternative for LoRA, which applies the idea of importance sampling to different layers in LLMs and randomly freezes most middle layers during optimization. Experimental results show that with similar or less GPU memory consumption, LISA surpasses LoRA or even full parameter tuning in downstream fine-tuning tasks, where LISA consistently outperforms LoRA by over 10%-35% in terms of MT-Bench score while achieving on-par or better performance in MMLU, AGIEval and WinoGrande. On large models, specifically LLaMA-2-70B, LISA surpasses LoRA on MT-Bench, GSM8K, and PubMedQA, demonstrating its effectiveness across different domains.

Figures (12)

• Figure 1: Training loss of LLaMA-2-7B model on Alpaca GPT-4 dataset with Full Parameter Training (FT), LoRA, GaLore, and LISA.
• Figure 2: Layer-wise weight norms during training of GPT2 and LLaMA-2-7B Model with LoRA and Full Parameters training.
• Figure 3: GPU memory consumption of LLaMA-2-7B with different methods and batch size 1.
• Figure 4: Single-iteration time cost of LLaMA-2-7B with different methods and batch size 1.
• Figure 5: Generated images using LoRA (left) and LISA (right) on Stable Diffusion v2.1 model and Stable Diffusion v1.5. First row: number of inference step = 2. Second row: number of inference step = 10.

Theorems & Definitions (1)

• Theorem 1