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

We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.

Q1. Cleaning optimizer state.

We thank the reviewer for this very insightful question regarding our observation in Figure 7. We believe this effect is specific to the block-coordinate nature of MISA and would likely not be beneficial in standard full-model training.

Our conjecture for why clearing the optimizer state helps in MISA's pre-training is based on the concept of "stale" momentum.

• In MISA's block-coordinate framework, a module is updated and then frozen. It may not be selected for an update again for a significant number of iterations. During this long interval, the other active parts of the model change substantially, shifting the overall loss landscape.
• When the module is eventually re-activated, its stored momentum from many iterations prior is now "stale"—it was calculated based on a very "old" model state. Reintroducing this outdated momentum could provide a poor estimate of the current gradient's direction and potentially hinder optimization.
• Therefore, clearing the optimizer state acts as a momentum reset for the newly activated block. This forces the optimizer to re-evaluate the gradient direction and track it tightly, which can be beneficial for model training.

This contrasts sharply with standard full-model training (e.g., with Adam):

• In standard training, all model parameters and their corresponding optimizer states (first and second moments) are updated in every single iteration.
• The momentum is therefore always "in-time" and provides a valuable, continuously updated estimate of the gradient's trajectory.
• In this context, clearing the optimizer state would be detrimental, as it would discard this crucial historical information and disrupt the convergence path.

In summary, we believe the benefit observed in Figure 7 is a unique characteristic of block-wise update methods, where optimizer states for inactive blocks can become outdated.

Q2.1 Scaled gradient norm.

We thank the reviewer for this insightful question.

The primary motivation for our normalization—dividing the gradient norm by the square root of the number of parameters (as defined in Appendix B.2)—is to create a scale-invariant importance metric.

Without normalization, the gradient norm would be heavily biased towards parameter-dense modules (e.g., FFNs) over smaller ones (e.g., attention heads), simply because they have more parameters contributing to the norm's magnitude. Our chosen normalization effectively measures the average intensity of the gradient signal per parameter, allowing for a more equitable comparison between modules of vastly different sizes. This ensures that our importance score reflects the relative significance of the update for a given module, rather than just its absolute size.

Q2.2 Different normalizations in different layers.

We agree with the reviewer that a single, global normalization scheme may not be optimal, and we thank you for pointing us to the excellent work by Yang et al. [1]. The idea of using more sophisticated, layer-specific normalization is indeed valuable. This could involve adapting the normalization based on a module's depth, its functional role (e.g., attention vs. FFN), or other structural properties. We consider this a promising direction for future research and will cite and discuss this work in our revised paper. Thank you again for the inspiring suggestion.

Q3. Sensitivity to inner iteration $T$. We conduct an ablation study on the effect of $T$, as shown in the table below. We found that different choices of $T$ does not affect the convergence of online training loss, consistent with the observations reported in the BAdam and LISA.

Reference

[1] Yang, G., Simon, J. B., & Bernstein, J. (2024). A Spectral Condition for Feature Learning.

We sincerely thank the reviewer again for their thoughtful comments. We hope the above responses adequately address their concerns and are happy to provide further clarification if needed.

We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.

Question 1. Trade-off between $T$ and $N$ We conduct an ablation study on the effect of $T$ when $T\times N$ is fixed to 9650, as shown in the table below. We found that different choices of $T$ does not affect the convergence of online training loss, which is consistent with the observations reported in the BAdam and LISA.

Question 2. Fine-tune a small set of blocks with MISA

The reviewer is correct that MISA is a memory-efficient method for full-parameter fine-tuning. This approach aims to match or exceed the performance of standard full fine-tuning, while producing a full model rather than a small, separate adapter.

However, the core principles of MISA are flexible and can be directly applied to the "constrained scenario" the reviewer proposes. Our paper already contains two key experiments that validate the viability and power of using MISA to create or enhance parameter-efficient adapters:

1. Creating novel "module-adapters" (Appendix D.5): We explored a simplified version of the reviewer's idea in our ablation study. As shown in Figure 10, we fine-tuned only a single module type (e.g., only $W_q$ across all layers) while keeping the rest of the model frozen. The significant performance gains observed in this experiment prove that fine-tuning a small, targeted subset of modules is a viable and effective strategy. This validates the core concept of creating a "module-adapter" that would result in a small, storage-efficient "diff" file. To address this, we conducted comparative experiments under three settings: (1) MISA fine-tuning selected modules, (2) MISA fine-tuning all modules, and (3) LoRA fine-tuning all modules. The results are shown in the following table. Since fully fine-tuning all modules is not always optimal for LoRA, its reported $W_{all}$ corresponds to fine-tuning only $W_q$, $W_v$, $W_{up}$, and $W_{down}$ (which yields the best performance), while MISA's $W_{all}$ refers to fine-tuning all modules.

As shown in the table, even when only specific individual modules are fine-tuned, MISA still achieves competitive performance compared to LoRA.

2. Enhancing existing LoRA adapters with MISA (Appendix C.2): We can also apply MISA's principles directly to LoRA. In our LoRA+MISA hybrid approach, we partition the LoRA matrices into modules, and use MISA's importance sampling to activate only a fraction of them during training. As shown in Figure 6, this method achieves higher accuracy than standard LoRA while saving ~8% memory. Crucially, the final output is still a standard LoRA adapter, demonstrating that MISA can create more performant adapters with greater training efficiency.

Together, these results confirm the value of the reviewer's suggestion. The MISA framework is not only powerful for full-parameter tuning but can also be used to create novel, efficient "module-adapters" or enhance existing adapter methods.

Question 3. Variance due to sampling mechanism

We agree that the stochastic sampling mechanism will inherently introduce some variance across different training runs. To ensure the reliability of our findings, we report MISA's results on LLaMA2-7B from multiple runs with different random seeds. The table below illustrates the averaged performance and standard deviation.

To validate that the stochasticity within our proposed importance sampling is crucial to MISA's performance, we conducted a new experiment comparing our method against a deterministic, cyclic sampling baseline. The results presented in the table below show that MISA with stochastic importance sampling significantly outperforms the deterministic approach, confirming that the stochastic element is essential and cannot be removed without impairing performance. Math and commonsense reasoning tasks are conducted on LLaMA3-8B.

Question 4. Peak memory comparison between LoRA and MISA.

The difference in peak memory usage between LoRA and MISA for varying context lengths (Figure 2) stems from their distinct approaches to activation memory management. LoRA, while introducing lightweight adapters per layer, does not reduce overall activation memory; it must store full activation values for backpropagation across all layers (including activations of its specific adapters). In contrast, MISA freezes inactive modules, thereby eliminating the need to store activations associated with their gradient computations. As sequence length increases, activation memory becomes the dominant factor in total memory usage, and MISA’s ability to skip activations for frozen modules yields significantly lower memory overhead compared to LoRA. We assume the number of layers is $L$, the number of attention heads is $a$, batch size is $b$, sequence length is $s$, dimension of hidden states is $h$, LoRA rank is $r$, and the threshold of trainable parameter ratio is $\delta$. According to our detailed memory analysis in Appendix F, the memory comparison between LoRA and MISA is as follows:

As shown in the Table above, a direct comparison of the formulas reveals that for long-text tasks, MISA's activation memory is guaranteed to be smaller than that of LoRA when $\delta < 7/12$. In practical scenarios, however, the value of $\delta$ generally stays below $10\\%$, which indicates that MISA boasts significantly higher memory efficiency than LoRA in long-text tasks.

Weakness 1. MISA does not lead to an adapter

The reviewer is correct that MISA is a memory-efficient method for full-parameter fine-tuning. This approach aims to match or exceed the performance of standard full fine-tuning, while producing a full model rather than a small, separate adapter. However, the core principles of MISA are flexible and can be directly applied to the "constrained scenario" the reviewer proposes.

• When MISA is employed to fine-tune only a small subset of (or even a single) modules, it results in what we term "module-adapters." This produces an adapter with parameters comparable to LoRA, as further detailed in our response to Question 2.
• Furthermore, MISA can be used to enhance LoRA adapters. This MISA + LoRA hybrid approach achieves "LoRA-adapter" functionality while simultaneously saving memory and incurring improved performance, see our response to Question 2.

Weakness 2. Modality model

Thank you for highlighting this important point. We acknowledge that our current work focuses exclusively on text modalities and does not include experiments on other modalities. This is indeed a limitation of the present study, which we will explicitly clarify in the limitation section of the paper. We promise to explore multi-modality applications (e.g., image-text, audio-text) as part in the camera-ready revision.

Weakness 3. Non-Transformer model

Thank you for pointing out this important oversight. We will clearly state in the limitations that our method is only evaluated on Transformer-based models, and its applicability to non-Transformer architectures is not addressed. Moreover, we will highlight exploring non-Transformer frameworks as a key direction for future work.

Weakness 4. Small scale experiment

In Section 5 of the manuscript, we have acknowledged the limitations in pretraining due to resource constraints. While not at the 7B+ scale, our pre-training results (Figure 4 and Table 6) on 130M and 350M models provide valuable proof-of-concept. Furthermore, we include an additional pre-training experiment on a 1B model using 13.1B tokens to further demonstrate MISA's utility, shown in the following table. Most importantly, our work's primary focus is on memory-efficient fine-tuning, where MISA's effectiveness and scalability have been extensively validated on large models such as LLaMA3-8B, Qwen2.5-7B, Mistral-7B, and LLaMA3-70B.

We sincerely thank the reviewer again for their thoughtful comments. We hope the above responses adequately address their concerns and are happy to provide further clarification if needed.

We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.

W1. Terminology. We sincerely thank the reviewer for highlighting this point. Our definitions are:

• Layer: A standard transformer layer (e.g., the entire multi-head attention and FFN). This is the coarse-grained unit of optimization in prior work .
• Module: Our proposed fine-grained optimization unit. A module can be regarded as a component inside the standard transformer layer. For instance, within the multi-head attention mechanism, we identify four such modules: $W_q$, $W_k$, $W_v$, and $W_o$. Similarly, in the feed-forward network, we define two modules: $W_{up}$ and $W_{down}$. These internal components collectively constitute a transformer layer.
• Block: The general term for the unit of optimization in the Block Coordinate Descent (BCD) framework. a block typically refers to an entire layer. However, in our work, a block can be either a single module or a collection of several modules.

With these definitions, the writing in Section 2.3 is unambiguous, which motivates our shift from a "layer-as-a-block" paradigm to a "module-as-a-block" paradigm. Proposition 2 further demonstrates that using modules as the optimization unit is superior to using layers. In the revision, we will add explicit definitions for "layer," "block," and "module," clarifying their interrelationships and ensuring global consistency throughout the manuscript.

W2. Algorithm 1. $\tau_n$ is the block selection index for step $n$. In Step 5 of Algorithm 1, all selected modules constitute the block $\tau_n$. Steps 5-13 of Algorithm 1 then operate on the modules within the block specified by $\tau_n$. For detailed information on block selection, please refer to Algorithm 2 in Appendix B.1.

W3. Notation. We'd like to clarify that there's no inconsistency in our notation. Vectorizing module matrices into vectors follows a standard convention in optimization theory.

• While modules are intuitively described as architectural matrices (e.g., the weight matrix $W_q$ in Section 2.3), it is common practice in optimization literature to $**vectorize these parameters for mathematical analysis**$. This simplification streamlines the notation for gradients and convergence proofs. For instance, most stochastic gradient descent convergence proofs in literature utilize vector representations, even though neural network weights are typically in matrix form.

• In the specific case of Proposition 2, we employ this convention to illustrate the benefits of our fine-grained block-coordinate update. Here, $\theta_b$ represents the parameters of a single layer, which is then decomposed into its constituent modules, denoted as $\theta_b = (\theta_{b,1}, \dots, \theta_{b,K})$. Each $\theta_{b,j}$ corresponds to the vectorized form of that module's matrix parameters. The primary purpose of this proposition is to formally demonstrate that optimizing over these finer-grained module-blocks is a superior strategy compared to optimizing over the coarse layer-block.

W4.1 Comparison with LoRA. We would like to respectfully clarify that we did include LoRA and other key baselines in our instruction tuning experiments. The results of this comparison are presented in Table 5 on page 9. Our results demonstrate that MISA consistently outperforms or performs competitively against LoRA and other state-of-the-art baselines on three different models (TinyLLaMA, LLaMA2-7B, and Mistral-7B) for instruction tuning.

W4.2 Comparison with LISA. MISA outperforms LISA for two fundamental reasons.

• Module-wise partitioning. LISA employs a coarse "layer-as-a-block" approach, treating all components within a transformer layer as homogeneous. In contrast, MISA recognizes the inherent heterogeneity of these components, underscoring the necessity of a finer "module-as-a-block" approach. Proposition 2 provides the theoretical justification for this, proving that our finer-grained partitioning enables a superior optimization step compared to any layer-wise strategy.
• Adaptive importance sampling. While LISA keeps the embedding and LM-head layers active, it samples the transformer layers—which account for the majority of parameters—uniformly at random. This non-adaptive strategy can be inefficient. MISA, however, employs an adaptive importance sampling mechanism (governed by Eq.(4) in the manuscript) that dynamically prioritizes modules with higher utility. This intelligent allocation of computational budget allows MISA to focus on the most impactful modules at each step, leading to better convergence.

Q1. Proposition 2. The reviewer is correct that a probability distribution over layers, $p_b$, is not directly a probability distribution over modules, $p_{b,j}$. Our statement that "a solution to problem (2) is also a feasible solution to (5)" is meant to convey that any layer-wise sampling strategy can be represented as a valid instance of a module-wise sampling strategy. Let's make this mapping explicit:

1. A layer-wise sampling strategy, which is the solution to problem (2), gives a probability $p_b$ for sampling each layer $b$. Implicitly, when layer $b$ is chosen, all $K$ of its modules are updated.
2. We can construct an equivalent module-wise probability distribution by distributing the layer's probability uniformly across its modules. That is, we can set the module sampling probability $p_{b,j} = p_b / K$ for all modules $j=1,...,K$ within layer $b$.
3. This constructed distribution ${p_{b,j}}$ is a feasible solution for problem (5) because it satisfies the necessary constraints:
   • $p_{b,j} >= 0$ since $p_b >= 0$.
   • The probabilities sum to one: $\sum_{b=1}^B\sum_{j=1}^K p_{bj} = \sum_{b=1}^B\sum_{j=1}^K p_{b}/K = \sum_{b=1}^B p_b = 1$

The core purpose of Proposition 2 is to show that the optimization space for module-wise sampling (problem 5) is a superset of the strategies considered by layer-wise sampling (problem 2). Because the layer-wise approach represents just one possible (and likely suboptimal) feasible solution within the module-wise framework, the optimal solution found by MISA in the space of problem (5) is guaranteed to be at least as good as, or better than, the optimal layer-wise solution of problem (2).

Q2. Comparison with LISA. Our central finding is that how the transformer blocks are fine-tuned is far more critical for downstream performance than whether the embedding and LM-head layers are updated. MISA's performance advantage stems directly from its two core contributions, which are absent in LISA: Fine-Grained Module Partitioning and Adaptive Importance Sampling (see our response to W4.2). We conduct the ablation study to validate these factors. As shown in the following table, using the same granularity as LISA (layer-wise) combined with our proposed importance sampling strategy already yields higher accuracy than LISA.

Q3. Value of $\delta$. In MISA's experiments, we primarily choose the value of $\delta$ in Algorithm 1 to align with the memory usage of baselines. For instance, setting $\delta=3\\%$ ensures GPU memory consumption matches that of BAdam, while $\delta=1\\%$ achieves lower consumption than all baselines (see Tables 1, 3, 4, etc.). This choice effectively balances efficiency and performance by constraining the number of active modules within hardware memory limits. In practice, $\delta$ adapts flexibly to the GPU memory constraints of the target hardware, allowing users to adjust it based on their specific computational environment without compromising the core mechanism of importance sampling. This flexibility renders MISA robust across diverse deployment scenarios, from memory-constrained edge devices to high-end training clusters.

Q4. Data combination. The combination of training data from all eight tasks into a single set for fine-tuning follows the experimental setup of LLM-Adapters[1] and DoRA[2]. This is a standard practice in multitask instruction tuning, designed to enable multitask generalization and ensure a fair comparison with baselines like LISA, LoRA, and GaLore under consistent conditions.

Q5. Comparison with LoRA. We would like to respectfully clarify that we did include LoRA and other key baselines in our instruction tuning experiments. The results of this comparison are presented in Table 5 on page 9. Our results demonstrate that MISA consistently outperforms or performs competitively against LoRA and other state-of-the-art baselines on three different models (TinyLLaMA, LLaMA2-7B, and Mistral-7B) for instruction tuning.

Reference

[1]Hu Z, Wang L, Lan Y, et al. Llm-adapters: An adapter family for parameter-efficient fine-tuning of large language models[J]. EMNLP, 2023.

[2]Liu S Y, Wang C Y, Yin H, et al. Dora: Weight-decomposed low-rank adaptation. ICML, 2024.

We sincerely thank the reviewer again for their thoughtful comments. We hope the above responses adequately address their concerns and are happy to provide further clarification if needed.

We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.

W1. Detailed profiling. The following table shows the detailed profiling, including all unique operations. MISA-related operations are performed every $T$ steps ($T=50$ by default). The table reports the average time consumption per step (in ms), showing that the overhead introduced by MISA is negligible compared to the forward and backward consumption. The profiling is conducted using Qwen2.5-7B on a single RTX PRO 6000.

W2. Memory and compute overhead. We define the hidden dimension as $h$, sequence length as $s$, batch size as $b$, and number of layers as $L$. Each layer has 7 modules, totaling $B=7L$ modules.

• Theoretical comparison. Analyses in Appendices F.5 and G.3 show that computing and storing the smoothed historical gradient norm ($G_b$) incurs an $O(7L)$ memory overhead and $O(h^2)$ computational overhead. In comparison, the memory and computational overheads for gradients are $O(7Lh^2)$ and $O(bsh^2)$, respectively. This results in relative ratios of $O(1/h^2)$ for memory and $O(1/bs)$ for computation. As the model scales (i.e., larger $h$ and $L$), these ratios will remain stable or decrease quadratically, rendering $G_b$'s overhead negligible.
• Empirical evidence. To validate this point, we conducted additional experiments. We trained models of different sizes on the Alpaca-GPT4 dataset with a batch size of 1, under both bf16 and fp32 precision settings. As shown in the table below, the computational and memory overhead of gradient norm is trivial. Storing the gradient norms of all modules requires less than 100 Bytes.

W3. Gradient norm as module importance. Using the gradient norm to gauge module importance is not a heuristic; instead, it follows from a mathematical derivation.

• According to the derivation in Line 130, the decrease in loss function can be lower bounded by $\mathbb{E}[f(\theta^n) - f(\theta^{n+1})] \ge \sum_{b} p_b \\|g_b\\|^2.$ To maximize this loss decrease, we should maximize the right-hand side, $\sum_{b} p_b \\|g_b\\|^2$. As shown in Lines 132-134, this implies choosing the block with the largest squared gradient norm. Intuitively, this strategy prioritizes blocks with larger gradient norms, as they contribute more significantly to optimization progress.

• To better evaluate the gradient norm in the presence of stochastic noise, this paper applies an moving average to the gradient norm (as illustrated in Eq. (4) of the manuscript). This effectively mitigates the impact of noise. Additionally, the sampling probability in Eq.(4) employs a softmax operation on the importance metrics, further smoothing the noise and enhancing training stability.

W4. Module partition.

• MISA is flexible. We wish to clarify that MISA supports both coarser and finer-grained block partitioning. The core algorithm development—including the block coordinate descent update (Lines 124-130), importance sampling (Lines 131-141), practical implementation details (Lines 144-155, Algorithm 1), and convergence analysis—does not depend on a specific, concrete partitioning like $W_q, W_k, W_v$. This means our framework can operate effectively with any valid partitioning, unequivocally demonstrating that it is not hardcoded.

• Justification for the chosen partition. The partitioning into units like $W_q,W_k,W_v$​ serves as a natural illustrative example in LLMs for two key reasons: (1)These are distinct matrix parameters within a transformer layer. They are inherently associated with weight gradients that participate in forward/backward passes, making them straightforward to compute and manage. (2) As shown in Figure 1, these modules exhibit significant heterogeneity in their gradient norms, which is crucial for the effectiveness of our importance sampling strategy.

• Partition ablation. We do not claim our specific partition is optimal. It is possible that better partitions might exist, and to explore this, we conducted an ablation study on partitioning granularity. In the coarse-grained setting, each Transformer layer is divided into the Attention block and the MLP block. In the fine-grained setting, each matrix within the Transformer layer is partitioned into to the size of attention heads($128\times 4096$). Compared to the module-wise granularity, the improvement brought by finer-grained strategies becomes marginal.

Q1. Module partition. While we do not claim our specific partition is universally optimal, our primary goal is to demonstrate that this fine-grained module-wise approach is a significant and superior advancement over the layer-wise partitioning used in prior work. Proposition 2 provides theoretical evidence for this superiority , and our strong empirical results confirm it in practice. Please refer to our response to W4 for more detailed discussion and ablation study.

Q2 and Q3. Importance score update frequency. According to Line 14 in Algorithm 1, the importance score is updated every $T$ inner iterations. We have conducted experiments analyzing how the importance scores change in training. We observed that the scores of several most important modules become stable early in training, while the later modules require a period of training to stabilize. In the later stages of training, the relative importance of modules changes very little. We also observe that when $T$ is too large, the convergence speed in the early training phase slows down. Therefore, it is beneficial to update the module importance scores more frequently at the beginning of training. In practice, however, the overhead of computing importance scores is very low (see W2), so more frequent updates of sampling probabilities do not introduce significant overhead. We will include more discussion on this in the revision.

Q4. Sensitivity to gradient norm. We thank the reviewer for this insightful question.

• Ablation study on importance metrics. We investigated the sensitivity of MISA to the choice of importance metric through the following ablation study on LLaMA3-8B. Our results consistently show that the gradient norm yields the best performance when compared to the weight norm and parameter count, thus validating it as a robust choice.

• Mitigating over-sampling of high-variance modules. We acknowledge the potential issue of over-sampling modules that exhibit high gradient variance but low long-term utility. MISA is explicitly designed to counteract this through its sampling strategy. Our objective function, defined in Eq. (2) of the manuscript, incorporates a KL divergence penalty that regularizes the sampling probabilities. This term prevents the sampling distribution from deviating excessively from a uniform distribution, effectively balancing exploitation (prioritizing modules with high gradient norms) and exploration (ensuring all modules have a non-trivial probability of being sampled).

L1. Test on larger models. In Section 5 of the manuscript, we have already acknowledged the limitations due to resource constraints. While not at the 7B+ scale, our pre-training results (Figure 4 and Table 6) on 130M and 350M models provide valuable proof-of-concept. Furthermore, we include an additional pre-training experiment on a 1B model to further demonstrate MISA's utility. As demonstrated in the table below, MISA outperforms standard Adam optimization in terms of perplexity in 1B model training. Most importantly, our work's primary focus is on memory-efficient fine-tuning, where MISA's effectiveness and scalability have been extensively validated on large models such as LLaMA3-8B, Qwen2.5-7B, Mistral-7B, and LLaMA3-70B.

L2. Module partition granularity. See the granularity ablation study in our response to W4 and Q1.

L3. Detailed profiling. See our response to W1.

We sincerely thank the reviewer once again for their thoughtful comments. We hope the above responses adequately address their concerns and are happy to provide further clarification if needed.

Dear reviewer ULBr，

I hope this email finds you well.​

First and foremost, I want to express my sincere gratitude once again for the time and effort you've already invested in reviewing our paper (ID: 15661). Your initial feedback was incredibly insightful, and I truly appreciate your commitment to the review process.​

I understand that you must be extremely busy with multiple responsibilities, but I wanted to gently follow up regarding our rebuttal. We submitted it within the given timeframe, and since we haven't received any further comments or indications of remaining concerns from you, I'm writing to kindly inquire if you've had the opportunity to review it. Your perspective is of utmost importance to us, as it will help us refine our work and ensure it meets the high standards of this conference.​

If there's anything we can do to assist you in reviewing our response more efficiently, please don't hesitate to let me know. I would be truly grateful if you could spare some time to share your thoughts with us at your earliest convenience.​ Thank you again for your dedication to academic excellence and for your invaluable contributions to this review process. I look forward to hearing from you soon.

Best regards,​

The authors of paper 15661