Title: Computation and Memory-Efficient Model Compression with  Gradient Reweighting
Venue: NeurIPS 2025
Accepted: yes
Keywords: model compression

deep neural networks

sparse training

deep learning
Abstract: Pruning is a commonly employed technique for deep neural networks (DNNs) aiming at compressing the model size to reduce computational and memory costs during inference. In contrast to conventional neural networks, large language models (LLMs) pose a unique challenge regarding pruning efficiency due to their substantial computational and memory demands. Existing methods, particularly optimization-based ones, often require considerable computational resources in gradient estimation because they cannot effectively leverage weight sparsity of the intermediate pruned network to lower compuation and memory costs in each iteration. The fundamental challenge  lies in the need to frequently instantiate intermediate pruned sub-models to achieve these savings, a task that becomes infeasible even for moderately sized neural networks.  To this end, this paper proposes a novel pruning method for DNNs that is both computationally and memory-efficient. Our key idea is to develop an effective reweighting mechanism that enables us to estimate the gradient of the pruned network in current iteration via reweigting the gradient estimated on an outdated intermediate sub-model instantiated at an earlier stage, thereby significantly reducing model instantiation frequency. We further develop a series of techniques, e.g., clipping and preconditioning matrix, to reduce the variance of gradient estimation and stabilize the optimization process. We conducted extensive experimental validation across various domains. Our approach achieves 50\% sparsity and a 1.58$\times$ speedup in forward pass on Llama2-7B model with only 6 GB of memory usage, outperforming state-of-the-art methods with respect to both perplexity and zero-shot performance. As a by-product, our method is highly suited for  distributed sparse training  and can achieve a 2 $\times$ speedup over the dense distributed baselines.


Peer reviews:
comments: SUMMARY:
This paper proposes an optimization-based pruning method for large language models (LLMs). The key contribution is a reweighting technique that allows the reuse of outdated sub-models to estimate gradients for structural parameters, reducing the need to frequently instantiate new sub-networks. The method also incorporates variance reduction techniques to stabilize training. Experimental results demonstrate memory savings, inference speedup, and competitive performance compared to state-of-the-art pruning methods. The paper also extends the method to a distributed sparse training setting.

STRENGTHS AND WEAKNESSES:
**Strengths**

1. The method allows using an outdated sub-model for multiple pruning steps, which significantly reduces computational cost.

2. The method achieves competitive or superior performance compared to several baselines in terms of perplexity, memory usage, and inference speed.

3. The paper is well organized and easy to follow.

**Weaknesses**

1. The introduction states that "Consequently, LLMs pose a unique challenge for pruning", but this claim lacks sufficient justification. A clearer explanation should be provided.

2. The theoretical parts in Section 4 are sometimes dense and hard to follow without a strong background in knowledge.

3. Comparisons with more recent optimization-based methods are limited. Most baselines are metric-based.

4. typos: caption of Figure 1 ("vanila" -> "vanilla")

QUESTIONS:
Please refer to Strengths and Weaknesses

LIMITATIONS:
The authors have not discussed the limitations and potential negative societal impact of their work. It is recommanded to be included in the final revision.
TITLE: Official Review
CLARITY: 2
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a computationally and memory-efficient pruning algorithm for LLMs, which addresses a key practical challenge in optimization-based structured pruning — the high overhead of sub-model instantiation. The core idea is a reweighting scheme, which reuses the old sub-model gradients to estimate the current gradient. The paper integrates multiple variance reduction techniques and demonstrates the method’s advantages in both standalone and distributed pruning scenarios.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. Novel reweighting-based pruning technique with strong theoretical and practical grounding.\n2. Extensive experimental coverage across LLaMA, OPT, and vision models.\n\nWeaknesses\n1. I'm wondering about the results beyond 50% sparsity and robustness under extreme pruning, which is lacking in the paper.\n2. SliceGPT reported that the WikiText2 of Llama2-7B at 30% sparsity is 8.12, which significantly differs from the one reported in this paper (23.31). The WikiText2 of Llama2-13B at 30% sparsity is also larger than the one in the original paper. The difference also occurs in the results of LLM-pruner, I'm wondering why there is such a gap. Is there some difference in the experimental setup? If so why not use the same setup as the original paper?\n3. Comparison with optimization-based baselines is limited. Only 50% sparsity is compared, more sparsity levels should be compared. Nuteprune reported that the WikiText2 of Llama2-7B at 50% sparsity is 12.29, which significantly differs from the one reported in this paper. Does this difference come from the lack of fine-tuning?\n4. Since the other methods do not appear to match the results presented in their own original papers. The quality of these results in a scientific context is low, because it is difficult to draw any conclusions from them. I hope the author can provide some explanation for the difference. \n5. I want to know why the authors only compare the results without fine-tuning? If the models pruned by gradient reweighting can't beat those SOTA methods after fine-tuning, the practical significance of gradient reweighting would seem to be less. I recommend that the authors compare the performance after fine-tuning.\n6. I wish the authors could provide statistical information rather than single experiments, which would help to increase the credibility of the results\n\nQUESTIONS:\nSee weaknesses.\n\nIf the authors can solve some of my confusion in experiments, the practical significance of the paper will be greatly improved, and I will be happy to increase my score.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenges of pruning large language models (LLMs), where conventional methods struggle due to the high computational and memory overhead of repeated sub-model instantiations. The authors propose a novel, resource-efficient pruning approach that introduces a reweighting mechanism to estimate gradients of the current pruned model using gradients from previously instantiated sub-models. This significantly reduces the frequency of model instantiation. Additional techniques such as gradient clipping and preconditioning are introduced to stabilize optimization and reduce variance. The method is empirically validated across domains, achieving 50\\% sparsity and 1.58$\\times$ forward-pass speedup on LLaMA2-7B using only 6 GB of memory, with improved perplexity and zero-shot performance compared to prior work. Furthermore, the technique demonstrates strong potential for distributed sparse training, achieving 2$\\times$ speedup over dense baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper presents a clear and concise mathematical formulation, including only the essential equations—particularly for the variance reduction mechanism—without unnecessary complexity.\n2. The ablation studies are thorough and well-structured, effectively isolating the impact of each proposed component on addressing key bottlenecks in pruning.\n3. The experimental evaluation is comprehensive, covering both algorithmic metrics (e.g., perplexity, zero-shot performance) and practical system-level outcomes (e.g., memory usage and actual speedup).\n\nWeaknesses:\n1. The manuscript repeatedly emphasizes the three main contributions throughout multiple sections, leading to redundancy and detracting from the overall clarity and flow of the writing.\n2. The empirical comparison with prior optimization-based pruning methods is limited, relying on only two baseline works. A broader and more diverse set of baselines would strengthen the validity of the claims.\n3. The trade-off introduced by the variance reduction techniques appears to slightly degrade model performance in Table 5. It would be helpful to evaluate whether variance reduction is strictly necessary—can the method maintain acceptable accuracy while further improving throughput and memory consumption if the variance reduction is omitted?\n\nQUESTIONS:\nIs the m in Equation 1 either 0 or 1? Can we prune part of the each sub-module?\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 1, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces a reweighting-based mask updating method for pruning DNNs, particularly Large Language Models (LLMs), with a focus on computational and memory efficiency. The primary challenge addressed is the computational complexity and memory redundancy associated with instantiating intermediate pruned sub-models. Instead of relying on back-propagation for the entire neural network, this method estimates the gradient of the pruned network. The authors highlight that the DNN feed-forward-only mechanism can also be generalized to sparse distributed training. The results demonstrate state-of-the-art performance in LLM and DNN compression compared to existing pruning algorithms.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths \n* The reweighted policy gradient estimator approximates the gradient of the mask distribution parameter using outdated sub-model instantiation results, significantly reducing model instantiation frequency.\n* Variance reduction in the pruning framework is mathematically proven to be bounded, demonstrating the framework's stability for practical application.\n* The framework runs in parallel with model training, making it compatible with existing methods for simultaneous model finetuning and compression.\n* As a byproduct, the framework is also compatible with sparse distributed training.\n* The reweighting-based pruning framework achieves state-of-the-art performance compared to existing pruning frameworks for DNNs and LLMs.\n* The presentation is logical, smooth, and clear, making it easy for future researchers to extend the work. Detailed discussions on theorems, implementation, and code are provided in the appendix.\n\n\n\n\nWeakness\n* The study highlights that pruning performance on Large Language Models (LLMs) remains largely unexplored. While the proposed method outperforms other LLM pruning techniques, the pruned network significantly underperforms the original LLM. This disparity might stem from LLMs' inherent sensitivity to noise. However, computational resource limitations prevented the application of larger, more redundant networks to validate this hypothesis. Consequently, further research is needed to determine if the trade-off between LLM complexity and efficiency justifies this approach to LLM pruning.\n* minor: line 159 compactess -> compactness\n\nQUESTIONS:\nSuper LLMs and Pruning Limitations\n* Resource Requirements: What are the potential resource requirements for super LLMs to demonstrate the limitations of pruning?\n* Network Growth and Redundancy: Theoretically, at what point does network growth indicate sufficient redundancy to warrant pruning?\n\nDistributed Sparse Training and Asynchronous Pruning\n* Full-Sync Optimization: Does the current distributed sparse training framework necessitate full-synchronization for optimization?\n* Asynchronous Training and Framework Adjustments: If asynchronous training is implemented, what adjustments are required? Would a separate framework be necessary to accommodate asynchronous pruning?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}