Lua-LLM: Learning Unstructured-Sparsity Allocation for Large Language Models

We are glad that the reviewer found that our work is clearly written and well-motivated. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the questions you have raised.

Q1: The downstream tasks can be made more comprehensive by adding common evaluation datasets like MMLU.

Q3: Is perplexity a good evaluation metric for comparisons between different pruning strategies?

A1: We agree with the reviewer that when comparing different types of models, there can be mismatchs between perplexity and downstream performance. However, we think that perplexity remains a straightforward and efficient evaluation metric for comparing different pruning methods on the same model type.

To demonstrate more comprehensive comparisons, we evaluate the 5-shot accuracy on the MMLU dataset for LLaMA2-7B and LLaMA3-8B model pruned with Wanda, OWL, and Lua-LLM, at 60% and 70% sparsity levels. The results (See Table below) show the superior performance of our method.

Table 1: Comparisons for 5-shot accuracy on the MMLU dataset.

Q2: The practical speedup observed is not very large even at 70% sparsity ratio. Could the authors discuss whether this limits the practical usage of unstructured pruning methods?

A2: We thank the reviewer for raising this meaningful discussion for the practical usage of unstructured pruning methods. Our inference evaluation is conducted with SpInfer framework. As discussed in its paper [1], SpInfer faces limitations during the prefill phase when batch size and sequence length are large, which leads to higher memory access overhead for the SpMM operations. Considering the relatively low performance degradation and substantial memory savings offered by unstructured pruning methods, optimizing SpMM operations on GPUs remains an important research topic for future work. Furthermore, future research directions include designing hardware-friendly sparsity patterns like V:N:M [2], which can leverage Sparse Tensor Cores on NVIDIA GPUs without significant performance degradation.

Q4: Could the authors discuss practical trade-offs between structured and unstructured pruning strategies in terms of end-to-end speed up and downstream performances?

A4: Structured pruning removes entire components, such as attention-heads in the Attention layers and neurons in the MLP layers, which can directly generate a smaller "dense" model and achieve significant inference speedup on the GPU Tensor Core architecture. However, it typically suffers severe performance degradation. Even if employing weight reconstruction methods such as LoRA-finetuning [3] and ADMM-based reformation [4], the performance loss is still difficult to compensate, which limits its sparsity levels under 50% typically.

Unstructured pruning removes less critical weight parameters at the element granularity, which offers the most flexibility and yields better post-training performance than structured pruning even at high sparsity levels. Despite the significant performance gains and memory savings, a major challenge for unstructured pruning is to achieve practical speedups on GPUs with substantial irregular memory access overhead, especially for large language models. Sparse inference framework like SpInfer [1] manages to optimize the SpMM operations for GPU Tensor Core architectures, yet its inference throughput is limited during the prefill stage.

N:M sparsity serves as a trade-off between flexibility and efficiency, yet only 2:4 sparsity is supported by Sparse Tensor Cores on NVIDIA GPUs. Moreover, recent studies have focused on designing more adaptive sparsity patterns. For example, DISP-LLM [5] proposes a dimension-independent structured sparsity pattern, which prunes the channels for each weight matrix independently. To enhance the practical usage of model pruning, we will focus on other adaptive sparsity patterns in future exploration.

Q5: How sensitive is the proposed method to the training dataset used (e.g., the size of the training set and its domain/distribution)?

A5: We conduct an ablation study on the LLaMA2-7B model with different size of C4 training dataset across different sparsity levels. The results (See Table below) demonstrate that the performance improves with larger training datasets. Notably, even with only 128 training samples, Lua-LLM demonstrates the ability to achieve superior results (60.82) compared to Wanda (75.42).

Table 2: Perplexity of sparse LLaMA2-7B models on Wikitext2 with different training samples.

We also conduct an ablation study on the LLaMA2-7B model with C4 and Wikitext103 datasets across different sparsity levels. The results (See Table below) reveal that our method on two different training datasets achieve similar performance, demonstrating the robustness and generalization capability of our method across different datasets.

Table 3: Performance of sparse LLaMA2-7B models with different training datasets.

Thanks for the attentive reading of the manuscript and constructive feedback. We hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are glad to discuss further questions and suggestions.

[1] Fan, Ruibo, et al. "Spinfer: Leveraging low-level sparsity for efficient large language model inference on gpus." Proceedings of the Twentieth European Conference on Computer Systems. 2025.

[2] Castro, Roberto L., et al. "Venom: A vectorized n: M format for unleashing the power of sparse tensor cores." Proceedings of the International Conference for High Performance Computing*, Networking, Storage and Analysis*. 2023.

[3] Ma, Xinyin, Gongfan Fang, and Xinchao Wang. "Llm-pruner: On the structural pruning of large language models." Advances in neural information processing systems 36 (2023): 21702-21720.

[4] Shen, Xuan, et al. "Search for efficient large language models." Advances in Neural Information Processing Systems 37 (2024): 139294-139315.

[5] Gao, Shangqian, et al. "Disp-llm: Dimension-independent structural pruning for large language models." Advances in Neural Information Processing Systems 37 (2024): 72219-72244.

We are glad that the reviewer found that our work is clearly motivated and novel. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the questions you have raised.

Q1: Significance concerns: Lua-LLM does not show that it's preferable to use one of these sparse models over a smaller, dense counterpart.

A1: We thank the reviewer for raising this meaningful discussion regarding performance comparisons between large sparse LLMs and small dense LLMs with similar parameter counts.

First, we would like to clarify that large sparse LLMs with unstructured sparsity are often better than small dense LLMs with similar parameter counts on downstream task performance, even without any weight updates, which is also discussed in Section 4.1 of Wanda's paper [1]. To verify the practical usage of unstructured pruning methods without weight updates, we evaluate the average zero-shot accuracies on 7 downstream tasks for the LLaMA2-13B model at 50% sparsity level with Wanda, OWL, and Lua-LLM. The results (See Table below) demonstrate that unstructured 50% sparse LLaMA2-13B with OWL (60.03%) outperforms dense LLaMA2-7B (59.12%), and our Lua-LLM further improve the accuracy to 60.91%.

Table 1: Performance comparisons between large sparse LLMs and small dense LLMs

Moreover, we understand the reviewer's concern that the 50% sparse LLaMA2-13B model has a higher perplexity than the dense LLaMA2-7B model, yet a potential reason for this mismatch is that the perplexity metric cannot reflect the practical performance difference between different structures of models. For example, LLaMA3-8B shows superior zero-shot accuracy than LLaMA2-7B (65.60% vs. 59.12%), but it has a larger perplexity than LLaMA2-7B (6.14 vs. 5.47). Notably, perplexity remains a straightforward and efficient evaluation metric for comparing different pruning methods on the same model type.

Table 2: Performance comparisons with perplexity and downstream accuracy.

Q2: Evaluation is limited to only two model families: LLaMA and OPT.

Q6: Is Lua-LLM robust to different importance metrics?

A2: We conduct an ablation study on LLaMA2-7B (Multi-Head Attention architecture) and Mistral-7B (Grouped-Query Attention architecture) models at 70% sparsity with different importance metrics (Magnitude, Wanda), evaluation metrics (PPL, MMLU), and allocation strategies (with or without Lua-LLM). The results (See Table below) demonstrate that:

(1) When integrated with Lua-LLM, the relative order among Magnitude and Wanda importance metrics is preserved under different models and metrics.

(2) For different importance metrics, model architectures, and evaluation metrics, our Lua-LLM consistently outperforms the uniform pruning baseline Wanda.

Table 3: Robustness evaluation for importance metrics, evaluation metrics, and allocation strategies.

Q3: The regularization term's benefit is not clear.

Q7: What was the final sparsity achieved by each model in Table 6?

Q8: What does it look like if this regularization term "doesn't converge" (line 322)?

A3: We thank the reviewer for raising this meaningful discussion about the effectiveness of the regularization term. If the regularization hyperparameter is not large enough, the regularization loss fails to converge to zero. This indicates that the pruning mask learned by our method does not meet the target sparsity, and the overall sparsity value is actually lower than the target one, since this can lead to better model performance and lower training loss. We consider the case as "doesn't converge", since it leads to an unfair comparison. When the regularization hyperparameter is sufficiently large, the regularization loss rapidly converges to zero, which strictly enforces the target overall sparsity.

Q4&Q5: What batch size and sequence lengths (input/output) are used?

A4: We use the testcase with batch size set to 8, input sequence length set to 32, and output sequence length set to 256. In this case, the inference throughput of SpInfer is not bottlenecked by the prefill stage, and the sparse model integrated can achieve practical speedup compared to the dense baseline. SpInfer faces limitations during the prefill phase when batch size and sequence length are large, which leads to higher memory access overhead for the SpMM operations. A potential solution [2] [3] is to extract dense blocks of non-zeros in the sparse weight matrices and use dense matrix-matrix multiplication to achieve high throughput, and we plan to leave this for our future study.

Q9: Is Lua-LLM inherently incompatible with techniques like SparseGPT?

A9: Lua-LLM is compatible with techniques like SparseGPT. Although SparseGPT iteratively selects block-wise masks and updates the other unpruned weights to compensate the pruning error, which is expensive to integrate into the training process, we can use our Lua-LLM to search for block-wise sparsity, and then apply the sparsity allocation to the pruning process of SparseGPT.

To verify this, we adopt a block-wise mask selection granularity for Lua-LLM, and search for the optimal sparsity allocation for these intra-layer groups. After the learning process, we consider the searched sparsity allocation as the block-wise sensitivity statistics to weight pruning, and apply it to SparseGPT, which uses the searched sparsity to select block-wise masks and updates the other weights. The results (See Table below) show that our allocation strategy achieves superior performance at different sparsity levels. However, the results also demonstrate that the weight update strategy adopted by SparseGPT is not sufficient to recover model performance at higher sparsity ratios, which underscores the importance of mask selection granularity. In A10, we further show that improving the mask selection meaningfully improves model quality after fine-tuning.

Table 4: Wikitext2 perplexity of sparse LLaMA2-7B models pruned with SparseGPT.

Q10: Are these layer-wise sparsity reallocations necessary in the face of continued pre-training or fine-tuning?

A10: We understand the reviewer's concern with the necessity of improving the mask selection process when considering further fine-tuning. Thus, we conduct an ablation study on sparse LLaMA2-7B pruned with Wanda and Lua-LLM for LoRA fine-tuning with Alpaca training dataset. Each model is fine-tuned on 1 A100 GPU for 3 epochs, which takes about 13 hours. We report the training loss and evaluate the model performance after fine-tuning. The results (See Table below) show that:

(1) The loss converges in around 3 epochs, ensuring a meaningful comparison.

(2) Compared to the final converged training loss of Wanda (in around 13 hours), our Lua-LLM achieves comparable values in 0.5 epoch (in around 2 hours).

(3) Compared to Wanda, our Lua-LLM improves model quality after fine-tuning.

Table 5: Training loss convergence of sparse LLaMA2-7B during the LoRA fine-tuning process.

Table 6: Wikitext2 perplexity of fine-tuned LLaMA2-7B models for different pruning methods.

Thanks for the attentive reading of the manuscript and constructive feedback. We hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are glad to discuss further questions and suggestions.

[1] Sun, Mingjie, et al. "A simple and effective pruning approach for large language models." arXiv preprint arXiv:2306.11695 (2023).

[2] Rumi, Masuma Akter, et al. "Accelerating sparse cnn inference on gpus with performance-aware weight pruning." Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques. 2020.

[3] Yang, Tao, et al. "Spmmplu: A compiler plug-in with sparse ir for efficient sparse matrix multiplication." 2023 60th ACM/IEEE Design Automation Conference (DAC). IEEE, 2023.

We are glad that the reviewer found that our work effectively preserves the model capability and achieves superior search efficiency. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the questions you have raised.

Q1: The inconsistency that training with soft masks while inference with binary hard masks, is not adequately addressed.

Q4: How to address the train-test mismatch between soft masks and hard binary pruning?

A1: Our Lua-LLM employs the straight-through-estimator (STE) technique [1], enabling gradient computation via a soft approximation during back-propagation when processing non-differentiable functions. As discussed in Section 2, STE techniques have also been adopted in other learning-based pruning methods such as MaskLLM [2] and DISP-LLM [3]. In this paper, we further propose a novel approximation method based on the sigmoid function, which provides an effective approximation for the Top-K mask selection function. For the sigmoid-based soft mask in Equation (7), we set the value of parameter λ to the number of elements per-row, which is large enough to provide a precise approximation. During the training stage, we use the sigmoid-based soft mask in Equation (7) and search for sparsity allocation with learnable threshold parameters. After the training process, we store these threshold parameters and use the hard mask selection function in Equation (6) to generate hard masks with pruning thresholds for the inference stage. We will incorporate the detailed interpretation of the straight-through-estimator technique into Section 4.2 of our revised manuscript.

Q2: The paper emphasizes row-wise pruning as a key advantage but fails to compare it against other granularities.

Q5: Please add ablation studies on different pruning granularities.

A2: We conduct ablation experiments under different pruning granularities, i.e. different comparison group selection strategies. The results (See Table below) demonstrate that row-wise comparison remains the optimal choice for the adaptive sparsity allocation scenario, aligning with Wanda's observations for uniform pruning strategies, which is discussed in Section 3.2 of our manuscript.

Table 1: Perplexity of Lua-LLM with LLaMA2-7B on WikiText-2 under different pruning granularities.

Q3: Lua-LLM retains Wanda’s design without incorporating explicit output-side importance signals, failing to address this limitation in its own sparsity allocation strategy.

Q6: How to address the output-side weight importance in your method?

A3: We thank the reviewer for raising this meaningful discussion for output-side importance signals. We would like to clarify that Lua-LLM do not manually allocate sparsity based on heuristic statistics (e.g., output-side importance signals), in order to avoid suboptimal pruning results. Instead, we adopt a NAS-based pruning strategy, which automatically capture the output-side sensitivity to weight pruning, using gradient-based optimization techniques. Our method decomposes the mask selection process into two components: intra-row and inter-row processing. (a) During intra-row weight comparison, all weights are within the same output channel. Consequently, calculating their importance metrics does not require considering output-side weight importance, and we can directly use Wanda's metric in Equation (2) to capture the weight importance with input-side signals. (b) For inter-row sparsity allocation, we do not manually allocate sparsity based on heuristic statistics, but with a NAS-based pruning strategy. This approach allows us to search for the optimal solution and control the overall sparsity via a regularization loss, eliminating the need for manual determination of per-row sparsity based on importance scores.

Thanks for the attentive reading of the manuscript and constructive feedback. We hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are glad to discuss further questions and suggestions.

[1] Bengio, Yoshua, Nicholas Léonard, and Aaron Courville. "Estimating or propagating gradients through stochastic neurons for conditional computation." arXiv preprint arXiv:1308.3432 (2013).

[2] Fang, Gongfan, et al. "Maskllm: Learnable semi-structured sparsity for large language models." Advances in Neural Information Processing Systems 37 (2024): 7736-7758.

[3] Gao, Shangqian, et al. "Disp-llm: Dimension-independent structural pruning for large language models." Advances in Neural Information Processing Systems 37 (2024): 72219-72244.

We are glad that the reviewer found that our work is technically novel and clearly written. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the questions you have raised.

Q1: Detailed analysis of the complexity.

Q4: Could you compare the complexity of performing pruning on Lua-LLM vs other methods?

A1: Following the suggestions, we compare the pruning time costs of Wanda, OWL, DSA, and Lua-LLM for LLaMA-7B model (See Table below). We agree with the reviewer that NAS-based pruning methods (DSA, Lua-LLM) are computationally more expensive than proxy-based methods (Wanda, OWL), and we would like to consider this as a meaningful trade-off for the significant performance improvements. Besides, we would like to clarify that our Lua-LLM (1 hour) is significantly faster than evolutionary methods like DSA (12+ hours), since we leverage a soft Top-K selection function to enable an efficient row-wise mask learning process.

Table 1: Pruning time comparisons for different pruning methods.

Q2: An algorithm could help with clarity.

A2: We thank the reviewer for the meaningful suggestion and we will incorporate an algorithm into our revised manuscript. Furthermore, we summarize the algorithm of our Lua-LLM into three stages as follows:

(a) Initialization. We compute the importance scores for weight parameters with Equation (2), and uniformly initialize the row-wise threshold parameter t as the target sparsity p.

(b) Mask Learning. We set the training loss with Equation (10). In the training process, we generate row-wise pruning masks with Equation (7) for forward-propagation, and update the learnable threshold parameters with back-propagation. After the training process, we save the learned threshold parameters.

(c) Pruning. We use the threshold parameters and mask selection function in Equation (6) to prune the original model.

Q3: What are the speedups of other pruning methods (with or without the SpInfer framework)?

A3: We evaluate the inference speedups of Wanda and Lua-LLM for OPT-6.7B model at 50%-70% sparsity levels. The experiment (See Table below) shows that the speedups of different pruning methods are quite the same for each sparsity level, which is compatible with the parameter counts. The results demonstrate that although adaptive unstructured pruning methods introduce a more irregular pattern for the sparse weight matrices, we can achieve meaningful speedup for the overall model with multiple performance optimization techniques employed in SpMM kernels like SpInfer, such as efficient sparse format and fine-grained execution pipeline.

Table 2: Inference throughput (token/s) speedups for different pruning methods.

Thanks for the attentive reading of the manuscript and constructive feedback. We hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are glad to discuss further questions and suggestions.