PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rui Pan, Xiang Liu, Shizhe Diao, Renjie Pi, Jipeng Zhang, Chi Han, Tong Zhang

arXiv:2403.17919v4cs.LGcs.AIcs.CLmath.OC

TL;DR

Large language model fine-tuning consumes substantial GPU memory, while LoRA does not universally reach full-parameter fine-tuning performance. LISA samples layers for updating and randomly freezes others; the authors report stronger downstream performance than LoRA with similar or lower memory consumption.

  • Problem

    LLM fine-tuning has substantial memory costs, while LoRA does not universally match full-parameter fine-tuning in large-scale settings.

    Source: §1 · pp. 1–2

  • Method

    LISA uses LoRA’s skewed layerwise weight norms to guide importance sampling and randomly freezes most layers during full-parameter optimization.

    Source: §1 · p. 2

  • Results

    10%-35%: LISA outperforms LoRA on MT-Bench and achieves better performance on multiple benchmarks, with gains reported across 7B–70B models and varied tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors present LISA as a promising alternative to LoRA, including for large-scale and domain-specific LLM training.

    Source: §1 · p. 2, §6 · p. 10

    LISA’s forward pass still requires the model in memory, a bottleneck the authors propose may be addressed with approaches like QLoRA.

    Source: Limitations · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Large-scale fine-tuning is computationally demanding, and LoRA does not universally match full-parameter fine-tuning. The paper introduces LISA, motivated by skewed layerwise weight norms, and reports stronger downstream performance with LoRA-like memory use.

1 Introduction

  • LoRA has not universally surpassed full-parameter fine-tuning, and it can falter on large-scale datasets during continual pretraining.The paper cites LoRA’s smaller number of trainable parameters as limiting its representation power.Source: §1 · p. 2
  • LoRA’s layerwise weight norms are skewed: bottom and/or top layers account for most update weights, while other self-attention layers account for little.The authors use this pattern to motivate sampling layers according to their importance.Source: §1 · p. 2
  • With less or similar memory consumption as LoRA, LISA selectively updates essential layers and leaves others untouched, including when training models with at least 65B parameters.The authors describe LISA as outperforming both LoRA and conventional full-parameter fine-tuning on downstream tasks.Source: §1 · p. 2
  • 10%-35%: LISA outperforms LoRA on MT-Bench and achieves better performance across multiple benchmarks.The authors also report gains across models from 7B to 70B and tasks including instruction following, medical QA, and math.Source: §1 · p. 2

2 Related Work

The related work describes memory constraints in large language models and reviews parameter-efficient and optimizer-based approaches. It positions LISA among layer-selection methods while noting LoRA’s advantages and performance gap with full-parameter tuning.

2.1 Large Language Models

  • Models with at least 65B parameters can make fine-tuning infeasible in low-resource settings because of GPU memory consumption.The section frames more efficient LLM training as a response to this constraint.Source: §2.1 · p. 2

2.2 Parameter-Effieient Fine-Tuning

  • PEFT methods include prompt learning, adapters, and selective updates; prior layer-freezing methods have not been widely adopted for LLMs because of complexity or incompatibility with memory-reduction techniques.Selective methods optimize part of the model without appending extra modules, making them closer to LISA.Source: §2.2 · p. 3

2.3 Low-Rank Adaptation (LoRA)

  • LoRA reduces trainable parameters and can merge its low-rank matrices into the original model, but its performance is not universally comparable with full-parameter fine-tuning.The section notes this gap in some tasks and large-scale pretraining settings.Source: §2.3 · p. 3

2.4 Large-scale Optimization Algorithms

  • MeZO, Sophia, and GaLore offer alternative efficiency strategies, but the section notes performance, complexity, or large-model evaluation concerns; Adam and AdamW remain predominant in LLM settings.The authors characterize LoRA variants with AdamW as the dominant large-size fine-tuning paradigm whose performance still needs improvement.Source: §2.4 · p. 3

3 Method

The method is motivated by LoRA’s uneven layerwise weight norms and uses probabilistic layer sampling to select which model layers to update. LISA retains full-parameter layers while limiting the number unfrozen during optimization.

3.1 Motivation

  • Under LoRA, embedding or language-model-head layers often have weight norms hundreds of times larger than intermediate layers, a pattern not salient under full-parameter training.Figure 2 plots layer ID from embedding to final layer against weight norm.Source: §3.1 · p. 4
  • The observed difference suggests that LoRA and full-parameter training emphasize layers differently, though the paper notes multiple possible interpretations.The authors state that LoRA’s representation is restricted by its low-rank space.Source: §3.1 · p. 4

3.2 Layerwise Importance Sampled AdamW (LISA)

  • LISA uses layerwise importance sampling to approximate LoRA’s update pattern in full-parameter training, assigning smaller unfreezing probabilities to layers with small LoRA weight norms.Its practical probabilities set the embedding and LM head to 1.0 and intermediate layers to γ/NL; γ controls the expected number of unfrozen layers.Source: §3.2 · pp. 4–5
  • Each sampling period freezes all layers except the embedding and LM head, samples γ intermediate layers to unfreeze, and runs AdamW for K iterations.Randomly sampling γ layers bounds the maximum number unfrozen during training.Source: §3.2 · p. 5

4 Experimental Results

Across moderate- and large-scale fine-tuning and continual pre-training, LISA matches or exceeds baseline performance while using less or comparable memory. Its ablations show trade-offs between sampling choices, memory use, and performance.

4.1 Memory Efficiency

  • For LLaMA-2-70B, LISA E+H+2L uses 75G peak memory versus 79G for LoRA Rank 128; on LLaMA-2-7B, it fits a single 24GB RTX4090.LISA also provides almost 2.9× speedup over full-parameter training and ∼1.5× over LoRA in single-iteration time.Source: Table 1 · p. 5, §4.1 · p. 6

4.2 Moderate Scale Fine-Tuning

  • LISA outperforms full-parameter training, LoRA, and GaLore on most moderate-scale evaluation tracks, with especially large gaps on instruction-following tasks.LISA also sometimes exceeds full-parameter training; the authors suggest restricted unfrozen layers may provide an implicit regularization effect similar to dropout.Source: Table 2 · p. 7, Table 3 · p. 7, §4.2 · p. 7

4.3 Moderate Scale Continual Pre-training

  • In mathematics-domain continual pre-training followed by GSM8K fine-tuning, LISA matches or exceeds full-parameter performance using half the memory.The continual pre-training data is a 1.5-billion-token subset of OpenWebMath.Source: §4.3 · pp. 7–8, Table 4 · p. 8

4.4 Large Scale Fine-Tuning

  • On LLaMA-2-70B, LISA is on par with or better than LoRA and surpasses full-parameter training on instruction-tuning tasks.The comparison also includes MT-Bench, GSM8K, and PubMedQA.Source: Table 5 · p. 8, §4.4 · p. 8

4.5 Ablation Studies

  • More sampled layers generally improve performance at higher memory cost, while the sampling period has an optimum threshold; across three seeds, performance varied by at most 0.13.The study varies sampling-layer count γ and sampling period K.Source: Table 6 · p. 8, §4.5 · p. 9

5 Discussion

Layerwise importance sampling retains convergence guarantees in the original loss, with LISA’s convergence result relying on stated convexity, smoothness, and boundedness assumptions. The discussion also suggests different layerwise emphases may suit memorization and reasoning tasks.

5 Discussion

  • Layerwise importance sampling preserves convergence guarantees in the original loss, unlike LoRA’s added parameters and changed loss objectives.The discussion says unbiased gradient estimates can retain guarantees despite increased variance, which can be reduced with an appropriate importance-sampling strategy.Source: §5 · p. 9
  • LISA’s convergence guarantee assumes a convex, smooth loss, a bounded convex optimization set, and bounded stochastic gradients.The guarantee follows existing RBC-Adam convergence results.Source: §5 · p. 9
  • LISA is observed to outperform LoRA more on memorization-centered tasks than reasoning-centered tasks, motivating a possible combination of their benefits.The discussion suggests width may matter for memorization and depth for reasoning, while presenting this as an intriguing observation rather than a settled finding.Source: §5 · p. 9

6 Conclusion

The paper proposes LISA, which randomly freezes LLM layers according to a probability, inspired by LoRA’s skewed layerwise weight-norm distribution. It reports improved downstream fine-tuning performance over LoRA across varied models and domain-specific tasks.

6 Conclusion

  • LISA randomly freezes LLM layers according to a given probability, using a memory-efficient freezing strategy inspired by LoRA’s skewed layerwise weight norms.The paper presents this as Layerwise Importance Sampled AdamW.Source: §6 · p. 10
  • LISA’s freezing strategy is reported to scale beyond 70B parameters with less or similar memory cost as LoRA.The paper describes LISA as a promising alternative to LoRA for LLM training.Source: §6 · p. 10
  • Across varied models, including LLaMA-2-70B, LISA is reported to improve downstream fine-tuning performance over LoRA and to work in domain-specific training.The conclusion characterizes the method as having potential as an alternative to LoRA.Source: §6 · p. 10

Limitations

LISA’s forward pass still requires the model in memory, and the paper notes that its intermediate-layer sampling strategies may not be optimal. In image generation and instruction-following experiments, the authors report advantages over LoRA and other fine-tuning baselines in specified settings.

Limitations

  • LISA’s forward pass still requires the model to be in memory, leaving significant memory consumption as a bottleneck shared with LoRA.The authors say approaches similar to QLoRA may compensate for this limitation, with further experiments planned.Source: Limitations · p. 10
  • E+H+2L and E+H+4L may not be optimal because intermediate layers are still sampled uniformly at random.The authors anticipate that incorporating data sources and model architecture into importance sampling may further improve optimizer efficiency.Source: Limitations · p. 10

A.1 Image Generation

  • LISA-generated images are reported to have more intricate detail and sharper clarity than LoRA images, including with fewer inference steps.The comparison used latent consistency distillation on Stable Diffusion v1.5 and v2.1; Figure 5 shows LoRA on the left and LISA on the right, with rows for 2 and 10 inference steps.Source: §A.1 · p. 17, Figure 5 · p. 18

A.2 Instruction Following Fine-tuning

  • Full-parameter fine-tuning reduced LLaMA-2-70B’s strong Writing performance, while LISA maintained strong performance across domains with a lower Alignment Tax.The paper describes Alignment Tax as the trade-off between performance and human alignment in instruction tuning.Source: §A.2 · p. 17
  • LISA surpasses LoRA, GaLore, and full-parameter tuning in most MT-Bench settings, with consistent advantages in Writing, STEM, and Humanities.Table 8 covers Writing, Roleplay, Reasoning, Math, Extraction, STEM, and Humanities.Source: §A.2 · p. 17, Table 8 · p. 19
  • For LLaMA-2-70B, LISA outperforms LoRA in all reported MT-Bench aspects and shows on-par or faster convergence across models.Table 9 reports the MT-Bench comparison; Figure 6 shows loss curves across models.Source: §A.2 · p. 17, Figure 6 · p. 18, Table 9 · p. 19

A.3 Continual Pre-training

In continual pre-training experiments on LLaMA-2-7B, LISA with eight sampling layers achieves accuracy comparable to full-parameter continual pre-training, while 16 layers exceed it.

A.3 Continual Pre-training

  • Eight sampled layers achieve accuracy comparable to full-parameter continual pre-training, while 16 sampled layers pass its accuracy.The comparison uses the LLaMA-2-7B model.Source: §A.3 · p. 19
  • The experiments use OpenWebMath for continual pre-training and the GSM8K train split for fine-tuning, varying LISA’s sampling-layer count γ across 2, 4, 8, and 16.The comparison is against full-parameter continual pre-training accuracy.Source: §A.3 · p. 19
  • The reported continual pre-training comparison evaluates LISA configurations on a 7B-sized LLaMA-2 model.The passage identifies LLaMA-2-7B as the model used.Source: §A.3 · p. 19

A.4.1 Sampling Layers γ

The appendix examines LISA’s sampling-layer count and sampling period, then tests how random layer selection affects training convergence.

A.4.1 Sampling Layers γ

  • With γ = 8, training loss stays below the γ = 2 and γ = 4 configurations, while the GSM8K results favor more sampling layers when the learning rate is reduced proportionally.The ablations use LLaMA-2-7B; the loss comparison uses Alpaca-GPT4 and the learning-rate comparison uses GSM8K.Source: §A.4.1 · pp. 19–20

A.4.2 Sampling Period K

  • For LLaMA-2-7B on the 52K-entry Alpaca-GPT4 dataset, K = 13 is reported as optimal for the loss curve and corresponding MT-Bench score radar graph, despite similar convergence points across tested periods.The compared periods are K = 122, K = 25, and K = 13.Source: §A.4.2 · p. 20

A.4.3 Sensitiveness to Randomness

  • Three runs with different random layer picks converge similarly, although their loss trajectories vary slightly.The runs use LLaMA-2-7B on Alpaca-GPT4 with K = 13 and γ = 2.Source: §A.4.3 · p. 20

A.5 Performance with Early Exiting

The authors test LISA with the DoLa early-exiting strategy on LLaMA-2-7B. They report that LISA does not harm the performance or representation of some evaluated layers and may improve effectiveness.

A.5 Performance with Early Exiting

  • The study tests whether LISA behaves differently when combined with DoLa, an early-exiting method cited by the authors.The cited prior work is Chuang et al. (2023).Source: §A.5 · p. 21
  • The evaluation uses LLaMA-2-7B trained on Alpaca-GPT4, with early exits at layers 0, 8, 16, and 32, and compares vanilla, vanilla with DoLa, FT with DoLa, and LISA with DoLa.The GSM8K test evaluation uses a subset of 100 questions.Source: §A.5 · p. 21
  • LISA with DoLa does not negatively affect the representation or performance of some model layers and contributes to some effectiveness improvements.The comparison is reported in Table 12.Source: §A.5 · p. 21

A.6 Comparison of Evaluation Loss

On Alpaca-GPT4, LISA’s validation-loss trend mirrors its training-loss trend, with no reported signs of overfitting. Additional Mistral-7B experiments show a layerwise weight-norm pattern similar to LLaMA-2-7B.

A.7 Additional Observations of Layerwise Skewness

  • On Mistral-7B, the bottom layer has a larger weight norm than other layers, matching the reported LLaMA-2-7B observation.Figure 12 shows layerwise weight norms during LoRA and full-parameter training.Source: §A.7 · p. 21, Figure 12 · p. 22

A.6 Comparison of Evaluation Loss

  • LISA’s validation-loss trend mirrors training loss and shows no signs of overfitting in the Alpaca-GPT4 split.The dataset was split 9:1 into training and validation sets with no overlap; Figure 11 compares FT, LoRA, GaLore, and LISA.Source: §A.6 · p. 21, Figure 11 · p. 22

B.1 Training Setup

The appendix describes the training setup, continual-pretraining data, hyperparameter search, and dataset licensing.

B.1 Training Setup

  • Experiments use LMFlow for FT, LoRA, and LISA, one epoch for fine-tuning and continual pretraining, and DeepSpeed offload on eight 48 GB NVIDIA Ampere GPUs.The learning-rate search spans 5 × 10^-6 to 3 × 10^-4; LoRA uses rank 128 or 256, while LISA typically uses γ = 2 and uses γ = 4 for 70B experiments.Source: §B.1 · p. 23

B.2 Continual Pre-training Dataset

  • The continual-pretraining subset comes from OpenWebMath, retains documents with Math_score above 0.95, and contains 1.5 billion tokens.Source: §B.2 · p. 23

B.3 Hyperparameter search

  • The hyperparameter search varies learning rate, sampling-layer count γ, and sampling period K, with optimal settings reported in Table 15.The search considers LoRA ranks 128 or 256, and generally selects γ = 2 for LISA or γ = 4 for 70B models.Source: §B.3 · pp. 23–24, Table 15 · p. 24

C Licenses

  • The named instruction-following and domain-specific datasets are released under MIT or Apache-2.0 licenses, while the GPT-4-generated dataset is designated for research purposes subject to its terms of use.Source: C Licenses · p. 24

Related papers

Loading 2403.17919v4…