MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling
Authors: Yuxi Liu
Renjia Deng
Yutong He
Xue Wang
Tao Yao
Kun Yuan
Venue: NeurIPS 2025 poster ()
Keywords: large language models
memory-efficient fine-tuning
block coordinate descent
importance sampling
Primary area: optimization

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.

MISA: Memory-Eficient LLMs Optimization with Module-wise Importance Sampling
Yuxi Liu" Peking University yuxiliu666dstu, pku.edu.ccm edu.
Renjia Deng"
Peking University 25012100780stu, pku. edu.
Yutong He Peking University yutonghe@pku. edu. cn
Xue Wang Alibaba DAMO Academy
xue www.waawaaaaaaaaaaaaaa com
Tao Yao
Kun Yuan
Shanghai Jiao Tong University taoyao tu. edu. cn
Peking University
kunyuan@pku. edu. cn
Abstract
The substantial memory demands of pre-training and fine-tuning large language models (LLMs) require memory-effiiien 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 suboptimal per- formance. 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), 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 (1/ K) convergence rate under non-convex and stochastic conditions, where K is the total number of block updates, and provide detailed memory analysis showcasing MISA's superiority over existing baseline methods. Experiments diverse learning tasks validate the effectiveness of MISA. Source code is available at: http://gtuub..oo///keeoo////////
Introduction
Large language models (LLMs) have emerged a cornerstone of modern artificial intelligence, driving groundbreaking advancements in diverse domains such as machine translation, commonsense reasoning, and mathematical problem-solving. Despite their impressive capabilities, fine-tuning these models for downstream tasks introduces significant challenges [45, 46, 49], primarily due to the substantial memory overhead required to store optimizer states, gradients, and intermediate activations. These memory footprints where the full-parameter fine-tuning of billion-scale models often exceeds the available hardware capacity, thereby impeding the practical deployment and scalability of LLMs [69, 60, 14].
particularly acute in resource-constrained environments,
"Equal contribution. Corresponniig author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

To address these challenges, parameter-effiiint fine-tuning (PEFT) methods, such as Low-Rank Adaptation (LoRA) [24] and its variants [35, 12], have gained considerable attention. These methods freeze the majority of pre-trained parameters and optimize only small, low-rank matrices integrated into transformer layers, leading to substantial reductions in memory usage. However, while these approaches are highly memory-effciiet, they inherently constrain the model's adaptability. By optimizing only a sparse subset of parameters, such methods often lead to suboptimal performance compared to full-parameter fine-tuning, as task-specific features crucial to downstream performance may reside within the frozen portions of the network [33, 66].
Recent research has explored block-coordinate descent (BCD) [55, 59, 62], a classical strategy for high-dimensional optimization, as a promising alternative for fine-tuning LLMs, resulting in the fayer-wise optimization methods such as BAdam [37], HIFT [36], LIFT [75], LISA [41], and BlockL.I. [47]. Unlike LoRA, these layer-wise methods optimize transformer blocks sequentially while keeping the remaining layers frozen. By iteratively updating all transformer blocks, layer-wise optimization effectively enables full-parameter updates, thereby preserving the expressive capacity of the original model. Empirical studies demonstrate that layer-wise optimization consistently outperforms LoRA in performance [37, 41J. Furthermore, by skipping gradient computations for frozen layers, layer-wise optimization eliminates the need to store their intermediate activations, offering greater memory efficiency than LoRA.
Motivating questions. All existing layer-wise LLM approaches [37, 36, 75, 41, 47] adopt the traditional BCD optimization paradigm, in which the model weights are divided into manageable blocks and updated iteratively. Although BCD) provides structured framework for handling the high dimensionality of LLM weights, fully realizing its potential in practical L.MM fine-tuning requires addressing the following fundamental open questions.
Q1. How to effectively partition the LLM's weights into blocks for optimization?
The choice of partitioning strategy can significantly impact optimization efficiency and convergence. Current methods [37, 36, 75, 41, 47] treat transformer layers as homogeneous units, overlooking the differing significance of internal modules-sii as multi-head attention, feed-forward networks, and normalization
each layer. By uniformly updating all parameters a layer, these methods risk over-adapting less impactful modules while under under-training critical ones, resulting in suboptimal optimization performance. Thus, more effective partitioning strategies must be explored.
Q2. How to effectively sample each block to achieve fast empirical convergence?
The performance of layer-wise methods is highly influenced by the block sampling strategies. However, current approaches primarily rely eyclic [37] or uniform [36] sampling patterns, which overlook the varying importance across layers. LISA [41] keeps the Embedding and L.M head layers active but assigns equal sampling probabilities to all transformer layers, ignoring their varying significance. BlockLLLM [47) prioritize more frequent updates to critical layers but ultimately focus on small set of fixed blocks, failing to adequately explore the remaining layers. An effective sampling strategy should maintain a balance between thorough exploration of the parameter space and efficient exploitation of the most promising optimization directions ensure rapid convergence.
Q3. How establish convergence guarantees under practical LLM settings, incorporating Adam optimizer, stochastic gradients, and multiple updates per sampled block?
Theoretical convergence of layer-wise LLM optimization remains understudied. Traditional BCD optimization literature typically focuses on block-wise GD/SGD convergence [59, 61, 22]. Recent works [37, 74, 41] provide convergence guarantees under restrictive assumptions, such as noiseless gradients or single updates per sampled block. These analyses fail to reflect practical layer-wise I.LM optimization, where Adam optimizer, stochastic gradients, and multiple updates per block are standard practice. A rigorous convergence analysis under these realistic conditions essential for broader adoption of layer-wise optimization methods.
Main contributions. To address the aforementioned open questions, we propose a novel Module- wise Importance SAmpling (MISA) method. MISA partitions the LLM's weights into smaller modules, which serve as blocks for optimization. Additionally, we assign importance scores to each module and develop an effective strategy that dynamically samples modules based on real-time i importance metrics. Our main contributions are summarized as follows:
CI. Module-wise optimization. We define a module as a matrix parameter within a transformer layer associated with weight gradients. Empirically, we observe that internal modules within

---

transformer layers exhibit heterogeneous importance. Theoretically, we demonstrate that decomposing each layer into smaller modules preserves more information in gradient, thus motivating module-wise optimization. This addresses Question 1. Furthermore, we find this fine-graines update strategy eliminates the need load full layer into memory, making more memory-efficient than layer-wise LLM optimization approaches [37, 36, 75, 41, 47].
C2. Improved importance sampling. Traditional sampling strategies often rely on heuristics, which can be suboptimal for LLM optimization. To improve this, we parameterize the gradient variance function of the sampling probability for each module and maximize it to optimize the sampling strategy. Additionally, we introduce a strategy to balance importance and uniform distributions, ensuring both comprehensive exploration of all modules and efficient exploitation of critical ones, addressing Question 2.
C3. Convergence guarantees.
demonstrate that MISA achieves a convergence rate of O(I/VK) K), where K represents the total number of block updates. Our theoretical guaran- tees derived under practical LLMs training scenarios, incorporating Adam optimization, stochastic gradients, and multiple updates per sampled block, thus addressing Question 3. Conventional BCD analysis relies the assumption that block gradients unbiased estimators of the full gradient-an an assumption that fails when performing multiple up- dates within the same block. Our new analysis addresses this limitation by establishing fundamental connections between block-level gradients and the full gradient.
Experimental results. MISA demonstrates strong empirical performance.
E1. Fine-tuning. We evaluated MISA on different LLMs across three benchmarks: Common- sense Reasoning [25], Math Reasoning [25], and Instruction Following, encompassing total of 16 datasets. We compared MISA against PEFT and layer-wise optimization methods, including LoRA [24], DoRA [35], BAdam [37], and LISA [41]. Under comparable memory constraints, MISA outperformed all baselines. An illustration of MISA's superiority in fine-tuning tasks is provided Table 1, with more detailed results presented in Section 4.
E2. Pre-training. We trained the LLaMA2 130M and 350M variant [32] on the C4 dataset [46]. In 350M model training, MISA achieved a perplexity of 22.11 after 2.7B training tokens, significantly outperforming GaLore's 24.34 [71]. and approaching Adam's 's 21.3.
Table 1: Comparison of fine-tuning methods on LLaMA3-8B across eight commonsense reasoning tasks. The "CharGPT" row presents results from ChatGPT obtained using Zero-shot CoT [58] with the GPT-3.5-turbo APL "Hella." " refers the HellaSwag datasct [65], and "Wino." " refers the Winogrande dataset [50]. Memory usage is reported without the application additional memory-saving techniques such as gradient checkpointing or flash attention. The symbol denotes the proportion of parameters updated in each training iteration.
Model
Method
Mem.(GB)
BoolQ
PIQA
SIQA
Hella.
Wino.
ARC-
ARC-c
OBQA
Avg?
CharGPT
73.1
85.4
68.5
78.5
66.1
89.8
79.9
74.8
77.0
FT
150.5
75.1
89.2
80.4
96.2
88.3
92.4
82.5
89.8
86.7
LoRA
35.7
70.8
85.2
79.7
92.5
84.9
88.9
78.7
84.4
74.6
89.3
95.5
LLAMA3-8B
LISA
56.3
74.6
88.1
81.5
96
86.4
92.5
81.7
86.2
85.9
BAdam
34.1
74.2
87.1
81.6
95
84.6
91.2
79.8
84,8
84.8
MISA(5 = 1%)
30.7
74.1
88.6
80.7
95
85.6
92
81
86.2
85.4
MISA(Ó = 3%)
34.4
75.4
90.5
81.4
95.9
88.2
92.2
82.4
86.6
2 Module-wise Importance Sampling
2.1 Problem Formulation
We consider the following problem:
min Eg(F((:E)] BERd
where F(0:C) {) is loss function that depends random variable ç, and 9 represents the set of d trainable parameters. Conventional algorithms update all elements simultaneously. leading to high memory and computational costs that hinder the training of LLMS on low-end hardware. In this work, we propose an alternative formulation the problem in (1). Specifically, we introduce block-wise representation of 0, defined as 0 = (0,, 02.00..... .... 0g), where O, E gd represents the b-th

---

block of weights. Here, we have d d -b= de. The resulting problem formulation is then given by:
min I (01, 01....0z
=- [ (e:
2.2 Block Sampling Strategy
At each time step, block-wise LLM training updates only a subset of parameters {0.). This approach inherently reduces the number of trainable parameters, resulting lower memory and computational requirements. However, it also introduces greater variability in the gradient's unbiased estima- tor, which can potentially hinder convergence performance. This section will derive effective importance sampling strategy to enhance block-coordinate optimization.
Block coordinate descent. Let g. = [V f("))l Rd denote the b-th block gradient of V f@"), where E [B] is the block index and n 2 0 is the iteration index. Additionally, we define the matrix Ub := [0;-.. .... -b: :0 E Rdxds where the b-th block of U, is the identity matrix L, E R& xdb and all other blocks zero. Suppose each block gradient g is sampled with probability ps. the BCD algorithm will iterate as follows
g=+1 t aUss.' where block is sampled with probability ps.
where is the learning rate. Assume f(0) is L-smooth, have
La²
f(0n+1)= f(0 aUigh) $ f(0") a(V)(("). Ubg
2 |1bghh
La
(b)
f(0") a(1
|1
S f(0") - a 1/29
2
where (a) holds because of Ug f f(")) and (b) holds when a S 1/LL. Taking expectation on b,
Elʃ(0") f(on+))]
a 2
1ol1,811
b=
Importance sampling. An ideal update g should maximize the expected decrease E[f( f(("+))), } ensuring the most significant descent in the objective function. This can be achieved by
B
S.I. Σm 1m=1, =1 pb 2 O.
max (p.18
Intuitively, this strategy prioritizes blocks with larger expected gradient norms, as they contribute more significantly to optimization progress. However, excessively prioritizing important blocks can lead to scenario where only a few blocks are updated frequently. This results in over-exploitation at the expense of exploration. To encourage broader exploration, we constrain the sampling probabilities to stay close to a uniform distribution by incorporating a Kullback-Leible (KL) divergence penalty:
B
B
max L.
(I/ Σ Pollo" (0.l:(l)--(//)KKT.L (B), s.t. b=1
Σp Pb=1, =1 Pb 2
(2)
{
d=1
where (n 1/B denotes the uniform distribution, and 0 is coefficient controlling the trade-off between exploration and exploitation. As 1 - 0, the KI. divergence penalty dominates, and each ps approaches uniform sampling; + +o0, the penalty vanishes, recovering standard importance sampling. The following proposition provides a closed-form solution to problem (2).
Proposition 1. The optimal solution to problem (2) is given as follows
exp (1ll,7 122)
Pb
Eb= exp (nll
which is the sampling probability of block b at iteration
Practical implementation. The sampling probability in (3) is impractical to implement, as gE [V [0")) represents a full-batch block gradient, which is inaccessible during LLM optimization. In block-wise LLM training, each sampled block is typically updated for T steps using the Adam optimizer before switching to another block. Let 9r = [VF(00.) cn.)], denote the b-th block stochastic gradient at outer iteration and inner update t. We approximate the full-batch gradient

---

norm |1% using the empirical average 1
T 1
11 resulting in a practical sampling strategy:
pb exp (TGR)
+ (1-18) Gb
3G
Σ 11, 112 If b is sampled; (4)
where GB
bb=1 exp (mG
otherwise.
Instead of relying solely on the most recent T mini-batch block stochastic gradients, GB aggregates all historical stochastic gradients to approximate the full-batch block gradient norm g in (4).
To eliminate the impact of differences in the number of parameters across blocks on gradient norm calculation, we actually use the scaled gradient norm in practice rather than the original gradient norm, and the specific definition is provided in Appendix B.2. We thus address Question by answering how to effectively sample blocks.
Remark 1 (MEMORY AND COMPUTATION OVERIEED)). The computation and storage cost to maintain {GT B 6- negligible compared to that of the block stochastic gradients. See Sections G.3 and F.5 for details.
Remark 2 (CLARIFICATION ON "LAYER". "MODULE", AND "BLOCK"), A layer is standard transformer component (e.g,, with MHA and FFN), module a fine-grained subcomponent within layer (e.g.. Wq. Ws. W.), and a block is flexible optimization unit in block coordinate descent that can be a layer, a module, or a group of modules.
2.3 Partition the Weights into Fine-Grained Modules Instead of Coarse Layers
Current approaches [37, 36, 75, 41, 47] often treat transformer layers as homogeneous structures, ne- glecting the distinct roles and contributions of their internal components-such
multi-head attention mechanisms, feed-forward networks, and normal- ization layers. As shown in Fig. 1, different mod- ules within the same transformer layer exhibit signifi- cantly varying gradient norms, underscoring their dif- fering levels of importance. Uniformly updating all modules within layer risks over-adapting less crit- ical components while under-training more essential ones, ultimately leading suboptimal performance.
Lager
Luper
0.0
0.0
W,
W/.
18c
W.
Wayt
Wap
Wiow
Figure 1: The gradient norm of different modules in different layers when fine-tuning LLaMA3-8B.
Motivated by this observation, we propose decomposing layers into distinct internal modules. this work, module is defined as matrix parameter within a transformer layer that is associated with weight gradient. For example, in the multi-head attention mechanism, we identify four modules-W, Wx. W.. and W.-while in the ferd-forward network, we define two modules-Wi up and Wdown. In other words, we treat each module as a sampling block, addressing Question 2. Proposition characterizes the advantages of module over layer-wise block sampling.
Proposition 2. Suppose Os in problem (2) contains K modules such that 0g = ( 0b.2,.-. ,0, which associates with block gradient 96 = (9b.1, 9b.2...... -b. We further introduce Paj as the sampling probability of the block gradient 9b.j- Under this setting, problem (2) transforms into
B K
B K
max {
ΣΣ D 1=1
Poj (p/
(Pb). s.t.
Poj =1 ΣΣ b=1
=oj,,
Pbj 0, (5)
b
Here, qBK I/BK) is the uniform distribution. Any layer-wise importance sampling (i.e., solution to problem (2)) is also a feasible solution to (5). Consequenily, module-wise importance sampling (i.e., the optimal solution to problem (5)) will yield larger objective function value than layer-wise importance sampling, making it a superior strategy to layer-wise sampling.
2.4 Module-wise Importance Sampling (MISA) Method
Based on the discussions in Sec. 2.2 and 2.3, we now introduce Module-wise Importance SAmpling (MISA). The pseudocode for MISA is provided in Algorithm 1. MISA is a double-loop algorithm: the outer loop, indexed by n, selects module, while the inner loop (Lines 7-12) updates the parameters of the selected module 7 times using the Adam optimizer, with each update indexed by (. As shown in Line 1, MISA partitions the model into modules rather than layers. Lines 14-15 illustrate that MISA follows the effective sampling strategy described in Proposition 1. By adjusting 77, we can
5

---

balance exploration across all modules with exploitation of the more important ones. In Line 16, the module weight 6m.7 undergoes an additional Adam step obtain 0 which facilitates the convergence analysis. In Line 17, we clear the optimizer states to ensure consistent memory efficiency throughout the training process.
Algorithm 1 Module-wise Importance Sampling (MISA)
Require: NN,,T, B, a, 8 > 0, ,1, E (0, 1) and the sampling block at outer loop
1: Partition the model into B modules (not layers);
b Partition weights into modules;
2: Initialize probability weights pl = ( ... B B):
3: Initialize the module gradient estimate G O for b [B] and let G (G.."..,,")
4: for = 1, ..., do
5: Sample L, modules (labeled with index T,,) according to P such that the ratio of trainable parameter less than 6. (See Algorithm for more details);
D Importance sampling;
6: Initialize m = 0, v.
7: for 1... do
8:
Sample batch of data and calculate block stochastic gradient 9²t for selected module Thi Update the corresponding first-order and second-orde momentum as follows:
9:
10:
m. + 31 3mm + (1 31)9%. w +- B2 + (1 - B2) (o²;0)²
Update the corresponding module as follows:
12:
0n. - --0.t-1-
Tm amm. /ttp/((// UTa +
13:
end for
Update GØ for each b [B] according to (4);
Track block gradient norm;
15:
Update Pb a+1 -
explnGE)
for each b E [B]]
Update sampling probability;
Σ/=1 exp(nG7)
gu+1,0 7m
Gz.T -a I-
81
m
Veni tc
17:
Jr.: T.
T UTa t None
Clear optimizer states;
18:
end for
19: Return AN .NNN
Table Comparison of existing memory-efficient optimization methods of LLMs.
LoRA[24]
IST[64]
GaLore[71]
OwLore[30]
BAdam[37]
LISA[41]
MISA
Full-rank Update
Importance-awar
Fine-grained memory
/
X
X
/
Gradient acocumulation
/
X
X
/
7
w/o SVD
/
X
X
/
/
Fine-Tuning
/
Pre-Traine vee-Trainng
Convergence guarantee
LISA's importance-aware strategy focuses only the emhedding layer and L.-heea layer, while transformer layers which account for the majority of model parameters still use uniform sampling. @: GaLore's 's convergence proof[71] ]] holds only when the gradient cxhibits a highly specialized structure, and BAdam's analysis(37) applies exclusively to the full-gradient regime. Their settings in the theoretical guarantees rely on restrictive assumptions.
@: MISA's framework is based on the standard assumptions unalyzing non-convex optimization problems, making more relevant to practical applications.
Table 2 provides a detailed comparison of existing optimization methods of LLMs. Notably, although 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 (proportional to vocabulary size), and training them would lead to significant increase in memory consumption.
2.5 Memory Analysis

---

We perform detailed memory analysis of MISA and present a comparative overview of memory consump- tion between block-wise optimization and subspace optimization methods, such as LoRA and Galore. Our findings demonstrate MISA's significant poten- tial for long-sequence fine-tuning tasks. As illus- trated in Fig. 2, for the LLaMA3-8B model, MISA significantly outperforms LoRA terms of memory efficiency when the sequence length becomes suffi ciently large. Notably, when r = 16, the trainable parameters in LoRA constitute only about 0.5% of the total parameters. In Appendix F, we demonstrate that (1) Layer-wise method acheives lower peak mem- ory than subspace method such as LoRA and Gal.ore in the long-sequence setting. (2) Layer-wise wise method can update more weights than LoRA in memory-constraine scenarios. (3) MISA achieves greater memory efficiency compared to the layer-wise method when the sampling ratio threshold ó is small.
32.5
(8-138)
MISA (6-55))
30.0
LoRA (=16)
27.5
25.0
24
GB
22.5
20.0
17.5
500
Soquence Length
Figure 2:
Comparison of peak memory con- sumption using MISA and LoRA tie-ttning on LLaMA3 8B across various sequence length.
Convergence Analysis
This section presents the convergence analysis for the proposed MISA method.
Limitations in existing analysis. Conventional BCD methods focus on the convergence of block- wise GD or SGD [59, 61, 22], but they do not generalize the Adam optimizer. More recent studies [37, 74, 41] establish convergence results for layer-wise Adam, but only under restrictive conditions- such as assuming noiseless gradients or allowing just single update per sampled block. These assumptions, however, diverge significantly from practical LLM training, where Adam is used with stochastic gradients and multiple updates are performed for each sampled block.
New challenges. Conventional BCD methods rely on the assumption that block gradients are unbiased estimators of the full gradient, assumption that holds only when single update is performed per sampled block-duu to the preservation of unbiasedness through random block selection. However, our repeated block updating strategy introduces two key challenges: (1) Biased Gradient Estimates. Performing multiple updates on the same block breaks the randomness of block selection assumed in conventional BCD, resulting in biased estimates of the full gradient. (2) Bias-Noise Interplay. The introduced gradient bias interacts with inherent gradient noise, potentially amplifying estimation errors that would otherwise be mitigated by unbiased estimators.
Analysis highlights. Our work tackles these challenges via three innovations: (1) Bias Propagation Analysis. We derive novel recursive relations (Corollary that characterize how gradient bias and stochastic noise accumulate over successive block updates. (2) Bridging Block and Full Gradients. We establish fundamental connections between block-level gradients and the full gradient (Corollary 3), despite the loss of unbiasedness. The additional Adam step in Line 15 of Algorithm plays a critical role ensuring a smooth transition from local block updates the optimization of global variables. (3) Convergence Framework. We develop new analytical tools to jointly control gradient bias and amplified stochastic noise, resulting in rigorous convergence guarantees (Theorem 1).
Main results. We theoretically justify the convergence of MISA as follows.
Theorem 1 (Informal). Assume the loss function is L-smooth, the stochastic gradient is unbiased and bounded, the gradient variance is upper bounded (see Appendix E for detailed assumptions). With appropriate learning rate ar, MISA achieves the following convergence rale (proof is in Appendix E):
Σ n=1
N
BJJ E(1
(1
1
N
/NT
N
where and denote the numbers of outer and inner-loop iterations, respectively.
Remark According to Theorem 1, MISA converges for any number T of inner local block updates, which constitutes a key advantage of our analysis. contrast, the analyses in [59, 61, 74] are limited to the case T = 1, while the analysis in [48, 5] requires T to be sufficiently large to ensure each block is adequately updated.
7

---

Experiments
This section evaluates the performance MISA across various fine-tuning and pre-training tasks. All our experiments are conducted on RTX 4090 24GB. We also conducted detailed ablation experiments, as presented in Appendix
4.1 Commonsense Reasoning
Settings. To validate the efficiency of MISA, we fine-tuned the LLaMA3-8B [14] and Qwen2.5-7B [63] models on the Commonsense Reasoning tasks. The dataset consists eight subsets: BoolQ [6], PIQA [3], SIQA [51], HellaSwag [65], WinoGrande [50], ARC [7], and OBQA [40]. Following the settings in [25], we combined the training data from all eight tasks into single training set for fine-tuning and then evaluated each model separately on the eight test sets. The baselines included
LoRA [24]. DoRA [35], and BCD-based methods, namely LISA [41] and BAdam [37]. We also
reported GPU memory consumption for each method to ensure a thorough comparison. No additional memory-saving techniques, such as checkpointing [4] or flash attention [9], were employed.
Results. Table presents the results for the commonsense tasks. Notably, DoRA's additional memory consumption arises from activations, while LISA incurs extra memory usage from fine-tuning both the embedding layer and the language modeling head layer. MISA outperforms all baseline methods on both the LLaMA-3-8B and Qwen2.5-7B models. Furthermore, with the trainable parameter ratio set to 1%, MISA achieves performance comparable to LoRA while saving approximately 10% of memory. In particular, MISA demonstrates significant improvements over BAdam and LISA in the Qwen2.5-7B model.
Table 3: This table compares fine-tuning methods Qwen2.5-7B models across eight commonsense reasoning tasks. "Hella." refers to HellaSwag [65], and "Wino." " to Winogrande [50]. The first row shows ChatGPT results from Zero-shot CoT [58] using the GPT-3.5-turbo API. Refer to Table for results of LLaMA3-SB.
Model
Method
Mcm.(GB)
BoolQ
PIQA
Hella.
Wino.
ARC-e
ARC-c
OBQA
Avg.f
FT
140.5
75.3
91
81.4
96.2
91.5
96.1
89.1
93.4
89.3
LoRA
34.2
75.3
90.1
80.8
96.2
86.7
96.3
91.4
88.2
Qwen2.5-7B
DoRA
55.6
74.9
90.4
81.2
96.1
87.1
96.1
88.9
91.9
88.3
LISA
$7.9
75
90.3
80.6
95.2
86.2
95.8
89.1
89.8
87.8
BAdam
33.7
74.5
90.3
95.5
86.1
96.1
89.6
88.1
MISA(5 = 1%)
30.3
74
90.1
82.4
95.5
87.1
96.1
89.3
90,2
88.1
MISA(6 3%)
34
75.4
90.5
81.6
95.7
87.2
96.3
89.5
88.5
4.2 Math Reasoning
Settings. To evaluate MISA on math reasoning tasks, we tested LLaMA3-8B and Qwen2.5-7B on four arithmetic benchmarks: GSM8K [8], SVAMP [42], AQUA [34], and MAWPS [28], following [25] for dataset settings. The models were fine-tuned on MATHIOK [25], which combines training data from AQUA, GSM8K, and MAWPS, incorporating LM- -enerattdd saain-of-thoughte [58] steps.
Results. Table 4 presents the performance of various fine-tuning methods on LLaMA3-8B and Qwen2.5-7B. MISA outperforms all baselines both models. Notably, MISA achieves signif- icant performance boost over LoRA and BAdam, which have similar GPU memory consumption. Compared to DoRA, MISA attains comparable performance with greater memory efficiency.
Table 4: Comparison fine-tuning methods on LLaMA3-8B and Qwen2.5-7B across four arithmetic reasoning tasks. The left table presents results for LLaMA3-888 8B, while the right table displays results for Qwen2.5-7B. The unit of memory consumption is GB.
Method
Mem. GSMSK SVAMP AQuA MAWPS Avg. Mem.
GSMSK SVAMP AQuA MAWPS Avg.
LLaMA3-SB
Qwen2.5-7B
LORA
35.6
70.1
78.1
45.7
93.7
71.9
34.3
80.6
86.6
65.7
92.4
81.3
DoRA
54
70.9
78.4
50.3
93.4
73.3
54.8
80.1
86.7
68.1
92.5
81.9
LISA
54,2
68
76.7
49.7
88.3
70.7
55.6
77.6
86.8
70.1
89.8
81.1
BAdam
33.5
69.5
77.6
47.6
93.3
72
33.6
80.2
85.9
65.1
92.4
80.9
MISA
1%)
30.4
68.5
77.3
49,3
90.1
71.3
29.8
77.5
85
65
89.5
79.3
MISA(6 3%)
34.1
71.3
78.5
51.2
93.3
73.6
33.6
81
88.1
66.1
92.9
82
8

---

4.3 Instruction Tuning
To demonstrate MISA's efficiency in instruction-followiin fine-tuning, we fine-tuned TinyLLaMA667], LLaMA2-7B(54) and Mistral-7B[26) on the Alpaca GPT-4 dataset, which consists of 52K samples generated by GPT-4 based on imputs from the Alpaca dataset [53]. Model performance was evaluated on MMLU[21], MMI.U-pro(57) and MT-Beecch [72], As shown in Table 5, MISA outperforms all baselines in most evaluation benchmarks, demonstrating its robustness and efficiency for instruction fine-tuning across different models. In addition, we compared the validation loss of MISA with other BCD methods, BAdam and LISA. shown in Figure 3, the x-axis represents training time in minutes. MISA and LISA have similar total training time, while BAdam is faster than both. Despite this, MISA exhibits much better convergence compared to LISA and BAdam.
Table 5: Comparison of different methods on MMLU, MMLU-pro and MT- Results of LISA, GaLore, and LoRA in MMLU and MT-Bench are cited from L[SA[41], while others are reproduced by us. MT-Bench evaluation was conducted using GPT-4. Memory was tested with batch size 2.
Model
Method
Mem.(GB)
MMLU (5-shot)
MML.Upp (5-shot)
MTT-Bench
LoRA
5.00
25.81
11.32
1.90
GaLore
5.00
25.21
10.96
TiayLLaMA
LISA
5,92
26.02
11.59
2.57
BAdam
4.49
25.27
11.44
2.42
MISA
4.36
25.40
11.65
2.73
LoRA
20.48
45,50
20.57
4.45
GaLore
19.25
45.56
20.19
4.63
LLMMA2-7B
LISA
22.79
46.21
20.85
4,94
BAdam
19.83
45.61
20.68
4.81
MISA
19.72
46.27
LoRA
21.96
61.78
32.96
4.41
GaLore
20.87
$7.87
31.87
4.36
Mistral-7B
LISA
26.90
62.09
32.83
BAdam
21.21
62.11
32.79
5.03
MISA
21.18
62.90
33.44
3.00
1.0
8.30
USAA
BA-dum
BAdkm
MSA
MISA
MISA
3.00
1:00
3.0
1.00
200
Figure 3: Validation loss of LISA, BAdam, and MISA across three epochs of fine-tuning Mistral-7B (left), LLaMA2-7B (middle) and TinyLLaMA (right) on the CPPc dataset. The x-axis represents training time (minutes).
4.4 Pre-training
We evaluate MISA for pre- pre-training LLaMA models [54] the C4 dataset [46]. We pre-trained the 130M and 350M variant of L.LMM2 [33]. Unlike fine-tuning, where the embedding and output layers remain frozen, trained them using the standard Adam optimizer. The number of training tokens is 2.75B. We report results for Adam, Grlere, and MISA. Table 6 reports the validation perplexity after 52K training steps. With high-rank subspace (i.e., 256, 8 25%), MISA outperforms GaLore and Adam on the 130M model. While on the 350M model, MISA outperforms GaLore and is slightly behind Adam. As shown in Fig. MISA converges better than GaLore, and is similar with Adam's 'ss performance in the later training stage. MISA's strong performance indicates that it can be viewed as a regularization of Adam, and many gradients are redundant during full full- parameter training.
Table 6: Validation perplexity of pre-training LLaMA 350M model on C4 dataset.
GaLore
MISA
Adam
r=32
r=256
6=3%
4=25%
Perplexity
24.63
44.88
28.27
36.07
23.81
LLaMA 130M
Memory(GB)
6.03
5.73
5.85
5.09
5.87
LLaMA 350M
Memory(GB)
Perplexity
21.30
13.58
11.81
39.43
24.34
12.14
30.62
9,99
22.11
11.63

---

Validation Perplexity
Adam
Adam
80
GeLore(m296)
45
MISA = 25%
40
MISN(5 25%)
50
36
30
20
10000
20000
30030
40030
50300
10000
20000
30000
40000
Steps
Steps
Figure 4: Pre-training dynamies for LLaMA 130M (left) and LLaMA 350M (right) on the C4 datasct.
5 Conclusion and Limitation
Conclusions. In this paper, propose MISA, a novel memory-efficient training method. MISA decomposes each layer into smaller modules and employs an importance sampling mechanism for module training. We provide detailed memory analysis demonstrating that MISA achieves high memory efficiency in long-sequenee fine- tinettnning Additionally, we show that MISA attains convergence rate of (1/vKK), where K is the total number of updates.
Limitations. Due resource constraints, pre-trained MISA at small scale, and whether MISA can scale to large-scale LLMs pre-training (e.g.. 7B, 70B, or larger models) remains to be validated. Additionally, MISA has only been verified on text-modal Transformer-basee LLMs (e.g., LLaMA series, Qwen2.5 series, Mistral series) and has not been extended to multi-modal models or non- Transformer architectures, its adaptability across different modalities and model structures needs further exploration.
6 Acknowledgements
Kun Yuan supported by the National Natural Science Foundation of China (NSFC) under Grants W2441021, 92370121 and 12301392. Tao Yao is supported by the National Natural Science Founda- tion of China (NSFC) under Grants W2441021, 72371172, 72342023, and 71929101.
References
[1] Anonymous. HiRA: Parameter-efficient hadamard high-rank adaptation for large language models. In The Thirteenth International Conference on Learning Representations, 2025.
[2] D. Biderman, J. Portes, J. J. G. Ortiz, M. Paul, P. Greengard, C. Jennings, D. King, S. Havens, V. Chiley,
J. Frankle, et al. Lora learns less and forgets less. arXiv preprini arXiv:2405.00677 09673, 2024.
[3] Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
[4] T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublincar memory cost, 2016.
[5] Y. Chen, Y. Zhang. Y. Liu, K. Yuan, and Wen. A memory efficient randomized subspace optimization method for training large language models. arXiv preprint arXiv:2502.07222 2025.
[6] C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 1905 2019.
[7] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1803:0180..0557, 2018.
[8] K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and J. Schulman. Training verifiers solve math word problems. arXiv preprint arXiv:21 14 2021.
[9] T. Dao. FlashAttention: Faster attention with better parallelism and work partitioning. In International Conference Learning Representations (ICLR), 2024.
[10] ALL Deepseek-v3 technical report, 2024.
10

---

[11] T. Dettmers, M. Lewis. S. Shleifer, and L. Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprini 10..022 2021.
[12] T. Dettmers, A. Pagnoni, A. Holtzman, and I. Zettlemoyer. Qlora: Efficient finetuning of quantized llms, 2023.
[13] T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer. Qiora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems, 36, 2024.
[14] A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, Letman, A. Mathur, A. Schelten, A. Yang. A. Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
[15] M. Elhoushi, A. Shrivastava, D. Liskovich, B. Hosmer, B. Wasti, L. Lai, A. Mahmoud, B. Acun, S. Agarwal, A. Roman, A. Aly, B. Chen, and C.-J. Wu. Layerskip: Enabling early exit inference and self self-speeulative decoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), page 12622-12642. Association for Computational Linguistics, 2024.
[16] A. Fan, E. Grave, and A. Joulin. Reducing transformer depth on demand with structured dropout, 2019.
[17] C. Fan, Li, X. Ao, F. Wu, Y. Meng, and X. Sun, Layer-wise model pruning hased on mutual information, 2021.
[18] X. Han, W. Zhao, N. Ding, Z. Liu, and M. Sun. Ptr: Prompt tuning with rules for text classification. A/ Open, 3:182-192, 2022.
[19] He, C. Zhou, X. Ma, T. Beng-Kirkpatrick, and G. Neubig. Towards unified view of parameter-efficient transfer learning. arXiv preprint 10.0036,,,, 2021.
[20] S. He, L. Ding, D. Dong, M. Zhang, and D. Tao. Sparseadapter: An easy approach for improving the parameter-efficiency of adapters. arXiv preprint arXiv:2210.04284 2022.
[21] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding. arXiv preprini wXiv:2009.033000 2020.
[22] M. Hong, X. Wang, M. Razaviyayn, and Z.-Q. Luo. Iteration complexity analysis of block coordinate descent methods. Mathematical Programming, 163:85- 114,,142222. 2017.
[23] N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. De Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly. Parameter-efficient transfer learning for nip. In International conference on machine learning. pages 2790-2799. PMLR, 2019.
[24] F. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, L. S. Wang. L. Wang, and W. Chen, Lora: Low-rank adaptation of large language models. arXiv preprint 1rXiv:2106.096)) 2021.
[25] Z. Hu, Y. Lan, L. Wang, W. Xu, E.-P. Lim, R. K.-W. Lec, L. Bing, and S. Poria. Llm Lim-adapters: An adapter family for parameter-efficieet fine-tuning of large language models. arXiv preprint arXiv:2304.019 2023.
[26] A. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. Chaplot, D. d. Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, et al. Mistral 7b, arXiv preprint arXiv:2310.06825, 2023.
[27] T. Jiang, S. Huang. Luo, Zhang, Huang, Wei, W. Deng. F. Sun. Q Zhang, D. Wang, et al. Mora: High-rank updating for parameter-efficit fine-tuning, arXiv preprint arXiv:2405.12130 2024.
[28] R. Koncel-Ked Kedziorski. Roy, A. Amini, N. Kushman, and Hajishirzi. Mawps: A math word problem repository. In Proceedings of the 2016 conference of the north american chapter of the association for computational linguistics: human language technologies, pages 1152- 1152,1157, 2016.
[29] B. Li. J. Chen, and J. Zhu. Memory efficient optimizers with 4-bit states, 2023. URL https://rrxiv. org/ubs/2309.005077
[30] P. Li, L. Yin, X. Gao, and S. Liu. Owlore: Outlier-weighed layerwise sampled low-rank projection for memory-effeient Ilm fine-tuning, 2024.
[31] X. L. Li and P. Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint 101 000900, 2021.
[32] V. Lialin, S. Muckatira, N. Shivagunde, and A. Rumshisky. Relora: High-rank training through low-rank updates. In The Twelfth International Conference Learning Representations, 2023.
11

---

[33] V. Lialin, N. Shivagunde. S. Muckatira, and A. Rumshisky. Relora: High-rank training through low-rank updates, 2023.
[34] W. Ling, D. Yogatama, C. Dyer, and Blunsom. Program induction by rationale generation: Learning solve and explain algebraic word problems. arXiv preprint arXiv: 1705.04146.. 2017.
[35] S.-Y. Liu, C.-Y. Wang. H. Yin, P. Molchanov, Y.-C. F. Wang. K.-T. Cheng, and M.-H. Chen. Dora: Weight-dccomposed low-rank adaptation. In Forty-first International Conference on Machine Learning, 2024.
[36] Y. Liu, Y. Zhang, Q. Li, T. Liu, S. Feng, D. Wang. Y. Zhang, and H. Schütze. Hift: A hierarchical full parameter fine- tuning strategy. arXiv preprint arXiv:2401.1520,, 2024.
[37] Q. Luo, H. Yu, and X. Li. Badam: memory efficient full parameter optimization method for large language models. Advances in Neural Information Processing Systems, 37:24926-24958,, 2024.
[38] K. Lv, H. Yan, Q Guo, H. Lv, and X. Qiu. Adalomo: Low-memory optimization with adaptive learning rate. arXiv preprint 23011930.. 2023.
[39] K. Lv, Y. Yang. T. Liu, Q. Gao, Q. Guo, and X. Qiu. Full parameter fine-tuning for large language models with limited resources. arXiv preprint arXiv:2306.09770 2023.
[40] T. Mihaylov, Clark, T. Khot, and A. Sabharwal. Can a suit armor conduct electricity? new dataset for open book question answering. arXiv preprint arXiv: 1889902779. 2018.
[41] R. Pan, X. Liu, S. Diao, R. Pi, J. Zhang, C. Han, and T. Zhang. Lisa: layerwise importance sampling for memory-efficient large language model fine-tuning, Advonces in Neural Information Processing Systems, 37:57018 57049, 2024.
[42] A. Patel, Bhattamishra, and Goyal. Are nlp models really able to solve simple math word problems? arXiv preprint arXiv:2103.07191, 2021.
[43] J. Pfeiffer, I. Vulié, I. Gurevych, and S. Ruder. Mad-x: An adapter-based framework for multi-task cross-lingual transfer. arXiv preprint arXiv:2005.0005. 2020.
[44] G. Qin and Eisner. Learning how ask: Querying Ims with mixtures of soft prompts. arXiv preprint arXiv:2104.0m 06599999..9........1 2021.
[45]
X. Qiu, T. Sun, Y. Xu, Y. Shao, N. Dai, and X. Huang. Pre-trained models for natural language processing: A survey. Science China technological sciences, 63(10):1872-1897 2020.
[46] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1667, 2020.
[47] A. Ramesh, Ganapathiraman, H. Laradji, and M. Schmidt. Blockllm: Memory-efficient adaptation of lims by selecting and optimizing the right coordinate blocks. arXiv preprint arXiv:2406. 2024.
[48] M. Razaviyayn, M. Hong, and Z.-Q. Luo. A unified convergence analysis of block successive minimization methods for nonsmooth optimization, SIAM Journal on Optimization, 23(2): 126-1153, 2013.
[49] B. Rozière, J. Gehring, F. Gloeckle, S. Sootla, I. Gat, X. F. Tan, Y. Adi, J. Liu, R. Sauvestre, T. Remez, J. Rapin, A. Kozhevnikov, I. Evtimov, J. Bitton, M. Bhatt, C. C. Ferrer, A. Grattafiori, W. Xiong, A. Dé- fossez, J. Copet, E. Azhar, H. Touvron, L.. Martin, N. Usunier, T. Scialom, and G. Synnaeve. Code llama: Open foundation models for code, 2024.
[50] K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-100 2021.
[51] M. Sap, H. Rashkin, D. Chen, R. LeBras, and Y. Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint arXiv:1904.097290 2019.
[52] G. Shi, Z. Lu, X. Dong, W. Zhang, X. Zhang, Y. Feng, and X.-M. Wu. Understanding layer significance in Ilm alignment, 2024.
[53] R. Taori, 1. Guirajani. T. Zhang. Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto. Stanford alpaca: An instruction followwig Ilaaa model. https://iitub.ccmm/auuu http://hhhbc.ccmm//aaaaaaaaaaada lab/stnffordd alpaca, 2023.
12

---

[54] H. Touvron, L. Martin, K. Stone. P. Albert, A. Almahairi. Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel. L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull. D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lcc, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, I. Molybog, Y. Nic, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A, Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, Xu, Z. Yan, I. Zarov, Y. Zhang. A. Fan, M. Kambadur, S. Narang. A. Rodriguez, R. Stojnic,
S. Edunov, and T. Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
[55] P. Tseng. Convergence of block coordinate descent method for nondifferentiable minimization. Journal of optimization theory and applicationn, (09:475-4944 2001.
[56] Y. Wang, S. Agarwal, S. Mukherjee, X. Liu, J. Gao, A. H. Awadallah, and J. Gao. Adamix: Mixture-of- adaptations for parameter-efficiee model tuning. arXiv preprint arXiv:2205. 12410, 2022.
[57] Y. Wang, X. Ma. G. Zhang, Y. Ni, A. Chandra, S. Guo, W. Ren, A. Arulraj. X. He. Z. Jiang, et al. Mmlu- pro: more robust and challenging multi-task language understanding benchmark. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024.
[58] J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al. Chain-of-thouge prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
[59] S. J. Wright. Coordinate descent algorithms. Mathematical programming, 151(1):3-34, 2015.
[60] L. Xu, H. Xie, S.-Z. J. Qin, X. Tao, and F. L. Wang. Parameter--fffcient fine-tuning methods for pretrained language models: A critical review and assessment. arXiv preprint arXiv:2312.12148, 2023.
[61] Y. Xu and W. Yin. A block coordinate descent method for regularized multiconvex optimization with applications to nonnegative tensor factorization and completion. SIAM Journal on imaging sciences, 6(3):1758- -7787999-7789, 2013.
[62] Y. Xu and W. Yin. Block stochastic gradient iteration for convex and nonconvex optimization. SIAM Journal on Optimization, 25(3):1686-11) 1716, 2015.
[63] A. Yang, B. Yang, B. Hui, B. Zheng, B. Yu, C. Zhou, C. Li, C. Li, D. Liu, F. Huang. G. Dong, H. Wei, H. Lin, J. Tang, J. Wang, J. Yang, J. Tu, J. Zhang. J. Ma, J. Xu, J. Zhou, J. Bai, J. He, Lin, K. Dang. K. Lu, K. Chen, K. Yang. M. Li. M. Xue, N. Ni, P. Zhang. P. Wang. R. Peng. R. Men, R. Gao, R. Lin, S. Wang. S. Bai, S. Tan, T. Zhu, T. Li. T. Liu, W. Ge. X. Deng. X. Zhou, X. Ren, X. Zhang, X. Wei, X. Ren,
Y. Fan, Y. Yao. Y. Zhang, Y. Wan, Y. Chu, Y. Liu, Z. Cui, Z. Zhang, and Z. Fan. Qwen2 technical report. arXiv preprint arXvv:2407. 10671, 2024.
[64] K. Yao, P. Gao, L. Li, Y. Zhao, X. Wang. W. Wang, and J. Zhu. Layer-wisc importance matters: Less memory for better performance in parfmeter-efflcientt fine-tuning of large language models, 2024.
[65] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: 1005:1000000830 2019.
[66] B. Zhang, 7.. Liu, C. Cherry, and O. Firat. When scaling meets Ilm finetuning: The effect of data, model and finetuning method. arXiv preprint 47X11.2402.17193, 2024.
[67] P. Zhang, G. Zeng, T. Wang, and W. Lu. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.022 02.34081..3.,,. 2024.
[68] Q Zhang, M. Chen, A. Bukharin, N. Karampatziakis, P. He, Y. Cheng, W. Chen, and T. Zhao. Adalora: Adaptive budget allocation for parameter-efficien fine-tuning, 2023.
[69] S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diah, X. L. X. Win, et al. Opt: Open pre-trained transformer language models. arXiv preprint arX:v:205.00166, 2022.
[70] Y. Zhang, C. Chen, Z. Li, T. Ding, C. Wu, D. P. Kingma, Y. Ye, Z.-Q. Luo, and R. Sun. Adam-mini: Use fewer learning rates to gain more, 2024.
[71] J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. Galore: Memory- Ilm training by gradient low-rank projection. arXiv preprint
2024.
[72] L. Zheng, W.-L. Chiang, Y. Sheng. S. Zhuang. Z. Wu, Y. Zhuang, Z. Lin, Z. Li. D. Li. E. Xing, et al. Judging llm-as-a-judge with me-bench and chatbot arena. Advances in Neural Information Processing Systems, 36:46595 46623, 2023.
13

---

[73] Z. Zhong, D. Friedman, and D. Chen. Factual probing is [mask]: Learning vs. learning recall. arXiv preprini arXiv:2104.05244 2021.
[74] Y. Zhou, M. Zhang. J. Zhu, R. Zheng, and Q. Wu. A randomized block-coordinate adam online learning optimization algorithu. Neural Computing and Applications, 32(16):12671-1268) 2020.
[75] L. Zhu, L. Hu, J. Lin, and S. Han. LIFT: Efficient layer-wise fine-tuning for large model models, 2024.
14

---

Appendix for "MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling"
Contents
Introduction
Module-wise Importance Sampling
2.1
Problem Formulation
2.2
Block Sampling Strategy
2.3
Partition the Weights into GrainGr Modules Instead of Coarse Layers
2.4
Module-wise Importance Sampling (MISA) Method..
2.5
Memory Analysis
3
Convergence Analysis
Experiments
4.1
Commonsense Reasoning
4.2 Math Reasoning
8
4.3
Instruction Tuning...
4.4 Pre-training
5 Conclusion and Limitation
10
Acknowledgements
10
A
Related Work
17
Algorithm Details
17
B.1 Module selection strategy
17
B.2 Definition of the scaled gradient norm
18
C Additional Experiments
18
C.1 Memory Efficiency When Scaling To 70B Model
C.2 Enhancing LoRA with MISA
18
C,3 Computation Efficiency
19
Ablation Study
19
D.I Ablation Study: Impact of Clearing Optimizer States
19
D.2 Ablation Study: Impact of Hyperparameters
19
D.3 Impact of Inner-Loop Iteration T
20
D.4 Ablation Study: Sampling Strategy
21
15

---

D.5
Ablation Study; Importance Scoring Methods
21
D.6
Ablation Study; Impact of Each Module
21
Theoretical Analysis.
23
E,1 Notation
23
................................
E.2 Proof of theorem
33
Memory Analysis
34
F.1 Memory analysis of laser-wise optimization method
35
F.2
Memory analysis of MISA and LoRA
................................
37
F2.1
Memory analysis of Module-wise BCD
37
F.2.2
Memory analysis of LoRA
38
................................
E3
Memory analysis of GaLore
................................
38
E4
Memory analysis of MISA
39
F.5
Memory analysis of MISA's indicators
39
F66
Peak memory comparison between MISA and layer-wise method
40
F7
Peak memory comparison between layer-wise method and LoRA
F88
Paras/Peak-Memory comparison between layer-wise method and LoRA
40
Computation Analysis
G.I
Computation analysis of layer-wise optimization methods
42
G.2 Computation analysis of module wise optimization methods
43
G.3 Computation analysis of MISA's indicators
43
Broader Impacts
43
I Experimental Hyperparameters
43
L.1 Hyperparameters for commonsense reasoning
44
1.2 Hyperparameters for math reasoning
44
L.3 Hyperparameters for instruction fine-tuning
45
1.4
Hyperparameters for pre-training
46
16

---

A Related Work
Parimeter-effiiiet fine-tuning. Adapter- based methods introduce new modules (e.g.. fully connected layers) into the backbone model and update only these modules during training [43, 23, 19, 56, 20), incurring additional computational costs during inference. Prompt- basep methods prepend random soft tokens to the input (typically as prefix) and train their embeddings while keeping the backbone L.M frozen [31, 18, 73, 44), also introducing inference overhead.
LoRA and its variants, including A [13], DoRA [35], and Ada [68], have become state-of- the-art PEFT methods. LoRA injects low-rank matrices into linear layers estimate weight updates while keeping the backbone LLM parameters frozen. Unlike adapter- and prompt-based methods, LoRA has no additional inference overhead, the low-rank matrices can be merged into the model weights. To address the limitations of the low-rank structure, recent methods such as ReLoRA [32], MoRA [27], and HiRA [1] have explored high-rank updates based on LoRA.
Memory-efficient optimization. LOMO [39] eliminates optimizer and gradient memory costs during training with standard SGD). Adal HaMO [38] extends LOMO by integrating a fused backward operation with an adaptive learning rate per parameter... Gal Lore [71] enhances memory efficiency by projecting gradients into low-rank subspace. Quantization techniques parameters, gradients, optimizer states, and intermediate activations [11, 29, 10]. Adam-mini [70] reduces memory usage by optimizing learning rate allocation in Adam.
widely applied model
Layer-wise optimization. Inspired by classical Block Coordinate Descent (BCD), layer-wise opti- mization has emerged promising approach for fine-tuning large language models. Methods such as LIFT [75]. HIFT [36]. and BAdam [37] update LLM layers sequentially, achieving performance comparable parameter- efffiiienee fine-tuning (PEFT) methods while significantly reducing compu- tational costs. LISA [41] identified a skewed weight-norm distribution across layers LoRA and introduced randomized layer update strategy, keeping the embedding and language modeling head layers activated. Similarly, BlockLLLM [47] prioritizes frequent updates to layers with larger gradient norms. However, none of these approaches incorporate module wise importance sampling.
Varying importance of modules in LLM fine-tuning, Previous studies have shown that different layers contribute unequally to LLM performance. LayerSkip [15], LayerDrop [16], and Layer-wise Model Pruning [17] demonstrated that the overparameterizztion of pre-trained models leads to layer redundancy, enabling high performance be maintained even after removing certain layers. IST [64] and ILA [52] further highlighted that layer importance varies during LLM fine-tuning. Additionally, studies PEFT have revealed that different modules within transformer blocks exhibit varying ranks [2, 68). Owl oore [30] identifies layer importance by detecting outliers in the model weights. However, these insights have not yet motivate module-wie wise importance sampling in LLM optimization.
B Algorithm Details
B.1
Module selection strategy
Algorithm 2 Module selection
Require: Active set Tn, Non-active set B, Sampling threshold 8, Probability distribution P", Total model parameters modell Counter S. 1: Initialize:7, 0, B- (M1, .2.,
..., MR). S
2: while * do
3: Sample module m from B based on the probability distribution P".
4: = m7.
5: if S + |m| < 8nmodee then
6: Tn = Tn +
7:
end if
8:
end while
9: Return: Tn
17

---

B.2 Definition of the scaled gradient norm
For given module m4, we define its scaled gradient norm be the ratio of the Frobenius norm gradient to the square root of the number of parameters Loille in practice, which can be considered m
definition of normalization.
C Additional Experiments
C.1 Memory Efficiency When Scaling To 70B Model
Fig. 5 demonstrates that when training 70B model, MISA is still more memory efffeietfficientt than LoRA under long sequence length settings, and the use Muss--attention9) has minimal impact on overall memory consumption.
50.0
USA (P-7
222
Mantary (58)
HEAA (PTU
200
Labk
Menooy
25.0
24
Mannery
20.0
162
120
530
1000
1530
100
525
1000
1502
100
500
1000
1500
Sequence Longth
Sequence
Sequence Longth
(a) LLaMA3-SB
(b) LLaMA3-70B
(c) LLaMA3-70B w/ flash-attention
Figure 5: Comparison of peak memory consumption LLaMA3-8B and LLaMA3-70B. (c) used flash-attention technique.
C.2 Enhancing LoRA with MISA
A key advantage of LoRA over MISA is its ability to train separate adapters for different downstream tasks. Instead of storing full model weights, LoRA only requires storage for task-specific adapters, significantly reducing storage costs. Since MISA is orthogonal to LoRA and can be used alongside it, we explored hybrid approach combining both methods. Given a model weight W with LoRA applied, it is represented as W W + AB, where A and B are treated as separate modules. We freeze the original model parameters and consider the LoRA adapters as the total set of trainable parameters ORA MISA dynamically activates LoRA adapters while ensuring that the number of active parameters each step remains below TiRORA X Furthermore, we found that optimizer states not the primary source of memory consumption for LoRA MISA. Thus, retain optimizer states for all LoRA adapters without clearing them. Our results show that MISA enhances full LoRA's performance while further reducing memory usage. As shown in Fig. 6, activating more than 30% of LoRA parameters achieves or even surpasses full LoRA's performance while saving approximately 7.9% of memory. We also compared IST [64] with MISA (see Table 7) and MISA outperformed IST.
73.0
LORAMMIBA
36.0
72.5
34.5
Accuracy ray
34.0
Memoory
33.5
71.0
33.0
70.5
32..5
70.0
10%
30%
50%
70%
90%
5
Figure 6: Accuracy and memory consumption LoRA + MISA fine-tuning with varying on math reasoning tasks.
18

---

Table 7: Comparison of MISA and IST on LLaMA-7B and LLaMA3- 8B across eight commonsense reasoning tasks. Results of LoRA+IST are cited from IST paper.
Model
Method
BoolQ
PIQA
SIQA
Hella.
Wino.
ARC-e
ARC-c
OBQA
Avg-i
LoRA
68.9
80.7
77.4
78.1
78.8
77.8
74.8
74.7
LoRA+IST
68.7
81.7
77.3
82.7
78.7
80.6
62.4
80.0
76.5
LLaMA-7B
LORA+MISA
70.4
81.5
78.9
83.4
78.9
81.2
63.1
79.7
77.1
LoRA
70.8
85.2
79.7
92.5
84.9
88.9
78.7
84.4
82.5
LORA+IST
72.7
88.3
94.7
84.4
84.6
LLaMA3-8B
LORA+MISA
74.4
88.9
81.8
95.3
86.1
91.3
81.4
86.5
85.7
C.3
Computation Efficiency
training step includes forward propagation, m backward propagation, and optimizer computations (including parameter updates and any algorithm-specific operations). Table reports the average time per step for each method. GaLore incurs the highest optimizer computation overhead because performs a full-model SVD (with very large time complexity) every 2000 steps. For MISA, we further conduct detailed operation-level profiling (under the same hardware and training configurations as Table 8) to quantify the overhead of its unique components (e.g., module sampling, gradient norm tracking). Results show that these components introduce negligible additional time (accounting for less than 0.05% of the total per-step time), as their key operations (c.g.. gradient norm aggregation, sampling probability calculation) are either lightweight or amortized over multiple inner-loop steps. Thus, MISA's computation efficiency is only behind BAdam's s, while it also achieves superior convergence with respect to training time (see Fig. 3).
Table 8: Comparison of training time for fine-tuning LLaMA2-7B on the Alpaca-GPT4 dataset. The Forward, Backward, and Optimizer time represent the average time per step. The last column shows the overall training time epochs.
Method
Forward(ms)
Backward(ms)
Optimizer(ms)
Averged cost per step(ms)
LoRA
70.3
104.5
4.5
179.3
GaLore
35.7
140.5
267.4
443.6
BAdam
34.3
42.9
1.5
78.7
LISA
36.4
84.1
2.9
123.4
MISA
35.2
74.6
1.8
111.6
D
Ablation Study
D.1
Ablation Study: Impact of Clearing Optimizer States
In line 17 of algorithml, we clear the optimizer state to save memory consumption. We explored method to maintain the optimizer states of active layers during training. When layer is frozen, its complete optimizer state is preserved and offloaded CPU memory. These saved states are then precisely restored to GPU when the layer is reactivated in subsequent training cycles.
As shown Fig. preserving optimizer states has significant impact fine-tuning loss but brings significant extra pretraining perplexity. However, we don't have particularly good explanation for this phenomenon. Our guess that in the early stage of pre-training, when the model still in the coarse-tuning stage, the dependence on the history of momentum is not that strong, so the clearing optimizer may be able avoid local minima and thus improve optimization.
D.2 Ablation Study: Impact of Hyperparameters
We primarily investigated the effects of three hyperparameters: the learning rate, MISA's 5, and
MISA's n7. 77 controls the trade-off between importance sampling and uniform sampling. When 77 0.
19

---

presuree optm
MISA proser opónn.
0.85
Validation Perplesity
MISA
0.00
0.55
0.35
0.50
1000
2000
4000
5000
7000
10300
20020
20000
50000
Steps
Steps
Figure 7: Fine- Fine-tuning Ioss (left) and pre-training perplexity (right) of LLaMA3-8B Fine-tunin used MATHIOK dataset, and pre-training used C4 dataset. MISA discards the optimizer states when switching activated layers, while "MISA w/ preserve optim." indicates retaining the optimizer states.
MISA reduces to uniform sampling. 5 determines the number of parameters updated at each step and scales proportionally with memory consumption.
Fig. presents the impact of v and 5. In practice, variations in v have a minor effect on model accuracy, whereas MISA is highly sensitive the learning rate. An excessively large learning rate significantly degrades performance of the model.
05.0
2..5
625
60.0
F1.
Loariin
MSA
Figure 8: The accuracy of LLaMA3-8B 8B on math reasoning tasks with different learning rate and r..
Fig. 9 illustrates the validation loss over three epochs of fine-tuning LLaMA2-7B on the Alpaca-GPTV GPT4 dataset for different We observe that larger å lead to more severe overfitting in the third epoch, further supporting the interpretation of MISA as regularization method of full parameter fine-tuning.
6 1%
3%
15%
9k1g
Figure 9: Validation loss of LLaMM2-7B on Alpaca-GPT4 dataset.
D.3 Impact of Inner-Loop Iteration T
We conduct an ablation study to evaluate the sensitivity of MISA to the inner-loop iteration count (defined in Algorithm I, representing the number of Adam updates for sampled module before switching to another module). Experiments are conducted on the Mistral-7B model fine-tuned on the Alpaca-GPT4 dataset, with other hyperparameters fixed as: learning rate = Ie-5, 8 = 3%, = 0.5, and batch size 16. Table 9 reports the validation loss and MMLU score across different T values.
As shown in Table 9, varying from 5 500 does not significantly affect the convergence of validation loss or the MMLU performance. This observation aligns with the findings in BAdam [37] and LISA [41], confirming that MISA's convergence is robust to the choice of T result supported by our theoretical analysis in Theorem 1, where MISA is proven to converge for any number of
inner-loop updates Notably, when > 200, we observe slight degradation in MMLU score (e.g., 46.01 at T 200 and. 46.27 at T 50), which may stem from over-fitting to the sampled module during extended inner-loop updates. Thus, we recommend setting T 30-50 in practice to balance computational efficiency and performance.
20

---

Table 9: Ablation Study of the Inner-Loop Iteration T for the Mistral-7B Model on Alpaca-GPT4
T
Validation Loss (1)
MMLU (5-shot, 1)
5
0.877
46.22
15
0.874
46.23
30
0.871
46.17
50
0.873
46.27
100
0.877
46.19
200
0.881
46.01
500
0.879
45,89
D.4 Ablation Study: Sampling Strategy
To validate the necessity of MISA's sampling strategy, we explored alternative approaches. The Uniform strategy samples modules randomly without computing importance scores, effectively replacing BAdam's layer-wise optimization with module-wise optimization. The Top-K strategy selects the most important modules while ensuring that the total sampled parameters remain below of all parameters, In contrast, the Bottom-K strategy selects the least important modules under the same å constraint. As shown in Table 10, the Uniform and Top-K strategies achieve similar performance but remain inferior to MISA, while the Bottom-K strategy significantly reduces accuracy. never sampled, which may explain
We observed that with the Top-K strategy, many modules its suboptimal performance. The effectiveness of MISA can be attributed to its balanced approach, combining elements of both the Uniform and Top-K K strategies. It prioritizes modules with higher importance scores while ensuring that all modules have a chance to be sampled.
Table 10: Results of different sampling strategies for fine-tuning I.L.M3-8BB on math and commonsense reasoning tasks. The reported scores indicate the average accuracy across all tasks.
Strategy
Math
Commonsense
MISA
73.6
86.6
Uniform
71.1625)
85.9
Top-K
71.2
86:00)
Bottom-K
69.6(4.0
82.1(445)
D.5 Ablation Study: Importance Scoring Methods
In (4), we identify each module's importance score via its gradient norm. To validate its effectiveness, we also evaluated two alternative scoring functions: the module's weight norm and its total parameter count. Table compares those importance scoring methods and demonstrates that MISA's sampling strategy significantly outperforms the alternatives, thereby validating the efficacy of MISA.
Table 11:
Comparison of different scoring function in MISA. The accuracy of commonsense resoning and math
reasoning benchmarks are averaged across all tasks.
Model
Method
MMLU
MMLU-pm Commonsense
Math
Weight Norm
64.5
35.9
85.7
71.9
LLaMA3-8B
Number of Parameters
63.7
36.2
85.9
72.7
MISA(Gradien) Norm)
65
36.5
86.6
73.6
D.6
Ablation Study: Impact of Each Module
To assess the impact of each module, we applied MISA to individual modules while keeping others frozen. Fig. 10 shows that the effectiveness of fine-tuning follows the order: W,, Wt < W, < W., W Watte Waps Wdown. Interestingly, when applying MISA all modules, the sampling frequency does not strictly follow this order, shown in Fig. 11. We think there may be three
21

---

reasons: (1) Sampling frequency does not directly reflect importance more critical layers may require less fine-tuning for optimal performance. (2) When training multiple modules simultaneously. the combined importance of those modules may not be simply the sum of each module's individual importance. (3) Module sizes vary significantly; instance, W gates Wap- Wdown contain over ten times many parameters Wa and W., Small-size modules are less likely to exceed the sampling threshold 6, leading to more frequent sampling.
W;
We
W
Woom
W,
Wu
Ww
46
AQUA
GSMBK
74
12
70
SVAMP
MAWPS
Figure 10: Accuracy comparison on math reasoning tasks when sample each individual module separately.
Sampling Frequency
80
14
W.
60
Wgatn
40
W.
Wcown
20
2
9
1011 171311151611
1920212222 25777222292222222722898988999
Layer
Figure 11: The frequency of module sampling across all layers when applying MISA to LLaMA3-8B.
Therefore, we further applied both uniform and MISA sampling strategies to individual modules in isolation, in order to assess which modules benefit most from MISA. As shown in Table when fine-tuning Wi W gate. W Woown and W Woown. MISA achieves substantially larger accuracy gains than uniform sampling.
Table 12:
Average accuracy of LLaMA3-8B math reasoning benchmark when fine-tuning only 1 module type with uniform and MISA sampling strategies. Red indicates an improvement more than 1%.
Model
Method
W
q
W
W,
Wo
W
gate
W
up
W Wdown
Uniform
70
67.1
70.5
71.2
69.2
71.4
70.5
LLaMA3-8B
MISA
69.5
68.8
70.5
71.7
71.5
72.5
72.1
22

---

E Theoretical Analysis.
E.1 Notation
Let f Rd + R be the objective function to minimize.
Let f = infoer= (0) be the optimal value of f.
Let V denote the estimator for the gradient of the queried points.
Let denote the learning rate.
Let n denote the block epoch, n 1, 2,..,,
Let denote inner iterations for a certain block, t = 0, 1, T.
Let gn. denote the sampled data at n-th block epoch and t-th inner iteration.
Let Tn denote the sampled block at n-th block epoch.
Let b denote the block's 's ID, b = ..,
Let 0nd denote the parameter when selected module Tn at n-th block epoch and t-th inner adam step.
Let g- denote the partial derivative conditioned on data erd at n-th block epoch and t-th inner adam step.
Let Tm
1 Ei=illes |/ be the cumulative gains during a sampling period.
Let I
Diag 1/2 (5 + 110%% [/l...tt Ilmin - E), Jme max (v% || max)
Let 0m denote the full parameter at n-th block epoch and t-th inner adam step.
0rl 0-.l== == -0p (07-1:"' 00 01-1 .... 0"-1,- On.
0--111"' ..... .... 0n 1,7
Let F, denote the filtration containing all historical variables at and before computing 0m.
F, = ( 5 gh:", m., n fn.t- gn.l, ITh 1 m v (n.0)
Let E, denote the expectation conditioned Ft.
To enhance the smoothness of the proof procedure, we utilized three techniques.
1. Second-order Momentum Inheritance
To establish Lemma I, we utilize a second-order momentum inheritance technique. Upon completing the (n 1)-th block epoch, initialize the second-order momentum for the 7ith block epoch as:
7n | for n 2 2
This initialization strategy maintains memory efficiency, with its additional overhead being negligible compared to gradient storage requirements.
2. Additional Momentum Step
The additional momentum step arises naturally from the variable transformation in (11), This carefully designed operation prevents momentum wastage in 7, during the n-th block epoch, significantly simplifying the analysis at block transition points. The step defined as:
B1
m n.T
0 <- 0
1
B1
n.T UT
+
23

---

3. Single-module Gradient Propagation
While MISA typically activates multiple modules simultaneously, our theoretical analysis focuses on the single-module for clarity. The multi-module scenario can be represented by introducing diagonal projection matrix Qn, transforming the stochastic gradient Algorithm to:
a
Qn a0n.t-1 (8m+)))
L
g
where Ln denotes the set of activated modules at the 7-th step. Our analysis naturally extends to this generalized case.
Algorithm 3 MISA: View
Require: _N,N B, 17 a, 0, 0, 81, 32 E (0, 1) and the sampling block Tn at outer loop n > Partition weights into modules;
1: Partition the model into modules (not layers);
2: Initialize probability weights pa = . })
3: Initialize the module gradient estimate G for b E [B] and let G = (G). ".))
4: for 1,.,, N do
5: Sample a module Tn according to p";
b Importance sampling;
6: Initialize mn.0
v..0 0, 1/ -1168-1.11 1.7 max if 2:
7: for 1, .... do
8:
Sample batch of data g
9:
Calculate block stochastic gradient 3h. 00 (6%.2-1) fene
for selected module Tn
10:
Update the corresponding first-order and second-order momentum as follows:
11:
mA + 31m7.-11 (1
9 wht <- B2 (32) (9Ph5)2
12:
Update 5.. as follows:
6n. max (um4 | | [max)
AMSGrad- type normalization;
Update the corresponding module as follows:
15:
m < gr.t-1 am"; /(v 5m. +&)
16:
end for
Update G for each b [B] according to (4);
Track block gradient norm;
Update P
a+1
expinGg) exp(1G7)
for each b E [B];
Update sampling probability:
ΣB
0u+1, Tn
0u. Tm
-aT
91
m
Additional momentum step;
20: JT.,1 M. VTn
None
Clear optimizer states;
21:
end for
Return @N_ .PNN
Assumption 1. The loss function is L-smooth. And when restricted on i-th block, it is Ly-smooth.
ØL
0L
(VT VI. (02) < T,l00 0²
|| S T. 10 0211, i = 1,... B.
00:
00,
Let max {1,
Assumption 2. (Unbiased and bounded stochastic gradient). For all inner iterates t n-th block epoch, the stochastic gradient g"h is unbiased and uniformly bounded by some constant G 2
Eg [o Vf (8H.1-1), 119":"1 G.
Assumption 3. (Bounded variance). For all inner iterates L at n-th block epoch, the variance of the stochastic gradient gh is uniformly bounded by some constant o² 2
E, [llg%. Vf (0mLe- 101xtt1)1133 2 <
24

---

Proposition 3. (Proof of proposition 1) Consider the optimization problem
B
max
G
rpiLppi Qn)
i=1
B
s.t. pi
2
0,
pi 1,
the solution
exp (?o
P
a+1
bb exp (To
where ?ju
Proof. We consider the equivalence problem
max TO0
G. KL(P..9))
1=1
Where 7o '7' construct the Lagrangian function C(P, A) of the optimization problem
B ((P.A) 7o Σ P.G.
B
B
Σ
p,log
+
x 2
Pi
- 1)
9B
==1
1=1
1=1
We take the partial derivative of pt and obtain
Oc
log 9B
logg; 1+ 1o + = =0, i= 1,,,
Op:
Which is equivalent to According to the normalization condition, we have
dB exp GM + - 1), i 1= B.
B
B
Σ
exp Gn + A -1)
1=1
Let Z = exp(1 - A), we have
Z Σ qB G
Which implies
P n+1 exp (T)
exp (no GT))
Σ 2 QB exp (no G!i
i== exp (7o G?!
Proposition (Proof of proposition Suppose that in each layer we have modules such that 0: = (0,,. P...... 0:x) ), the sampling probability of each layer be Pa, each module be Pigj the gain of each module be Gej at n-th step and certain iteration, we have
B
K
Σ GT 2
S
G
K
Proof. We consider p: Σx= Pi,j, which means the sum of each module's sampling probability is equal the sampling probability of the corresponding layer.
25

---

Noting that in layer-wise setting, the sampling probability of each module is k
Next, will prove
G Pt
Gippes
K
Noting that gain Ge is equal to the sum of each individual module gains G we have
Σ G
Σ Cn
(6)
==1
Where Pr, Pi.. Σ 2y=1 Ps.j =1
By symmetry, we can assume without loss of generality that Ch2 is ordered as follows.
G S G S... G.
From Proposition 3, it's obvious that
Ps, S P4.2 S S P..K
The left side of inequality (6) is arithmetic average, the ride side is weighted average of G.
From Chebyshev's ' inequality, incquality (6) holds. Summing over we obtain
B K
P. S
GRSPis
Corollary 1. The sampling probability (p. MEN Ji in algorithm 3 has uniform lower bound
Proof. From Assumption 2, for Vie B), N
have
ΣT 1=
11,,
n.t 112
TG2
Gn
G
T
T
We obtain
G = (1 - G
NN + 3G.
& (1 - B) G; + 3G
(7)
Let p= 1- 3. 3G.
Updating equation (7) for K times, we obtain
GN+K 1-p + + GN
1 - p
Noting that 0 < < 1, for Vk E N we obtain
G G < G² + (1- (G. G²)
Let क.
max ISnSN
[G 6 (1 N (G - G²)}
We have
cxp (nCrr)
< B Σ)=1 exp (nGr)
=P-""
Be
We get the lower bound I. Be
Let min((n (x;), which is the uniform lower bound of {p, MEN
Ji=1,2......
Lemma 1. Let 1n.t_ Diag (/2 +e, |1 Im +e),
Arml
m. t = 1, .... T 1m,e-1
Th
26

---

We have the following upper bound estimation
N
T
2
2
S
A
A
S
m=1
Where
Arm.0 1--11 m
Tn T4.
Proof. From algorithm 3, we have vh. =1,6,_11 Iaa I, 2 2.
For notational convenience, we denote Ir Uz. as The un in this proof.
T = Diag
1/2 + E, 06 Ilmin + +
= Diag
1/2 (/1: Ilmin +&)
1 max +
Diag 1/2 1/2 ||max + +
3 Diag
(// + e, Illmin + PT-1
which implies that AT, = FT-1 I is positive semi-definite. Hence, ||/' = Amar (Ali) 2 0
Using the convexity of Amax over symmetric matrices, have
Amax (T--1-Ti)
([txx((z-1) + Amax(-f())
(max (fi-(((T))))) - Aminn (Li))
M I
1 min +
0 max +
1
6 mmmi t
cr max +
1 mmx t
Ür m +
NT
1
1 min +
1i min +
0 max te
1h mix
+
Villmmi +
v T Imin +
v Voullmas +
1 max +
27

---

For the second sum of squared norms, notice that for scalars a 2 b 2 O, it holds that (a -)2 - b)(a
(a +
b) a² b². Therefore, the derivation can be repeated without the square roots:
r
T |
Amax
Amax (Tt-1) + Amax (-fi)) =
> (maxx Amin Amin((i)))
tm1
t=
Ammx (Tz-1))² (Amin (Ti))²
0 1llin +E
0g Imax +
10 Ilain +
1 mmax
| Imax +
0 max
+
T
1
+
Ut min + e
Uv Uellmin
Üt - max +
+
+
0o min +
@r Imn +
|mollmax +
Tr Imax +
Applying the above estimation, we obtain
T
Σ
AAr Tm
1
10% mmin
110%0 Ilmia +
| max +
v |10% mm + 6
110%
Ilmin +
V
1100 Imax + E
10
N-1
2
Σ
2
2
2
1.0
m
10
Ilmax +
+ 1,0 UFuti
Immi +
|
N.T
||max
2
2
V8
|
min
Where the last inequality is because ++110 n+11
m T. Similarly, obtain
Σ n=1 1=1
A 1
2 E
Which completes the proof.
Lemma 2.
3 fim n.t-1 Tn
81 31
By[l mm G,
Σ
E,
112
<
PT
E, [1 max] s G2
B
Proof.
[[(1
81) Σ
8
s (1 B1) Σ
Taking expectation, obtain
E,
E.Ilmm
R.d |
S
(8)
28

---

From equation (8), we have
$ m1
n.1-1
31
81 G²
7m
11
) ml 1mm.t-111" 112 1
81
B1
31
We get the bound
m n.t-1
B1
Σ
Et
112
31
31
1=1
Next, we bound ||
=I((1 - 82)
Σ
83
(g 2
S - 82)
$2 1 C s G
1=1
1=1
We obtain
22 Fm
maxx S G
Corollary 2. Suppose Assumptions 1, 2, 3 are satisfied. When the learning rate meets the condition /E 1 inner iterations n-th block epoch satisfies: 4Cos. C2T.S.
that min {v =
Lo?Lo²
To CO
[x f(x;--]]
B.ll Ta [(07+
2Co 1=0
82
2a 2a CR
+2a(1 + C)
Σ
E, [|ArT/
+
(9)
1=1
Proof. First, define series {-",]]
0, ...,
31 m
nht Ta
31
Noting that for & = 1,2,..., T, we have
m.. n.0
0, n.0 0m.0 gut gu.t-1
arl m.t n.f
m
(10)
We will prove the following recurrence relation
31
al al n.t
aA
m.t
31 m
n.t1 t= 1,2, ..., T
(11)
Tm
For notational convenience, we denote Oc, Ht, m.li. as On git mi.d Inm V (0 be-
low.
From the definition of (a we have
31
24 0, -al, B1
m
31
- -alime - al'
m
1- Bi
- -at 1- B1
m4
B1m- + (1
B1) gt.
0, aTr
1- 31
B1
0,-1 aTig.
at:
B1
29

---

Noting that
$1
01-1= 01-1-00---0(-1
1 B1 m4-1
We obtain
B1
Ft 4t-1 alig: aAl;
m
31
Next, we apply smoothness assumption of the loss function f over the iterates x7, t = I,.., T.
f(x)) - =(==-)++ xe-1) + I1 xe-1112.
Taking expectation, we obtain
- aB, (Vo) ( ([-))--))
+ a² L
Figt AP: (
81
Er
1- T- m
Expanding further, we have
Ex[f(x:)] - [x[ 5 -oE, ((V (()-))) P(gii) (I)
aEE, (V)(()-1) ( rigi)l
(II
olE, ( (f( A1;
(1 (1- 31 -
(III)
a² L E,
|riss Ali
B1
|| (IV)
Mc-1
1- Bi
We will estimate each part step by step.
Bounding term
== aEE I (0, (10
OR (0.-1) (1 [ [|]]]]]]
-admin (f--1) E, (0.-))) + aG A[ll
E, [|[ (0.- 1| a [ ||]
Co
Where Co
Amin (T--1,, the first inequality is because assumption 2, the second
vC
inequality is because lemma 2.
Bounding term II
Il = aE, (V (Vf((,))) ffx I). (17 - Vf (x--1) Ar.9))
< alE (V/ (0. -
|1T--1
[3ime-1
||| |1|ATi9.1|
1 B1
T2C
VV (0 (0111)131
E,
V (0.-1) - (+
E, ||AAi
2s
2p
Vē
ag
E, [[ (8.- 1] +
L2
E,
|T,_1]
(81me-1)
113
d²C,
E, |||A|||]
2.
2p
1- 8,
Ve
C2LL22
a LC2
s ap
E. [[ (011))11]
E, [|AA.]]
2pe
30

---

where C1 B
I- p is positive constant, we used the fact |
1
S Vea the first
11 +e and third inequality is because assumption the second inequality is because Sauchy-Schward inequality.
Bounding term III
B1
I!I = oE,,
(f(x Ar;
m
1- B1
[
mc-1 )>
+ aBe { (V((()-)) Vf (0)-1), Ar,
8 8,m4-1
S alE,
AT
1 B1
1 - B1
$ ol,,,
IVV (0.- (00-1))llAA)
311--111
||
+ a. I,E. |1
[31m4=1]
|||AT;
31m
1 31
1- $1
1 31
a² C
S a [|
E, ||AA|ll]
Where the third inequality is because assumption 1. Bounding term IV
IV = Q²T, 2 Es
(3.
)
12
Tig: ATe
(1- 31 Mc-1
11
$ a² LE + + LE
Ari
m
1 31
2L
11me-1
Ee [llgr VS (0(-)) + V (0t-1)) [)]]] + LB2
|
113
3
B1
(E. [|V (0:-1) | + - - a² CLLGEE |
|12/711
a² ' B. |IV (0:-1))1 To? +o -o' ||
Where we used assumption 3 that Se is unbiased with bounded variance o². We obtain the bound of each part
I S Co F.
[[V (0.-1)113 C.
ap E. [ (0.1))
II
E, [[Aill]]
2pe
Ve
a L2CLLG2 VE
III < aC G2 AM
E, IIANII
IV V a²1 R, [IV] (0.-1)
L= &
+ a a GIG [[ATT]]]]
Substituting each part into the above single-step step estimation, we obtain
Eeff(x;) f((,-))) (-G +
aL
ap
L²Lo² E.([V5()-=))1++
+
+ 2ps
8
2:
E
+ (((L+G))C2 + a CI)C
E, +|]A
V8
Updating T times, we obtain
(- (-co Co 2:
a T ap)
T Ta20f1202
E,f((r) f(xo)]
+
Σ
+ 2pe
31

---

T + +(a(L+C)00> + a F([+G))GGG
T E. [| + a C& E2 |||!!1|?
Σ
ve
Choosing 206 S min 1éb1
J Cib)
applying lemmal, we obtain:
T-1
[((0)) + To + TaLCL
(((((T)) -(xo)] S
a
E4
20o
62
t=
T E. ||]
Co2CIGG.
+2a(1+C1)G*
Corollary 3. For each inner iteration index 7
= 0, 0,1,..., at n-th block epoch, the expectation of
1 (12 is bounded by |V r. (0%_2))
12
E2 V vn 5 (0n - - 22
E, [[ f (6%:7)
(12)
Where C2 21.G
Proof.
E, r.. f (17]
< 2E, [|V T. f + || n. f (0m - Vr. (0n.21) ]²{]]]]]]]]
< 2E, [ r. (0% || 2E, V f (0 - v n. 12
2E, | n.. (0% |1 +2L B
110 0 ||?
Where
TT
11- y....
0--1,T 1.T 0n. 0--1.T
1.T ... 0z 1,T
gn.u
1,7
0--11
17
0
0Th 0--11
1. ... 0 1-1
the second inequality because the partial derivative vector V...f (0m:21) ) is part of the gradient vector v (0%)1)
Next, we will bound E, [Iom.0 - 11]
T
E, |10m00 - - 1|2
= E,
r m n.j 1|
j=1
T
E,
[l
8
T2 G2
T²a²G²
Where we used assumption 3 and lemma 2, the first inequality is because Cauchy-shwarz inequality. We obtain
2T-
E. T.. (0m-0) [2] < 2 [[[ f (0 )|2]
8
Let Cg 2L
we have
E, V r, f (12)
2T
< E, T..ff )1.
2
32

---

E.2 Proof of theorem 1
Theorem 2. (MISA Convergence) Suppose that Assumptions 1, 2, 3 are fulfilled, learning rate
satisfies a & min
c
MISA converges at the rate of O VNT
x))
4Cot. CIL)
Σx
Ex 1(0
4Cov Do(TLo Lo² + 2L.2
1-
81 )21
+ (
N
VN
Tave
Where Do = f(x110) f(x²), C1
81 1-
Co
VC*+F"
Proof. From algorithm 3, we have
31 In.T m
gn+1.0 8T. 0
B1
From the definition of the sequence { f in (10), we obtain
[f 1.]]] [g = E, [f()"))))
Which implies
[. = [x(f(==")))
(13)
Substituting (12) and (13) into (9) obtain
aT [, - [(;2,01 <
Ta
E. F.. )1+4
4Co
6²
T
2a CI
T²TOC2
+2a(1 + C
E, ||AA|l| Σ
4Co
It's obvious that the norm of full gradient is equal to the sum of individual block's gradient norms.
IV (0 11
Σ V (1=
Taking expectation, have
E
V (g (1
Σ E. [ (Bn.T) [|2]
Using the conclusion of Corollary 2 and taking expectation with respect to Tri we have
B Σ
Eg,ra [f(x==1.0) - -f(x000]
ph
[f(
1.00) - ff(=)
6=1
aT'
B
Lo +
TOLCL
p.E ||| | (Bnu))
4Co
22
b=1
202CLL
TTT
21 + G2
pm
E, ||
4Co
To
TaSCIIZCUGR
S
4Co
B 20 + C1) Σ
2a²O
P Σ E, [|
E
4Co
33

---

of
Lo
Ta
=B. (
4Co
8²
CLL
E. | 2a(1+ + C1)
E
4Co
Where bn arg max { {Σ E || {P,)} bE
is the sampling probability of each block at n-th block epoch, the first inequality is because Corollary 2, the second inequality is because Corollary 1.
Let Do
C3L2C,G" H DI = Lod D2 = x² 4Co
D3 = 2(1 + C,)G2, D4 = 2C9LC2
Ds
Ca. we
400
obtain
-102)) - 3 To aD3
E2 [[Ar m
+ D&a² T
Updating the above equation for N times, we obtain
[x f(x110]]
aoDoo
E.
V f(0m))))) N +
+ aaD
MI E. | by
+ +N + +NDga
From lemma 1, we know
N
T
2
E,
bn |]
Ve
Let Dg = (f((*) )x, we obtain
--=1 B.ll| V f(0mm0))
DB
aD1
a²D²
2D3
aD4
T
+ Do
N
aNTD.
Do
NT
TDo
Do
Do
TD1 + DA
(D2 2
a
NTDD,
TDo
Do
Choosing
Do N.TD:+D.)
get the conclusion
E,
DVD +
+0 1
N
VNTDo
4Cov To + 2LG2 B
)
T-
+0 )
Tave
F
Memory Analysis
The memory analysis presented in this study is conducted under frozen embedding layer and language modeling head layer configurations, with memory quantification comprehensively defined as the cumulative composition of parameters, gradients, optimizer states, and intermediate activation.
Notation. We suppose that the model follows the standard transformer architecture, where each layer consists of six modules, i.e., We E R Wx E R Wv E R Wo
ghxh W1 E h
W2 E IR 11 Let L denote the number of the transformer layers, and
34

---

denotes the number of attention heads. is the vocabulary size of the model. refer to the sequence length. b represents the training batch size. The embedding representation dim is
Table 13: Parameter dimension
Parameter
Property
Memory consumption
X
Activation
8 h bsh
Wo Q
Model parameter
(h * h/a) h²
Activation
a (8 (s* h/a) * bsh
WR
Model parameter
a* (h hh/a) 112
K
Activation
8* (h/hh/a) = bsh
WV
Model parameter
a+ = h
V
Activation
a (*(#h/aa) b = bsh
S
Activation
a a*s*s*b = abs²
O
Activation
a (*(sh/aa) * bsh
WO
Model parameter
a+(h+//a) h
Activation
s*h#b bsh
Z
Activation
s*h b bsh
Z1
Activation
** b = 4bsh
WI
Model parameter
h 4h 4hi2
Iww++b>0 +
Activation
8 4h b = 4bsh
W2
Model parameter
4h 4h2
F1 Memory analysis of layer-wise optimization method
In this section, we will give detailed memory analysis for layer-wise optimization methods in backward propagation based on LLaMA's standard architecture.
Lemma 3. 0f 00
91 (3% U).
as 04
01 85
y2
Where Layernorm(Us (U), softmar(A).
Proof. First, we compute aʃ
3U The computational flow of =
can be written as
1
=
U.
h
1=1
+8 (U., 2
4
Taking the partial derivative U, we obtain
as 1 05
0s 0Zm
Ug - oj 0f of
m1))
(Us
aUg 03 0Zg
dZa
k=
We observe that the computation of af ø0
depends solely on 7Z and U. For notational convenience, we
denote 0f as
0f 0Z
U)
we compute
0f ĐĄ
35

---

Let Sty
Any
we have
Ei-1
0, if m +
Aun cnis
0Sig
= Sin if m = i and n +
(Eam exax)²
AAm
ss (E e
Sij (1 -1- (1)S:)). if m = and n =
E c m
We obtain
0s
0s (-Si
0ʃ
Σ
Sig (1 - Sig)
AAi
USik
iSi
kØj
0f
0f Sik
Sij
Σ
iS
i5
k=1
Which indicates
05
0f
s)
0A
92
0s
Below, we present the forward and backward propagation flows of certain transformer layer.
Q XWQ
af aW4
X 0s 8Q
K XWR
85
X 0f 8K
0W*
V = XWv
0X
0f
01 W
AMR
Of
X
AV
of
A OK
0f
8f
KT
8f
of S DA
dah
DQ
$ = soʃ 1ofmmar(A A)
04
01
922
0S:
01
Forward
0 = SV
03
01
70
0f
VT 01
0V
ST
a1 50
Backward
Attn = OW
01
WT 8f O"
Df BAtln
00 JAttn
OW
U Attn + X
01 7Allh
00
LoyerNormm(U)
00
05
91
(Z.UU) U) 00
01
Zi ReLU (ZW; b2)
0f
af
L2w w. >0 WT, OW,
af
ZT
Of 87
87
87.
Z2 2,W2 + b2
85
81
WT.
8f JW.
ZT
81
J2,
T2,
JZ.
Next, we will analyze layer-wise optimization methods' memory based on the forward and backward propagation flow a transformer layer.
Activation memory of frozen layers. To enable backpropagation, the activation gradients of Transformer layer must be stored, including the tensors Q, V. K, S, and Tzw The total activation memory is calculated as: 3bsh + abs² + 5bsh abs² + 8hsh.
Activation memory of activated layers. In the layer-wise optimization method, the weight matrices of unfrozen layers Wo, WK, Wv, Wo, W1, W2 require computation, necessitating the storage of X. Q, K, V, O, S. U, U,I2wi+bi>0- Z and Z, The total activation memory here is: 5bsh + abs? 10bsh abs² 15bsh.
Parameter memory layer-wise optimization method. Irrespective which Transformer layer is activated, the model parameters must always be stored. This includes the weight matrices of full
36

---

Transformer layer Wo, WK. Wv. Wo. W1, W2. The total parameter memory is: L(4h² + 8h?) 124
Optimizer state and gradient memory of frozen layer. For frozen layers, neither optimizer states nor gradients need to be stored, resulting in zero memory usage for these components.
Optimizer state and gradient memory of activated layer. Activated layers require storing both gradients and optimizer states. The total memory for this 3 x 12h 36h².
Table 14: Memory of layer-wise optimization method
Layer State
Parameter Memory
Activation Memory
Optimizer State and Gradient Memory
Activated layer
12
abs² + 15bsh
36/²
Frozen layer
12
abs² + 8hsh
In practice, assume that the activated layer's ID is we obtain the following conclusion
Parameter: 12h T
Optimizer State 24h?
Gradient 12h?
Activation (L- - i)bbb? + 8bsh) + (abs? + 15bsh)
Total: (L - ()bb² + 8bsh) + (abs² + 15bsh) + 12h + 36h
The rationale for the activation memory is as follows: when the activated layer ID is i, backpropagation only needs to propagate to the i-th layer. Consequently, the activations of layers j 1,2, 2,.,... can be saved.
Based on this, we derive the peak memory of the layer-wise method.
Laass 8bsh) 7bsh + 12h L 36h
F22 Memory analysis of MISA and LoRA
In this section, we focus on the scenario where only a specific module within a layer is activated each step. Based on this, we derive a memory comparison between the module-wise optimization method and LoRA.
F.211 Memory analysis of Module-wise BCD
Based on the computation flow described above. activating different individual modules incurs additional memory consumption compared to the module-wise optimization method.
Table 15: Memory of Module-wise BCD
Unfrozen Module
Extra Storage
Extra Activation Memory
Extra Optimizer State And Gradient Memory
We
X, We
bsh
3h2
WK
X, WK
bsh bsh
3h
Wv
X, Wy
3/.2
Wo
O, Wo
bsh
3h2
W1
W,
bsh
12h
W2
Z1. W2
4bsh
12h2
Let us consider the case where the activated module is Wo as an example. yields the following components:
The memory analysis
Activation memory of frozen modules
37

---

From the layer-wise memory analysis, we know that for frozen layer, the activation memory requirement is abs² + 8bsh.
Activation memory of the activated module
To enable backpropagation through WQ. we need store the additional activation X, resulting in a total memory of abs² + 9bsh.
Parameter memory in layer-wise optimization
The layer-wise approach requires storing all parameters of each layer, amounting 12 L memory.
Optimizer state and gradient memory for frozen layers
Similar to the layer-wise optimization case, the module-wise approach requires no additional memory for frozen layers.
Optimizer state and gradient memory for the activated layer
Only the optimizer states and gradients for Wq need to be stored, requiring 3h memory.
The peak memory consumption for the module-wise optimization method, when WiQ is the activated module, is therefore:
Laass* 8bsh) + bsh 12h L 3h
F.2.2 Memory analysis of LoRA
Consider weight matrix Wq with LoRA adaptation, which can be expressed as:
We = Wo + BqAq, where BQ E Raxr Aq E B
The backward propagation gradients are computed as:
OC dc DW A
0C B DWQ
aAq
This formulation leads the following memory requirements:
Embedding layer activations:
Requires storing parameters size sh.
Adapter gradients and optimizer states: For each layer's adapters Ag and Bo, this requires 3aL 2hr = 6hr memory.
Full model parameters:
The base model parameters require 12 memory.
During backpropagation through Wo. LoRA needs to store additional activations X, Aq, and Bo, requiring L(abs² + 9bsh + 2hr) memory.
Activation storage:
The peak memory consumption when applying LoRA to We is therefore:
L (abs² 2 + 9bsh 12h² + 8hr)
From the analysis above, we can similarly generalize the memory storage of different modules.
F.3 Memory analysis of GaLore
Consider activating WQ as example, GaLore maintains projection matrix Pq Rahr for Wa and optimizer states in the projected subspace of shape x h. It holds that:
GaLore needs store gradient, projection matrix and optimizer state of each layer's Wo- which needs 4hr L memory storage.
Gal Lore needs to store full model parameters, which needs 12 memory storage.
For backward propagation, Gal.ore needs to store extra activation X, leading to L(abs² 9bsh) activation memory in total.
38

---

We obtain the peak memory of GaLore when the activated module is Wq:
L(abs* 9bsh 12h + 4hr).
From the analysis above, can similarly generalize the memory storage of different modules. Then we obtain the comparison between LoRA, GaLore and Module-wise BCD.
Table 16: Peak memory comparison between LoRA, GaLore and Module-wise BCD
Activating Module
LoRA
GaLore
Modulewise- BBCD
L(abs² + 9bsh
L(abs² + 9bsh
L(abs² + 8hsh
Wq
+12h + 8hr)
+12h² + 4hr)
+12h²) + bsh + 3h²
L(abs² + 9bsh
L(abs² + 9bsh
L( abs² 8bsh
WK
+12h² + 8hr))
+12h² + 4hrr
+12h²) + bsh + 3h2
L(abs² + 9bsh
L(abs² + 9bsh
L(abs² + 8bsh
Wv
+12h² + 8hr)
+12h² + 4hr)
+12h²) + bsh + 3/2²
L(abs² + 9bsh
L(abs² + 9bsh
L(abs² + 8bsh
Wo
+12h² + 4hr)
+12h² + 8hr)
+12h²) ) + bsh 3h2
L(abs² + 9bsh
L(abs² + 9bsh
L(abs + 8bsh
W1
-1262 + 20hr)
+12h² + 13hr)
+ +12h²) + bsh + 12h?
L(abs² + 12bsh
L(abs2 12bsh
L(abs² + 8hsh
W2
-12h2 + 20hr)
+12h + 13hr)
+12h² 4bbsh + 12h?
L(abs² 4 15bsh
L(abs + 15bsh
L(abs" + 8bsh
All
12h + 72hr)
+12h2 42hrr
+12h² 7bsh 36/²
F44
Memory analysis of MISA
Let æ, y. and denote the number of activated modules for Wo WK WWyW W1, and W2, respectively, Based on Table 15, we derive the memory estimation for MISA as follows:
Activation Memory
Compared layer- wise optimization, MISA requires additional activation storage of:
bsha bshy 4bshz
Parameter Memory
The full model parameters require:
121
Optimizer States and Gradients
The memory needed for storing optimizer states and gradients of activated modules is:
2/ + 12h? 12h²=
Given trainable parameter ratio threshold 8, where the number of trainable parameters is 12h² L8, we formulate the optimization problem to determine the maximum memory consumption M:
Constraint:
h² + 4h2 + 2) S 12
Objective:
M L(abs² + 8bsh + 12/2) + bsh(z + y + 4z)
+ h3r 12y 122)
The optimal solution yields the peak memory consumption of MISA:
(abs² + 8bsh + 12 12bshd + 36/ 28)
(14)
F.5 Memory analysis of MISA's indicators
In this subsection, we denote B be the total number of modules.
According to equation 4, MISA needs to maintain G pp.
39

---

Memory of Gp: For the frozen block bo, Gj follows the previous step. For the unfrozen
block b, since its gradient has been saved, only need store + Σt "12 which is negligible. Therefore, the total memory overhead is O
Memory of ph: According to the definition of P in equation 4, the previous P% pp-1 can discarded at n-th sampling period. Therefore, the memory occupation O(B).
Therefore, the storage of the importance indicator O(B) is negligible compared with the gradients, which is L).
Memory efficiency even as model scales. The memory overhead of storing the smoothed historical gradient norm G. is O(B), which is negligible compared to the O L) overhead the gradients themselves (with relative ratio of O(1/h²) when L O(B), in LLaMA models). As model dimensions and L increase, this ratio decays quadratically, making the overhead of storing G. negligible.
E.6 Peak memory comparison between MISA and layer-wise method
Lemma 4. Let 8 be the trainable parameter ratio threshold, when Ths+ 36A memory efficient than layer-wise method.
MISA is more
+2hhL
Proof. From conclusion (14), we only need to prove
L(abs + 8bsh + +2/? + 12bshd + 36h2d) 6) < L(abs + 8hsh + 1242) + 7bsh + 364
Ths + 36h
8 <
12bsL + 36hL
It's obvious that when 1, memory of MISA always smaller than layer-wise method.
F.7 Peak memory comparison between layer-wise method and LoRA
Lemma 5.
As sequence length gets larger, layer-wise optimization method will become more memory efficient than LoRA and Galore. When 36h-72rL peak memory of layer-wise optimization method is always smaller than LoRA and Galore.
76L-7b
Proof. In this proof, we only consider the case when LoRA target and Galore target are all modules. From table 16, when peak memory of layer-wise method is smaller than LoRA and Galore, have L ( abs² + 8bsh + 12h²) + 7bsh + 36/² < L abs² + 15bsh + 12h + 42hr)
1
7bsh 36h² 7bsh + 42rhL
36h - 42rL < s (7bL - 7b)
1
36h 42rL
8
7bl. 7b
F.8 Paras/Peak-Memrry comparison between layer-wise method and LoRA
Lemma 6. In the memory-consistent case, layer-wise optimization method can update more parame- ters than LoRA when the inequality holds 3rL
Proof. From table16, we only need to prove
18rhL.
12h.2 < L aabb² + 8bsh + 12h2) + 7bsh + 36h
L abss + 15bsh + 1242 + 72rh)
40

---

3r L (abs² + 8bsh + 12h2) + 7bsh + 36/²] < 2h (abs² + 15bsh + 12h² + 72rh)
abs² (2h Lr bsh [30h 24rL 21r] + (36r 24h 36r
When
2h 37.
holds, each part of the last inequality greater than 0, which completes the proof.
G Computation Analysis
In this subsection, adopt the same notation and setting as those in the memory analysis. It is well known that subspace optimization algorithms,, such as LoRA, are not computationally efficient. Therefore, in this section, we will only analyze the computational complexity of layer-wise and module-wise optimization methods. Here we present the computation analysis in detail. We will analyze the exact floating point operations (FLOPs) required for backward propagation through single Transformer layer. To begin with, we should remind that the matrices production with size m X n and r need 2mnr FL.OPs.
From the forward and backward propagation flows and analysis in F.1, We can estimate the computa- tional overhead under layer-wise setting.
Let:
d: batch size
sequence length
h: hidden dimension
a: number of attention heads
d h/a: per-head dimension
Wq. WK. Wv E B
Wo E TAXA
W1 hx4h
W₂ W
Backward FLOPs per layer
1. Feed-Forward Network (FFN)
Grad @W2 2bs 4h = 8bsh2
Grad 0Z1 2bs h 4h = 8bsh?
Grad @W1 2bs h . 4h = 8bsh2
ReLU mask 4bsh
Total: 24bsh² + 4bsh
2. Output Projection Wo
Grad @Wo 2bsh? 2
Grad d0 2bsh2
Total:
4bsh² 2
3. Q/K/V Projection
Each of @We, 0Wk, oW
2bsh2
Total
6bsh2
41

---

4. Attention Score and Softmax
QK 2
Grad w.r. Q, K: 2bs²h
Softmax backward 2bas²
Total
4bs + 2bas
5. Attention Output Context (SV)
S 80 h
asv
2bs² h
Total:
4bs
6. Layer-Norm Backward
10bsh
Total Backward FLOPs (One Layer)
34bsh + +bb + 2bas + 14bsh
G.1
Computation analysis of layer-wise optimization methods
We assume only one layer unfrozen each iteration under a block coordinate descent (BCD) strategy.
Activated layer
For an activated layer, cannot omit the computation cost of the entire layer during backpropagation. Thus, the total FLOPs without weight gradient computations is:
34bsh2 8bs2 h. + 2bas + 14bsh
Frozen layer
For frozen layer, omit part of the computational cost of calculating the weight gradients during backpropagation.
We exclude the following terms:
We, WK, Wy: 6bs/1
Wo: 2bsh?
W1, W2: 16bsh?
Thus, the total FLOPs without weight gradient computations is:
10bsh 8bs2 + 2bas + 14bsh
Total computation cost
(L - 1)(10bsh2 + Shs + 2has + 14bsh) + 34bsh? + 8hs2h + 2bas + 14bsh
42

---

G.2
Computation analysis of module-wise optimization methods
Let the total number activated modules Wo. WK. Wv. Wo be æ, number of modules Wi. W2 be y, the expectation of computation cost be C, from analysis of G, obtain the simple optimization Problem.
1/ 4h < 12h2
C L(10bsh 8bs"h + 2bas 2 + 14bsh) + 2bsh + 8bsh"
We obtain
Cmaz = T.(10ssh² + 8bs2h h + 2bas² + 14bsh) + 24
It is obvious that when 8 < 1. module-wise methods demonstrates greater computational efficiency compared to layer-wise optimization methods.
G.3
Computation analysis of MISA's indicators
In implementation, first update Gā and then compute Pğ according to Equation (4).
Computation overhead of Gg: The dominant term of the computation cost the calculation for li 1995"11 112, which needs (1h) flops.
Computation overhead of p:: The computation overhead of p only depends on the number of blocks B, which needs O(B) flops.
Therefore, the computational overhead of the importance indicators O(B+h²) is negligible compared with the gradients, which is (bsh?²)
Computational efficiency even as model scales. The computation of the smoothed historical gradient norm GG introduces an overhead of (hh)). In contrast, the computational overhead for gradients is Obbsh", yielding a relative ratio of O1/bbs. As the model scales up, (i.c., with increasing h and L), this ratio remains stable, making the overhead of computing G, overhead negligible.
H Broader Impacts
This paper presents work that aims advance the field of machine learning. The proposed method, MISA, is designed to efficiently train LLMs with low memory consumption. This will facilitate further research regarding memory-efficient LLM training in the future. We believe that this work will not cause significant societal consequences.
Experimental Hyperparameters
For all baselines and our proposed method, we conducted extensive hyperparameter searches. The learning rate searched in (2e-4, le-4, 5e-5, le-5, 5e-6, 3e-6, le-6). For LoRA and GaLore, we explored ranks {8, 16, 32), and we found that a rank of 16 or 32 consistently yielded better performance than For MISA, rj was searched in (0.1, 0.5, 1). The table below presents the optimal hyperparameter settings.
The batch size in the following tables is represented miero batch size gradient accumulation.
43

---

1.1 Hyperparameters for commonsense reasoning
LLaMA3-8B
Qwen2.5-7B
Hyperparameters
LoRA
DoRA
LoRA
DoRA
Rank
32
16
32
16
64
32
64
32
Dropout
0.05
Optimizer
AdamW
Learning rate
le-4
le-4
2e-4
2e-4
Batch size
4x4
Warmup Steps
100
Epochs
3
Target module
W WR, Wo, W Wdown
Table 17: Hyperparameters LoRA and DoRA in commonsense reasoning tasks.
LLaMA3-8B
Qwen2.5-7B
Hyperparameters
MISA
BAdam
LISA
MISA
BAdam
LISA
Optimizer
AdamW W
Learning rate
le-5
5e-6
3c-6
5e-6
5e-6
3e-6
Batch size
4x4
Warmup Steps
Activated parameters
1%, 3%
layer
layer
1%, 3%
layer
1 layer
MISA's
1
T
50
Epochs
Table 18: Hyperparameters of MISA, BAdam and LISA in commonsense reasoning tasks.
1.2 Hyperparameters for math reasoning
LLaMA3-8B
Qwen2.5-7B
Hyperparameters
LoRA
DoRA
LoRA
DoRA
Rank
32
32
32
32
32
Dropout
0.05
Optimizer
Adam W
Learning rate
2e-4
2e-4
le-4
le-4
Batch size
4x1
Warmup Steps
0
Epochs
3
Target module
W W u. Wap.
ups W Wdown
Table 19: Hyperparameters of LoRA and DoRA math reasoning tasks.
44

---

LLaMA3-8B
Qwen2.5-7B
Hyperparameters
MISA
BAdam
LISA
MISA
BAdam
LISA
Optimizer
AdamW
Learning rate
5c-6
5c-6
5c-6
le-5
le-5
5c-6
Batch size
4x1
Warmup Steps
Activated parameters
1%, 3%
layer
layer
1%,
3%
I layer
layer
MISA's
0.5
T
50
Epochs
Table 20: Hyperparameters MISA, BAdam and LISA in math reasoning tasks.
I.3 Hyperparameters for instruction fine-tuning
LLaMA2-7B
Mistral-7B
Tiny-LLaMe
Hyperparameters
LoRA
GaLore
LoRA
GaLore
LoRA
GaLore
Rank
32
32
32
32
32
32
64
64
64
64
64
Dropout
0.05
0.05
0
0.05
0
Optimizer
Adam W
Learning rate
2e-4
3e-6
1e44
1e-6
2e-4
1e-5
Batch size
2x8
Warmup Steps
0
Epochs
3
Target module
Wg, W., W., W., W up, W Wgntes Wdown
Table 21: Hyperparameters of LoRA and DoRA Alpaca GPT4.
LLaMA2-7B
Mistral-7B
Tiny-LLaMA BAdam
Hyperparameters
MISA
BAdam
LISA
MISA
BAdam
LISA
MISA
LISA
Optimizer
AdamW 5e-6
Learning rate
le-5
le-5
5e-6
le-5
le-5
Se-5
le-5
le-5
Batch size
2x8
Warmup Steps
0
Activated parameters
3%
luyer
layer
3%
layer
layer
4.545% 0.5
layer
layer
MISA's
0.5
T
50
Epochs
3
Table 22: Hyperparameters of MISA, BAdam and LISA Alpaca GPT4.
45

---

1,4 Hyperparameters for pre-training
Hyperparameters
MISA
Hyperparameters
GaLore
Optimizer
A
Optimizer
A
Learning rate
1e-3
Learning rate
1e-3
Batch size
32x256
Batch size
226
Warmup Steps
0
Warmup Steps
MISA's 6
3%, 25%
Gal.oee's 's rank
32, 256
MISA's 77
300
GaLore's
T
50
T
50
Sequence Length
256
Sequence Length
256
Training Steps
52000
Training Steps
52000
Table 23: Hyperparameters of MISA pre-training
Table 24: Hyperparameters of GaLore pre-training LLaMA2 350M.
LLaMA2 350M.
46

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
See the abstract and introduction part.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discussed the limitations in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
47

---

Justification: See Appendix E.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Hyperparameters are Appendix The code available the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
48

---

Answer: [Yes]
Justification:
All the code and data are available in the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https//niipscccppbbicc guddes/CodeSubmiisiooooool
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: //nips.cc/pp //nipsccc/publl////eecc wwwp://pbbiciguiddd//ooeeeeeeeee SeaedeeeeeeeeeSuuuiiiiooooioio for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provided all our experimental settings in Section and Appendix I.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We conducted extensive hyperparameter searches for all baselines and our proposed method. Experimental results are tested multiple times to ensure stability and reliability. See Appendix I.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
49

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See 4. Experiments are conducted on RTX 4090 GPUs.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:://eeeiiss.ccppubbeettissss
Answer: [Yes]
Justification: We have reviewed the code of ethics and confirmed that our research complies with its guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Appendix H.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
50

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper presents work whose goal is advance the field of LLM optimiza- tion. All the datasets and models are public and we have checked them all. Our paper poses such risks mentioned in the question.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We have cited the original paper that produced the assets properly.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
assetts are released, the license, copyright information, and terms of use package should be provided. For popular datasets, paperswithcode, seterrswihhcode.comdataaaeee has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
51

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: All the code are available in the supplementary material, along with explicit README.md.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
52

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs not used in our core method development.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non non-standard components.
Please refer to our LLM policy http://lourr..cccooottrro///ss)) for what should or should not be described.
53


Reviews:
[{"review_id": "y6jDb9Wgnt", "reviewer": "Reviewer_ULBr", "summary": "This 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.", "questions": "* 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?", "limitations": "* 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "The authors addressed my main concerns regarding profiling, theoretical justification for the importance metric, and ablation studies on partitioning granularity during the rebuttal. Thus, I would like to increase my score.", "rebuttal": "We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.\n\n**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.\n\n|Forward|Backward|Adam update|Sample modules|Track block gradient norm|Update sampling probability|Additional Momentum Update|Clear optimizer states|\n|-|-|-|-|-|-|-|-|\n|11.433|24.359|0.591|0.010|0.992|$\\le10^{-5}$|0.012|0.002|\n\n**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. \n- **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. \n- **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.\n\n|Model|precision|Avg. Step Time (ms)|Grad. Norm Time (ms)|Total Memory (GB)|Grad. Norm Memory (GB)|\n|-|-|-|-|-|-|\n|Qwen2.5-7B|bf16|37.40|0.99|15.68|$\\le 10^{-7}$|\n|Qwen2.5-7B|fp32|126.25|1.96|35.40|$\\le 10^{-7}$|\n|Qwen2.5-14B|bf16|78.09|1.89|34.08|$\\le 10^{-7}$|\n|Qwen2.5-14B|fp32|260.50|4.15|65.98|$\\le 10^{-7}$|\n\n**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.\n- According to the derivation in Line 130, the decrease in loss function can be lower bounded by \n$$\\mathbb{E}[f(\\theta^n) - f(\\theta^{n+1})] \\ge \\sum_{b} p_b \\\\|g_b\\\\|^2.$$\nTo 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**.\n\n- 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.\n\n**W4. Module partition.** \n- **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.\n\n- **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.\n\n- **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.\n\n|Granularity|Math|Commonsense|\n|-|-|-|\n|Module-wise(MISA)|**73.6**|86.6|\n|Attention and MLP(coarser-grained)|72.5|85.4|\n|Head-wise(finer-grained)|73.1|**86.8**|\n\n**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.\n\n**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.\n\n**Q4. Sensitivity to gradient norm.**  We thank the reviewer for this insightful question. \n- **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.\n\n |Method|MMLU|MMLU-pro|Commonsense|Math|\n |-|-|-|-|-|\n |Weight Norm|64.5|35.9|85.7|71.9|\n |Number of Parameters|63.7|36.2|85.9|72.7|\n |**MISA(Gradient Norm)**|**65**|**36.5**|**86.6**|**73.6**|\n\n- **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).\n\n**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.\n|LLaMA 1B|Adam|GaLore(r=512)|MISA($\\delta = 25\\\\%$)\n|-|-|-|-|\n|Validation perplexity($\\downarrow$)|15.56|15.64|**15.27**|\n\n**L2. Module partition granularity.** See the granularity ablation study in our response to **W4** and **Q1**. \n\n**L3. Detailed profiling.** See our response to **W1**. \n\n--- \nWe 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."}, {"review_id": "YWlnY9Hu6M", "reviewer": "Reviewer_CWR5", "summary": "In 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.", "questions": "1. 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.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The authors have satisfactorily addressed my concerns, both theoretical and experimental. While some clarifications (especially around notation and algorithm indexing) should be better integrated into the main text, the core ideas are sound, and the method is well-motivated. I will accordingly increase my score.", "rebuttal": "We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.\n\n**W1. Terminology.** We sincerely thank the reviewer for highlighting this point. Our definitions are: \n- **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 .\n- **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**. \n- **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.\n\nWith 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.\n\n**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.\n\n**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.\n\n- 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 $\\textbf{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.\n\n- 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.\n\n**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. \n\n**W4.2 Comparison with LISA.** MISA outperforms LISA for two fundamental reasons. \n - **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.\n - **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.\n\n**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:\n 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.\n 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$.\n 3. This constructed distribution ${p_{b,j}}$ is a feasible solution for problem (5) because it satisfies the necessary constraints:\n    - $p_{b,j} >= 0$ since $p_b >= 0$.\n    - 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$\n\nThe 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).\n\n**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.\n\n|Method|Math|Commonsense|\n|-|-|-|\n|LISA|70.7|85.9|\n|Layer-wise + **Importance Sampling**|73.2|86.1|\n|**Module-wise** + Uniform Sampling|72.1|85.7|\n|**Module-wise+Importance Sampling(MISA)**|**73.6**|**86.6**|\n\n**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.\n\n**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. \n\n**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. \n\n**Reference**\n\n[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.\n\n[2]Liu S Y, Wang C Y, Yin H, et al. Dora: Weight-decomposed low-rank adaptation. ICML, 2024.\n\n---\nWe 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."}, {"review_id": "DFeWcpMUYS", "reviewer": "Reviewer_HSph", "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.", "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?\n\nAlso 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].\n\nWhat 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?\n\n1. A Spectral Condition for Feature Learning (2024) by Yang et al.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "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.", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.\n\n**Q1. Cleaning optimizer state.** \n\nWe 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.\n\nOur conjecture for why clearing the optimizer state helps in MISA's pre-training is based on the concept of \"stale\" momentum.\n- 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.\n- 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.\n- 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.\n\nThis contrasts sharply with standard full-model training (e.g., with Adam):\n- In standard training, all model parameters and their corresponding optimizer states (first and second moments) are updated in every single iteration.\n- The momentum is therefore always \"in-time\" and provides a valuable, continuously updated estimate of the gradient's trajectory.\n- In this context, clearing the optimizer state would be detrimental, as it would discard this crucial historical information and disrupt the convergence path.\n\nIn 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.\n\n**Q2.1 Scaled gradient norm.** \n\nWe thank the reviewer for this insightful question.\n\nThe 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.\n\nWithout 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.\n\n**Q2.2 Different normalizations in different layers.** \n\nWe 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.\n\n**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.\n\n|T|5|15|30|50|100|200|500\n|-|-|-|-|-|-|-|-|\n|Validation Loss($\\downarrow$)|0.877|0.874|0.871|0.873|0.877|0.881|0.879|\n|MMLU($\\uparrow$)|46.22|46.23|46.17|46.27|46.19|46.01|45.89 |\n\n**Reference**\n\n[1] Yang, G., Simon, J. B., & Bernstein, J. (2024). A Spectral Condition for Feature Learning.\n\n\n--- \n\nWe 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."}, {"review_id": "bZB0xQCUFk", "reviewer": "Reviewer_ov6T", "summary": "The 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).", "questions": "1. 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).", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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", "was_revised": true, "final_justification": "I maintain my positive score. The proposed method is a valuable contribution to the community and is proven to work via sufficient experiments. Furthermore, the authors have addressed my concerns and questions thoroughly. \nAt the same time, the scale of improvements is not groundbreaking, which is hard for me to assign higher score.", "rebuttal": "We are grateful for the reviewer's insightful feedback and constructive comments. We address each point in detail below.\n\n**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.\n\n|T|5|15|30|50|100|200|500|\n|-|-|-|-|-|-|-|-|\n|Validation Loss($\\downarrow$)|0.877|0.874|0.871|0.873|0.877|0.881|0.879|\n|MMLU($\\uparrow$)|46.22|46.23|46.17|46.27|46.19|46.01|45.89 |\n\n**Question 2. Fine-tune a small set of blocks with MISA** \n\nThe 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.\n\nHowever, 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:\n\n 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.\n\n|Method|$W\\_v$|$W\\_o$|$W\\_{\\text{gate}}$|$W\\_{\\text{up}}$ |$W\\_{\\text{down}}$|$W\\_{\\text{all}}$|\n|-|-|-|-|-|-|-|\n|MISA|70.5|71.7|71.5|72.5|71.7|**73.6**|\n|LoRA|69.9|71.2|71.5|71.1|70.9|71.9|\n\nAs shown in the table, even when only specific individual modules are fine-tuned, MISA still achieves competitive performance compared to LoRA.\n\n 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.\n\nTogether, 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. \n\n**Question 3. Variance due to sampling mechanism** \n\nWe 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.\n|Method|MMLU|MMLU-pro|MT-Bench|\n|-|-|-|-|\n|LISA|46.21 $\\pm$ 0.12|20.85 $\\pm$ 0.09|4.94 $\\pm$ 0.14|\n|LoRA|45.50 $\\pm$ 0.07|20.57 $\\pm$ 0.08|4.45 $\\pm$ 0.15|\n|MISA|46.27 $\\pm$ 0.09|20.69 $\\pm$ 0.16|5.13 $\\pm$ 0.22|\n\nTo 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.\n|Method|MMLU|MMLU-pro|MT-Bench|Math Reasoning|Commonsense Reasoning|\n|-|-|-|-|-|-|\n|Ours |46.27|20.69|5.13|73.6|86.6|\n|Cyclic|46.13|20.64|4.89|72.1|85.7|\n\n**Question 4. Peak memory comparison between LoRA and MISA.** \n\nThe 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.\nWe 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: \n||LoRA|MISA|\n|-|-|-|\n|Peak Memory |$L(abs^2 + 15bsh + 12h^2 + 72hr)$|$L(abs^2 + 8bsh + 12h^2 + 12bsh\\delta + 36h^2\\delta)$|\n|Activation memory|$L(abs^2+15bsh+18hr)$| $L(abs^2+8bsh+12bsh\\delta)$|\n\nAs 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.\n\n**Weakness 1. MISA does not lead to an adapter**\n\nThe 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.\n- 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.\n- 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.\n\n**Weakness 2. Modality model**\n\nThank 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. \n\n**Weakness 3. Non-Transformer model**\n\nThank 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.\n\n**Weakness 4. Small scale experiment**\n\nIn 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.\n\n|LLaMA 1B|Adam|GaLore(r=512)|MISA($\\delta = 25\\\\%$)\n|-|-|-|-|\n|Validation perplexity($\\downarrow$)|15.56|15.64|**15.27**|\n---\n\nWe 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."}]

Author rebuttal:

Meta-review: This paper proposes MISA (Module-wise Importance Sampling Algorithm), a novel memory-efficient optimization method for full-parameter fine-tuning and pre-training of LLMs. Unlike prior layer-wise approaches that freeze entire transformer blocks, MISA decomposes each layer into smaller modules (e.g., attention heads, feed-forward components) and assigns gradient-based importance scores. A weighted sampling strategy then activates modules for updates, reducing gradient variance and yielding greater memory savings than layer-wise sampling. The method is theoretically grounded with convergence guarantees under non-convex stochastic settings and is validated by strong empirical results across commonsense reasoning, mathematical problem solving, instruction tuning, and pre-training tasks.

The reviewers recognize several key strengths of this work. MISA introduces a well-motivated module-wise partitioning of transformer layers, allowing fine-grained memory control and improved gradient preservation. Its importance-based sampling strategy, tied to gradient norms, is innovative and effectively reduces gradient variance. The method is theoretically grounded, with convergence guarantees under non-convex stochastic optimization, including with adaptive optimizers like Adam. Reviewers also praised the clarity of presentation, including logical flow, detailed analysis, and comprehensive experiments and ablation studies across tasks. Overall, the work is seen as methodologically novel, theoretically sound, and empirically strong, with all reviewer concerns addressed in the rebuttal and supplementary experiments. So we accept this work.
Decision: Accept (poster)