MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling

AI summary

Overview

Research area: Memory-efficient optimization for large language model (LLM) pre-training and fine-tuning, specifically block-coordinate / layer-wise optimization methods.

Technical level: Intermediate. The paper combines empirical LLM fine-tuning with block-coordinate descent theory, importance sampling, and convergence analysis, but the core idea can be understood without deep optimization background.

Scope: The paper proposes MISA, a fine-grained module-wise importance sampling method for training LLMs that partitions transformer layers into smaller modules, samples them according to learned importance scores, and provides convergence guarantees under practical Adam-based training conditions.

What This Paper Is About

Training or fine-tuning LLMs requires enormous memory for optimizer states, gradients, and activations, which often exceeds available hardware. One popular workaround is layer-wise optimization, which updates one transformer block at a time while freezing the rest; however, these methods treat every module inside a layer as equally important and still must keep at least one full layer active in memory. MISA addresses both problems by breaking layers into smaller modules (such as the attention matrices W_q, W_k, W_v, W_o and the feed-forward matrices W_up and W_down), assigning each module an importance score, and sampling modules by a weighted random mechanism that balances exploitation of important modules with exploration of the rest.

Key Contributions

1. Module-wise optimization. The authors define a module as a matrix parameter within a transformer layer that is associated with a weight gradient. They show empirically that modules inside the same layer have heterogeneous gradient norms, and demonstrate theoretically that decomposing a layer into smaller modules preserves more gradient information than layer-wise partitioning. This finer granularity also removes the need to load a full layer into memory, improving memory efficiency over layer-wise methods such as BAdam, LIFT, HIFT, LISA, and BlockLLM.

2. Improved importance sampling. The sampling strategy is derived by parameterizing gradient variance as a function of the sampling probability for each module and maximizing the expected decrease in the loss. A Kullback-Leibler (KL) divergence penalty toward the uniform distribution is added to control the exploration-exploitation trade-off via a coefficient η. Proposition 1 gives a closed-form softmax solution, p_b^n = exp(η‖g_b^n‖²) / Σ_j exp(η‖g_j^n‖²). In practice, the full-batch block gradient norm is replaced by an exponential moving average (β) of the average stochastic block gradient norms over T inner steps, using a scaled gradient norm to account for differing parameter counts across modules.

3. Convergence guarantees under realistic conditions. MISA is shown to achieve a convergence rate of O(1/√K) in the abstract (K being the total number of block updates), with Theorem 1 stating the informal bound (1/N)Σ E[‖∇f(θ^{n,0})‖²] = O(1/√(NT) + 1/N), where N and T are the numbers of outer and inner-loop iterations. The analysis incorporates the Adam optimizer, stochastic gradients, and multiple updates per sampled block, unlike prior layer-wise analyses that assume noiseless gradients or only a single update per block.

4. Memory analysis and method comparison. The paper provides a detailed memory analysis showing MISA's advantage over baselines, including that layer-wise methods achieve lower peak memory than subspace methods such as LoRA and GaLore in long-sequence settings, that layer-wise methods can update more weights than LoRA under memory constraints, and that MISA is more memory-efficient than layer-wise methods when the sampling ratio threshold δ is small. Table 2 compares MISA against LoRA, IST, GaLore, OwLore, BAdam, and LISA across properties including full-rank updates, importance awareness, fine-grained memory, gradient accumulation, absence of SVD, and support for fine-tuning, pre-training, and convergence guarantees.

Main Findings

• Modules within a layer differ sharply in importance. Figure 1 shows that different modules in different layers of LLaMA3-8B have significantly varying gradient norms during fine-tuning, which is the empirical motivation for module-wise rather than layer-wise partitioning.

• Module-wise sampling provably dominates layer-wise sampling. Proposition 2 shows that any layer-wise importance sampling solution is a feasible solution to the module-wise problem, so the optimal module-wise solution yields a larger objective value and is therefore a superior sampling strategy.

• Commonsense reasoning gains. On LLaMA3-8B across BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC-e, ARC-c, and OBQA, MISA with δ=3% reached the best average of 86.6 with 34.4 GB memory, compared with full fine-tuning at 150.5 GB and 86.7, LoRA at 35.7 GB and 82.5, DoRA at 54.1 GB and 85.2, LISA at 56.3 GB and 85.9, and BAdam at 34.1 GB and 84.8. On Qwen2.5-7B, MISA with δ=3% reached 88.5, ahead of full fine-tuning's 89.3 (at 140.5 GB), LoRA's 88.2, DoRA's 88.3, LISA's 87.8, and BAdam's 88.1.

• Memory savings at 1% trainable parameters. With the trainable parameter ratio set to 1%, MISA achieved performance comparable to LoRA while saving approximately 10% of memory. On LLaMA3-8B, MISA with δ=1% used 30.7 GB versus LoRA's 35.7 GB; on Qwen2.5-7B, 30.3 GB versus 34.2 GB.

• Math reasoning gains. Fine-tuned on MATH10K and evaluated on GSM8K, SVAMP, AQuA, and MAWPS, MISA with δ=3% achieved the best average on both models: 73.6 for LLaMA3-8B (versus LoRA 71.9, DoRA 73.3, LISA 70.7, BAdam 72.0) and 82.0 for Qwen2.5-7B (versus LoRA 81.3, DoRA 81.9, LISA 81.1, BAdam 80.9), at 34.1 GB and 33.6 GB respectively.

• Instruction tuning gains. On the 52K-sample Alpaca GPT-4 dataset, MISA achieved the highest MT-Bench scores among compared methods for TinyLLaMA (2.73 at 4.36 GB), LLaMA2-7B (5.13 at 19.72 GB), and Mistral-7B (5.19 at 21.18 GB). MISA also posted the best MMLU and MMLU-pro numbers on LLaMA2-7B (46.27 and 20.69) and Mistral-7B (62.90 and 33.44), and the best MMLU-pro on TinyLLaMA (11.65).

• Pre-training competitiveness. Pre-training LLaMA2 130M and 350M on C4 for 52K steps over 2.75B tokens, MISA on the 350M model reached a perplexity of 22.11 after 2.7B training tokens, better than GaLore's 24.34 and approaching Adam's 21.3. On the 130M model with a high-rank subspace (r=256, δ=25%), MISA outperformed both GaLore and Adam. The authors suggest MISA can be viewed as a regularization of Adam and that many gradients are redundant during full-parameter training.

• Convergence holds for any number of inner updates. Remark 3 states that MISA converges for any T inner local block updates, whereas prior analyses are limited to T=1 or require T to be sufficiently large.

• Training time trade-off. MISA and LISA have similar total training time, while BAdam is faster than both; despite this, MISA shows better convergence than LISA and BAdam across three epochs on Mistral-7B, LLaMA2-7B, and TinyLLaMA.

• Embedding and LM-head layers are excluded during fine-tuning. Although LISA identifies the embedding and LM-head layers as important, MISA does not train them in fine-tuning tasks because their parameters are too large (proportional to vocabulary size), and training them would significantly increase memory consumption.

Methodology in Plain English

The approach starts from block-coordinate descent, a classic technique that splits a high-dimensional optimization problem into smaller blocks updated one at a time. The authors change the granularity of those blocks: instead of treating a whole transformer layer as one block, they treat each matrix parameter inside a layer as its own module that can be sampled and updated independently. This matters because the paper measures gradient norms during LLaMA3-8B fine-tuning and finds that modules within a single layer carry very different amounts of optimization signal.

To decide which module to update, MISA frames sampling probabilities as an optimization problem: choose probabilities that maximize the expected decrease in the loss, with a KL penalty pulling those probabilities toward a uniform distribution. The penalty, scaled by η, prevents the method from repeatedly updating only a few high-gradient modules and ignoring the rest. The closed-form answer is a softmax over each module's gradient norm squared.

Because the true full-batch gradient is unavailable during LLM training, MISA approximates each module's gradient norm with an exponentially weighted average of the stochastic gradient norms observed during that module's T inner Adam steps, updating the estimate only for the module that was just sampled and keeping prior values for all others. A scaled gradient norm is used to keep modules with different parameter counts comparable. The algorithm is a double loop: the outer loop samples a set of modules whose trainable parameter ratio stays below δ, and the inner loop performs T Adam updates on the selected module. After the inner loop, an additional Adam step is applied for the convergence analysis, and optimizer states for the selected module are cleared to keep memory usage consistent throughout training.

The theory tackles a subtle problem: when you update the same block multiple times, the block gradient is no longer an unbiased estimator of the full gradient, as conventional BCD analysis assumes. The authors derive recursive relations that track how gradient bias and stochastic noise accumulate over successive block updates, establish connections between block-level and full gradients (with the extra Adam step playing a key role in transitioning from local to global updates), and build an analytical framework that jointly controls the bias and the amplified noise.

Experiments cover fine-tuning on LLaMA3-8B, Qwen2.5-7B, TinyLLaMA, LLaMA2-7B, and Mistral-7B, plus pre-training of LLaMA2 130M and 350M, all on RTX 4090 24GB GPUs with no gradient checkpointing or flash attention used.

Why This Matters

This paper pushes block-coordinate optimization for LLMs to a finer granularity and, importantly, backs the change with a convergence analysis that matches what practitioners actually do: Adam, stochastic gradients, and repeated updates per sampled block. It also challenges the implicit assumption in prior layer-wise work that all parameters inside a transformer layer deserve equal treatment, and it quantifies when the extra granularity pays off in memory.

Real-world applications:

• Fine-tuning large models on consumer-grade hardware. All experiments run on a single RTX 4090 24GB, showing that 7B- and 8B-scale models can be adapted without datacenter-class GPUs.
• Long-sequence fine-tuning. The memory analysis shows MISA increasingly outperforms LoRA in memory efficiency as sequence length grows.
• Domain and task adaptation of open models. The commonsense, math, and instruction-tuning results span 16 datasets across three benchmark families, covering common downstream use cases.
• Pre-training on constrained budgets. The LLaMA2 130M and 350M C4 experiments indicate the method is relevant beyond fine-tuning, where memory pressure is highest.

Industry relevance centers on cost and accessibility: methods that reduce optimizer state, gradient, and activation memory let organizations train and customize models on cheaper hardware, shorten iteration cycles, and deploy domain-specific LLMs without full-scale infrastructure. Because the partitioning is at the module level, the approach fits naturally into existing transformer architectures and standard optimizer implementations.

Future Directions

• Automating the exploration-exploitation trade-off. The coefficient η controls the balance between uniform and importance-weighted sampling, and δ controls the trainable parameter ratio. The paper does not report a principled procedure for setting these values across models and tasks, which is a natural next question.
• Revisiting large embedding and LM-head modules. MISA deliberately excludes the embedding and LM-head layers during fine-tuning because their parameter counts scale with vocabulary size. Whether memory-efficient ways exist to include these high-importance modules remains open.
• Combining with orthogonal memory savers. The experiments explicitly avoid gradient checkpointing and flash attention; the paper does not report how MISA composes with quantization, checkpointing, or other memory-reduction techniques.
• Scaling beyond the tested regimes. The largest models evaluated are LLaMA3-8B and Qwen2.5-7B for fine-tuning and LLaMA2 350M for pre-training. Behavior at much larger scale, on mixture-of-experts architectures, or on longer context windows is not reported.

Target Audience

This paper is most useful to machine learning researchers and engineers working on memory-efficient training and fine-tuning of large language models, particularly those interested in block-coordinate descent, importance sampling, or convergence theory for adaptive optimizers. It is also relevant to practitioners with limited GPU resources who need to fine-tune 7B-scale models on a single consumer GPU, and to theorists studying convergence guarantees under realistic Adam-based, multi-update block optimization.

Authors’ abstract

The substantial memory demands of pre-training and fine-tuning large language models (LLMs) require memory-efficient optimization algorithms. One promising approach is layer-wise optimization, which treats each transformer block as a single layer and optimizes it sequentially, while freezing the other layers to save optimizer states and activations. Although effective, these methods ignore the varying importance of the modules within each layer, leading to suboptimal performance. Moreover, layer-wise sampling provides only limited memory savings, as at least one full layer must remain active during optimization. To overcome these limitations, we propose Module-wise Importance SAmpling (MISA), a novel method that divides each layer into smaller modules and assigns importance scores to each module. MISA uses a weighted random sampling mechanism to activate modules, provably reducing gradient variance compared to layer-wise sampling. Additionally, we establish an \(\mathcal{O}(1/\sqrt{K})\) convergence rate under non-convex and stochastic conditions, where $K$ is the total number of block updates, and provide a detailed memory analysis showcasing MISA's superiority over existing baseline methods. Experiments on diverse learning tasks validate the effectiveness of MISA. Source code is available at https://github.com/pkumelon/MISA.