Layerwise Importance Sampled AdamW (LISA)

• The paper introduces LISA, a novel fine-tuning method that selects important layers via importance sampling to achieve on-par or superior performance compared to full-parameter AdamW.
• It employs selective freezing of layers to dramatically reduce GPU memory usage by updating only a dynamically chosen subset during backpropagation.
• Empirical evaluations show that LISA converges faster and boosts performance metrics in tasks like MT-Bench, GSM8K, and PubMedQA across various LLMs.

Layerwise Importance Sampled AdamW (LISA) is a memory-efficient fine-tuning strategy for LLMs that leverages the empirical skewness of weight-norm changes across model layers during adaptation. Unlike Low-Rank Adaptation (LoRA), which inserts learnable low-rank adapters at each layer, LISA applies importance sampling to select a small, dynamically changing subset of layers for AdamW optimization, randomly freezing the remainder. This approach maintains or exceeds the fine-tuning performance of both LoRA and full-parameter AdamW, while matching or reducing memory requirements associated with optimizer state and parameter updates (Pan et al., 2024).

1. Background and Motivation

Full-parameter AdamW fine-tuning of LLMs requires substantial GPU memory, as it necessitates storing all gradients, first and second moment buffers per parameter, and activations. For example, a 7B parameter model typically needs at least 60 GB of GPU memory, posing a barrier for researchers without access to large hardware resources.

LoRA significantly reduces trainable parameters by introducing low-rank adapters into each linear layer. However, in many large-scale settings, such as continual pre-training and instruction tuning, LoRA can underperform relative to full fine-tuning, as its parameter search is confined to a low-rank subspace. Detailed examination reveals that LoRA’s layerwise weight-norm changes are highly skewed: only the embedding and head layers exhibit substantial updates, while intermediate self-attention blocks receive minimal changes. By contrast, full-parameter tuning yields more uniform layer updates.

This consistent skewness motivates a strategy that prioritizes “important” layers—those experiencing larger updates—by allocating computational resources preferentially to them while freezing less critical layers. The LISA algorithm operationalizes this insight through stochastic, importance-driven parameter updates.

2. The LISA Algorithm: Core Mechanism

LISA (Layerwise Importance Sampled AdamW) is defined by layerwise importance sampling and selective freezing within the AdamW optimization framework. The algorithm proceeds in the following steps:

• For a model with $L$ transformer layers (including embedding and head), total iterations $T$, sampling interval $K$, and importance sampling probabilities $p_1, ..., p_L$, initialize parameters, and AdamW moment buffers.
• At every $K$ steps, sample an “active set” $\mathcal{A}$ of layers according to the fixed distribution $p_l$, with special treatment to always include embedding and head layers. The complement set $\mathcal{M}$ is frozen.
• Conduct forward computation through all layers for activation purposes, but during backpropagation, set gradients of frozen layers to zero: $\nabla_{\theta_\ell}\mathcal{L} = 0$ for all $\ell \in \mathcal{M}$.
• Apply AdamW updates only to active layers. For each $T$0:

$T$1

$T$2

$T$3

$T$4

$T$5

$T$6

• Frozen layers maintain their current parameter values and moment estimates.

Formally, with active set $T$7 and frozen set $T$8:

$T$9

$K$0

3. Importance Sampling for Layer Selection

A critical component of LISA is the construction of the importance sampling distribution over layers. The theoretical importances $K$1 can be formulated as either the L2 norm of layer parameters, $K$2, or the expected L2 norm of the gradient, $K$3. These are normalized to derive sampling probabilities:

$K$4

Empirical investigation reveals dominance of embedding and head layer norms, so their probabilities are fixed at 1.0. The remaining probability is distributed uniformly among $K$5 selected intermediate layers per interval. For practical purposes, exactly $K$6 non-embedding, non-head layers are selected at each $K$7 steps, while embedding/head are always active.

4. Memory Complexity and Efficiency

Let $K$8 denote the number of scalar parameters, $K$9 the layer count, and $p_1, ..., p_L$0 the LoRA rank per layer. The memory usage patterns for competing approaches are:

Method	Parameter Size	Optimizer State	Activations	Adapter Overhead
AdamW (full tuning)	$p_1, ..., p_L$1	$p_1, ..., p_L$2	proportional to $p_1, ..., p_L$3	None
LoRA (rank $p_1, ..., p_L$4)	$p_1, ..., p_L$5 + $p_1, ..., p_L$6	$p_1, ..., p_L$7 (for adapter states)	proportional to $p_1, ..., p_L$8	$p_1, ..., p_L$9
LISA (with $K$0 active)	$K$1	$K$2 (for active layers only)	proportional to $K$3	None

For $K$4, LISA’s optimizer-state memory is reduced to $K$5, significantly less than the $K$6 required by full AdamW and generally smaller than LoRA’s adapter overhead. Empirical results indicate that, with typical settings ($K$7–$K$8, $K$9–$\mathcal{A}$0), LISA uses within 5–10% of LoRA’s peak memory (see Table 1 of (Pan et al., 2024)).

5. Empirical Evaluation

Experiments evaluate LISA, LoRA, and full-tuning across models including GPT2-Small, TinyLlama (1.1B), Phi-2 (2.7B), Mistral-7B, LLaMA-2-7B, and LLaMA-2-70B. Tasks encompass instruction following (Alpaca GPT-4 finetuning, measured by MT-Bench), mathematics (GSM8K), and medical QA (PubMedQA). Key findings include:

• MT-Bench (LLaMA-2-7B): LISA ($\mathcal{A}$1, $\mathcal{A}$2) achieves 5.42, surpassing full-tuning (5.18) and LoRA ($\mathcal{A}$3, 4.86), for improvements of +11% vs LoRA and +4.6% vs full-tuning.
• MT-Bench (LLaMA-2-70B): LISA ($\mathcal{A}$4, $\mathcal{A}$5) yields 7.05, exceeding full-tuning (6.66) and LoRA (6.52).
• GSM8K (LLaMA-2-70B): Accuracy improves from 59.4% (LoRA) to 61.1% (LISA).
• PubMedQA (LLaMA-2-70B): Accuracy increases from 90.8% (LoRA) to 91.6% (LISA).
• Gains are larger on smaller models: TinyLlama MT-Bench average rises from 2.03 (LoRA) to 2.78 (LISA, +37%); Mistral-7B from 4.71 (LoRA) to 5.23 (LISA, +11%).
• LISA converges faster, as shown in training loss trajectories, and can outperform full-tuning in aspects sensitive to alignment, such as writing and humanities.

6. Ablation and Sensitivity Analysis

Performance trade-offs are analyzed with respect to the number of active layers ($\mathcal{A}$6) and the sampling interval ($\mathcal{A}$7):

• Increasing $\mathcal{A}$8 leads to improved MT-Bench scores but higher memory consumption.
• Reducing $\mathcal{A}$9 (more frequent re-sampling) accelerates convergence up to an optimal point.
• Sensitivity to sampling randomness is minimal: across three random seeds, MT-Bench score variance is ≤0.13.

7. Implementation and Practical Considerations

LISA is compatible with any PyTorch-style training loop by zeroing gradients for frozen layers or toggling $p_l$0 on parameters. Recommended hyperparameters:

• Learning rate: $p_l$1 for LISA (and LoRA) on 1B–7B models; $p_l$2 for full-tuning.
• Number of active layers: $p_l$3 (7B), $p_l$4 (70B).
• Sampling interval: $p_l$5 between 3–10 is effective; up to $p_l$6 for large-scale runs.
• AdamW settings: $p_l$7, $p_l$8, $p_l$9, weight-decay=0.1.
• Fixed probabilities for embedding/head layers; uniform allocation for remaining layers.
• Can be combined with DeepSpeed ZeRO-Offload or inference-time quantization (e.g., QLoRA) for additional memory savings.

A plausible implication is that LISA constitutes a tractable alternative to adapter-based or full-parameter fine-tuning frameworks for LLMs, particularly in GPU-limited environments. By exploiting skewed utility across transformer layers, it achieves improved or on-par downstream performance with materially reduced memory footprint (Pan et al., 2024).