Title: MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling
Venue: NeurIPS 2025
Accepted: yes
Keywords: large language models

memory-efficient fine-tuning

block coordinate descent

importance sampling
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 **M**odule-wise **I**mportance **SA**mpling (**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 training steps, and provide a detailed memory analysis showcasing MISA's superiority over existing baseline methods.  Experiments on diverse learning tasks validate the effectiveness of MISA.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces MISA (Module-wise Importance SAmpling) - a memory efficient Transformer optimization scheme. In the method, all model’s components (Q, K, V, O, W_up, W_down, W_gate matrices in all layers) are assigned a uniform probability of being sampled for training. After sampling, the selected components are optimized for a preset number of optimization steps, and after that, their probability of being sampled is updated based on the gradient in the component. This sample-and-train cycle occurs a preset number of times. The memory savings in the method are due to the fact that only the optimizer parameters of the currently selected blocks need to be stored in VRAM.\nThe method convergence rate is shown via theoretical analysis.\nThe method’s efficacy is measured in both fine-tuning and pretraining and against numerous baseline methods (LoRA, DoRA, LISA, BAdam, IST).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. Extensive experiments and ablations\n1. The paper is easy to read\n\nWeaknesses:\n1. The proposed method (unlike e.g. LoRA) does not save storage space needed for the storing the diff\n1. No experiments for modalities different than text\n1. No mentioning of non-Transformer architectures (even in Limitations)\n1. Very small scale (especially the number of tokens) in the pretraining experiments\n\nQUESTIONS:\n1. What is the tradeoff between T and N in Algorithm 1, i.e. if the total number of steps (T*N) is fixed, how could one approach choosing the appropriate T and N?\n1. Have the authors considered a constrained scenario in which only some of the blocks are modified, resulting (similarly to LoRA) with a diff that is smaller than the whole model?\n1. Have the authors explored the variance in the results inherent due to the sampling mechanism used? Can this stochastic element be eliminated from the algorithm?\n1. Why do LoRA and MISA differ in peak memory usage for different context lengths (Figure 2).\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

comments: SUMMARY:
The authors propose a new efficient optimization method based on module-level importance sampling. The key idea is to update each module probabilistically based on the computation budget and the importance score which comes from the gradient norm. The authors demonstrate empirically that models can be trained with a fraction of the memory and computation of full model training while achieving competitive performance. A theoretical convergence guarantee is also supplied.

STRENGTHS AND WEAKNESSES:
The paper is very clear with extensive experiments and ablations. The idea is simple and appears to be scalable. A detailed analysis of the computational and memory requirements are provided with comparisons to many popular baselines like LoRA.

QUESTIONS:
I found the figure 7 intriguing. Is a similar observation true in full model training, namely that cleaning optimizer state improves performance early in training?

Also is there some intuition for why gradient norm should be divided by the square-root of the number of parameters? It might make sense to have different normalizations in different layers, as per [1].

What is the impact of the number of inner iterations $T$? How sensitive is this hyperparameter and what is the rule of thumb for setting it in practice?

1. A Spectral Condition for Feature Learning (2024) by Yang et al.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 3
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn the proposed method, score based importance sampling is utilized for selecting modules for pretraining and fine-tuning. The authors provide convergence rate under the nonconvex stochastic optimization setting.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nModule-wise optimization is motivated by demonstrating that smaller modules of parameters preserve more gradient information. Moreover, this strategy eliminates the need for storing the parameters of an entire layer. Furthermore, the derive the sampling probabilities for the modules as a function of the gradient norm and the convergence rate.\n\nWeaknesses\n\n1. Although the idea of module based optimization is interesting, the presentation lacks clarity. The authors keeping switching between the use of terms like layers, blocks, and modules which makes it confusing to follow. For example, the description in Section 2.3 uses the term layers but Proposition 2 uses modules in the blocks. Moreover, the use of the terms layers and blocks are changed across sections without consistency. This dilutes the impact of the main contributions of the paper.\n\n2. Algorithm 1 is also confusing because as Steps 5-13 are for the modules within a block but which block is being selected is not shown. Steps 5-13 should also depend on the block index.\n\n3. There is inconsistency in the notations as well. For examples, in Sec 2.3, the authors mention that each layer has distinct internal modules where each module is a matrix parameter. However, in Proposition 2, the authors use $\\theta_b=(\\theta_{b,1},\\theta_{b,2},\\ldots,\\theta_{b,K})$, that encompasses the module parameters, to denote the parameters of a block (assuming this is same as a layer) which is a vector. \n\n4. The paper lacks adequate experiments. For example, techniques like LoRA should also be considered for instruction tuning experiments. Moreover, in several experiments, it is not clear why MISA performs better than LISA in terms of the accuracy.\n\nQUESTIONS:\n1. In Proposition 2, the authors state that a solution to problem (2) is also a feasible solution to (5). Based on my understanding, the solution to the problem (2) would give the probabilities for the block. However, the solution to problem (5) would give probabilities for the modules. Therefore, solution to problem (2) is not sufficient to solve problem (5).\n\n2. The authors also state that \"LISA finds the embedding and LM-head layers to be very important, MISA does not train them in fine-tuning tasks because their parameters are too large.\" However, I observe that in many of the simulation results, MISA is performing better than LISA, why is that the case?\n\n3. How are you deciding the value for delta in Algorithm 1?\n\n4. In Experiments Section, the authors mention that \"we combined the training data from all eight tasks into a single training set for fine-tuning.\" Why was this combination needed in this paper?\n\n5. In Section 4.3, the authors should also include results with respect to techniques like LoRA to show the effectiveness of the proposed approach.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes MISA (Module-wise Importance Sampling Algorithm), a novel and memory-efficient optimization method designed for full-parameter fine-tuning and pre-training of large language models (LLMs). Unlike prior layer-wise or low-rank adaptation approaches such as LoRA, MISA decomposes each transformer layer into smaller modules (e.g., attention heads, feed-forward components). It then computes gradient-based importance scores and samples these modules for update using a softmax-based strategy that balances importance and exploration. The method is theoretically grounded with convergence guarantees under non-convex stochastic settings and is supported by strong empirical performance across commonsense reasoning, mathematical problem solving, instruction tuning, and pre-training tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n* The motivation is well-articulated. The paper clearly identifies memory bottlenecks in full fine-tuning and highlights the limitations of existing layer-wise or low-rank methods.\n* The idea of fine-grained module-level optimization is compelling. Shifting from coarse layer updates to smaller modules allows for better control over memory and seems to preserve more task-relevant parameters.\n* Theoretical contributions are strong. The authors provide a convergence guarantee even when using stochastic gradients, Adam, and repeated updates per module—a realistic and practically relevant setting.\n\nWeaknesses:\n\n* While Section C.3 provides per-step timing comparisons (Table 8) and shows that MISA is more computationally efficient than LoRA and significantly faster than GaLore, the breakdown does not include the cost of repeated module updates ($T$ inner steps) or the overhead from adaptive sampling. While overall step time is reasonable, it remains unclear how much of MISA’s time is spent specifically on its unique operations (e.g., importance computation, sampling, module-specific Adam updates).\n* MISA requires maintaining a smoothed historical gradient norm ($\\mathbf{G}_b$) for every module, which introduces additional memory and compute overhead—especially as the number of modules increases. Although the authors state that this overhead is negligible, they do not provide a quantitative breakdown (in memory footprint or compute usage), which is particularly important when scaling to larger models or when operating in mixed-precision, low-memory environments.\n* The use of gradient norm as a proxy for module importance is somewhat heuristic. In practice, gradient magnitude may not always reflect the true contribution of a module to generalization, especially in overparameterized settings. This could lead to prioritizing noisy or spurious updates.\n* The partitioning into modules (e.g., Wq, Wk, Wv, etc.) is static and hardcoded. There is no discussion or ablation on whether this granularity is optimal, nor is there any mechanism to adapt it dynamically based on model behavior during training.\n\nQUESTIONS:\n* In Algorithm 1 (Line 1) you state: \"Partition the model into B modules (not layers)\". However Section 3.1 briefly mentions using subcomponents such as Wq, Wk, Wv, Wo, Wup, and Wdown. However, the paper does not justify why this particular partitioning is optimal. Could you clarify how this module granularity was chosen, and whether alternative decompositions (e.g., more fine-grained or head-wise) were explored?\n* How often do the module importance scores (i.e., gradient norms) change during training? Are the high-importance modules relatively stable, or do they vary significantly over time? A clearer understanding here would help assess whether recomputing them at every outer loop is computationally justified.\n* How often do the module importance scores (i.e., gradient norms) change during training? Are the high-importance modules relatively stable, or do they vary significantly over time? A clearer understanding here would help assess whether recomputing them at every outer loop is computationally justified.\n* Have you investigated how sensitive MISA is to the choice of gradient norm as the importance metric?\nSpecifically, could modules with consistently high variance but low long-term utility be over-sampled?\n\nLIMITATIONS:\n* The method has only been tested on small-scale pretraining (130M and 350M models). It remains unclear how it performs during full-scale pretraining on 7B or 70B models, which is where memory constraints are most severe.\n* The fixed module granularity may restrict adaptability across architectures or training regimes.\n* While the paper provides per-step timing comparisons (Table 8), it does not isolate the time spent on MISA-specific operations (e.g., sampling, importance score computation, and multi-step inner updates). A more detailed breakdown of these components would help fully assess where MISA’s computational efficiency gains originate, especially as model size scales.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}