Source-linked AI summary
ALPS: Improved Optimization for Highly Sparse One-Shot Pruning for Large Language Models
Xiang Meng, Kayhan Behdin, Haoyue Wang, Rahul Mazumder
TL;DR
LLM scale makes pruning costly, and existing approaches often rely on heuristics that may produce suboptimal compression tradeoffs. ALPS directly optimizes layerwise pruning with ADMM and PCG, and reports improved pruning and model-performance results over existing methods.
Problem
LLM scale makes pruning computationally challenging, and existing LLM pruning often uses heuristics that may yield suboptimal compression tradeoffs.
Source: §1 · p. 1Method
ALPS directly solves layerwise ℓ0-constrained one-shot pruning with ADMM, then fixes the support and refines weights using modified PCG.
Source: §1 · p. 2, §3.1 · p. 4Results
Takeaways & Limitations
ALPS provides an optimization-based one-shot unstructured pruning approach with convergence guarantees and reported gains over existing pruning methods.
Source: §5 · p. 9The paper focuses on one-shot unstructured pruning; extending ALPS to structured pruning constraints and quantization is left for future work.
Source: §5 · pp. 9–10
Abstract
from arXiv · showhide
The impressive performance of Large Language Models (LLMs) across various natural language processing tasks comes at the cost of vast computational resources and storage requirements. One-shot pruning techniques offer a way to alleviate these burdens by removing redundant weights without the need for retraining. Yet, the massive scale of LLMs often forces current pruning approaches to rely on heuristics instead of optimization-based techniques, potentially resulting in suboptimal compression. In this paper, we introduce ALPS, an optimization-based framework that tackles the pruning problem using the operator splitting technique and a preconditioned conjugate gradient-based post-processing step. Our approach incorporates novel techniques to accelerate and theoretically guarantee convergence while leveraging vectorization and GPU parallelism for efficiency. ALPS substantially outperforms state-of-the-art methods in terms of the pruning objective and perplexity reduction, particularly for highly sparse models. On the OPT-30B model with 70% sparsity, ALPS achieves a 13% reduction in test perplexity on the WikiText dataset and a 19% improvement in zero-shot benchmark performance compared to existing methods.
1 Introduction
LLM scale makes compression important, but scale-driven reliance on heuristics in one-shot pruning may yield suboptimal compression; ALPS instead directly optimizes layerwise pruning.
1 Introduction
- LLMs' storage and compute demands are severe, while one-shot pruning seeks to preserve accuracy without expensive full-dataset retraining.One-shot methods compress a pretrained model using a small amount of data, such as a few thousand samples.Source: §1 · pp. 1–2
- Because LLM-scale methods often use heuristics, they may deliver suboptimal compression-accuracy tradeoffs versus optimization-based approaches.The introduction names SparseGPT, Wanda, and an iterative mask-growing approach as examples of heuristics.Source: §1 · p. 2
- After support stabilizes, modified PCG optimizes weights on that support; vectorization and GPU parallelism yield a 20x-200x speedup over vanilla backsolve.The PCG post-processing solves the restricted quadratic problem.Source: §1 · pp. 2–3
- Figure 1 depicts the pipeline from layerwise ℓ0-constrained pruning, through ADMM support selection, to PCG weight refinement.The figure assigns these stages to its left, middle, and right panels, respectively.Source: Figure 1 · p. 2
- ALPS formulates layerwise pruning as an ℓ0-constrained optimization problem and uses ADMM to find support while updating weights.The method extends operator splitting to the non-convex, non-continuous problem and provides convergence guarantees.Source: §1 · pp. 2–3
- At 70% sparsity on LLaMA3-8B, ALPS reduces WikiText test perplexity by 29% and improves zero-shot benchmark performance by 4%-13%.The paper also reports a 20x-200x speedup over the vanilla backsolve approach.Source: §1 · p. 3
2 Related Work
Related work distinguishes pruning by sparse-network structure and training requirements, reviews post-training unstructured methods, and situates ALPS among LLM-scale heuristics and prior ADMM uses.
Network pruning
- Pruning methods vary by structure and training requirements; this work focuses on one-shot unstructured pruning.Unstructured pruning removes individual weights, while structured pruning removes structures such as channels, filters, or attention heads.Source: §2 · p. 3
Post-training unstructured pruning
- Post-training unstructured methods include importance scoring, second-order loss approximations, and layerwise reconstruction objectives.The paper adopts layerwise reconstruction error as its pruning objective.Source: §2 · p. 3
Unstructured pruning in LLMs
- LLM pruning methods address scale with approaches including SparseGPT's partial updates, Wanda's weight-activation criterion, and DSnoT's iterative mask growth and pruning.Boža [2024] determines optimal weights on a given support and extends the approach to support-update heuristics.Source: §2 · pp. 3–4
ADMM in network pruning
- Prior pruning work has used ADMM, while ALPS distinguishes itself by jointly optimizing weights and sparsity patterns at LLM scale and introducing a penalty update scheme.The scheme is described as ensuring convergence both practically and theoretically.Source: §2 · p. 4
3 ALPS: Effective LLM pruning in One-shot
ALPS formulates layer-wise pruning as an ℓ0-constrained reconstruction problem and solves it with operator splitting, followed by weight refinement on the stabilized support. Its penalty update balances support search with convergence, and the paper provides convergence guarantees for the non-convex, non-continuous problem.
3.1 Problem formulation
- The layer-wise objective minimizes output reconstruction error under a target sparsity constraint, with a ridge term keeping pruned weights near the original weights.The formulation uses a sparse weight matrix and calibration-sample input activations.Source: §3.1 · p. 4
3.2 Operator-splitting for layer-wise pruning
- ADMM alternates weight and sparsity-projection updates, followed by a dual update, to identify the support while updating weights.The reformulation introduces a copy of the weight matrix and couples the two variables through an equality constraint.Source: §3.2 · pp. 4–5
- The penalty parameter starts small and increases according to support changes, balancing search for a good support with faster convergence as ρ grows.The paper reports that small ρ can slow convergence, while large ρ may compromise solution quality despite early support stabilization.Source: §3.2 · p. 5
- The algorithm is guaranteed to converge on a non-convex, non-continuous ℓ0-constrained problem; the iterates W and D approach a shared limit.The paper presents this as a novel convergence result for such problems, to the authors’ knowledge.Source: §3.2 · p. 5
3.3 Efficiently refining weights after support stabilization
- After the support stabilizes, PCG fixes that support and refines the weights; this avoids solving separate linear systems for each output column.The paper notes that direct backsolve would require Nout systems of size O(Nin), costly at LLM dimensions.Source: §3.3 · p. 6
- The PCG refinement solves the weight system in one pass with support projection at each iteration, using vectorization and GPU matrix multiplications for efficiency.The method is designed to avoid separate per-column solves.Source: §3.3 · pp. 6–7
4 Experimental Results
Across single-layer and model-level evaluations, ALPS reduces reconstruction error and outperforms competing methods, with advantages reported at higher sparsity and for N:M patterns.
4 Experimental Results
- Evaluations cover OPT models from 1.3B to 30B, LLaMA2 models at 7B and 13B, and LLaMA3-8B, using C4 calibration data and perplexity and zero-shot benchmarks.The zero-shot tasks are MMLU, PIQA, LAMBADA, ARC-Easy, and ARC-Challenge.Source: §4 · p. 7
4.1 Reconstruction error on a single layer
- 20%–40%: ALPS-selected supports yield lower reconstruction error than other methods when each support is refined to optimality.This comparison isolates support quality by solving the post-processing problem to optimality for each method.Source: §4.1 · p. 8, Table 1 · p. 8
- PCG post-processing achieves reconstruction errors comparable to the optimal solution while running 20x–200x faster.The comparison uses magnitude pruning to determine the support.Source: §4.1 · p. 8, Table 1 · p. 8
4.2 Pruning OPT and LLaMA models
- 60%: ALPS achieves a perplexity reduction on WikiText2 versus other methods at 80% sparsity; it outperforms competitors above 50% sparsity.The performance gap widens as sparsity increases.Source: §4.2 · p. 8, Figure 3 · p. 9
4.3 N:M sparsity
- ALPS outperforms other methods on most datasets for N:M pruning, with larger improvements than unstructured pruning at the same sparsity.The results cover OPT-30B and LLaMA2-13B at 2:4 and 4:8 patterns.Source: §4.3 · p. 8
5 Conclusion
ALPS is an optimization-based framework for one-shot unstructured LLM pruning, combining operator splitting with PCG post-processing. Experiments show improvements over existing pruning methods, while structured pruning constraints and quantization remain future extensions.
5 Conclusion
- ALPS combines operator splitting for the ℓ0-constrained layer-wise pruning problem with PCG post-processing, a penalty-parameter update scheme, and convergence guarantees.The post-processing uses vectorization and GPU parallelism.Source: §5 · p. 9
- Experiments confirm that ALPS outperforms existing pruning methods on both the pruning objective and the performance of pruned models.Source: §5 · p. 9
- Structured pruning constraints and quantization are identified as future extensions for understanding the approach's strengths and scope.Source: §5 · p. 10
A Proofs of Theorem 1
The proof establishes convergence by combining lemmas with Algorithm 1's initialization and showing the generated sequence is Cauchy.
A Proofs of Theorem 1
- The proof uses H = X⊤X + λ2I and G in its theorem setup.The passage truncates the definition of G.Source: A Proofs of Theorem 1 · p. 16
- The main theorem proof combines Lemma 2 with Algorithm 1's initialization.The sequence under discussion is generated by Algorithm 1.Source: A Proofs of Theorem 1 · p. 16
- The proof concludes that D(t) is Cauchy and that W(t) converges to the same limit.It first obtains convergence of D(t), then uses the vanishing difference between W(t+1) and D(t+1).Source: A Proofs of Theorem 1 · p. 16
A.1 Proof of Lemma 1
The appendix establishes Lemma 1's inequalities and uses them in Lemma 2, where a telescoping argument completes the bound.
A.1 Proof of Lemma 1
- The D-update rule and projection operator establish inequality (7).Summing inequalities (16) and (18) verifies inequality (8).Source: §A.1 · p. 17
A.2 Proof of Lemma 2
- Lemma 2 begins by applying Lemma 1 and then combines inequalities (20) and (22)..Source: §A.2 · pp. 17–18
- The Lemma 2 proof defines a_t from the norms of D(t) and V(t), then uses telescoping to complete the proof.a_t = ∥D(t)∥F + ∥V(t)∥F /ρt.Source: §A.2 · p. 18
B.1 Experimental setup
Experiments use sequential, layer-wise pruning with calibration activations, specified baselines, and a preprocessing transformation for scaling.
B.1 Experimental setup
- Experiments ran on an Intel Xeon Gold 6248 with 20 CPUs and one NVIDIA V100 GPU.The machine had 192GB CPU RAM and 32GB CUDA memory; PyTorch was used for implementation.Source: §B.1 · p. 18
- For each sparsity level, the ℓ0 constraint k is set to ⌊NinNouts⌋, and layers are pruned sequentially.Each layer's input activations are generated by the preceding pruned layers on N calibration samples.Source: §B.1 · p. 18
- Preprocessing defines E = Diag(X⊤X + λ2I)−1/2 and transforms W to W′ = E−1W for scaling.The original solution is recovered by setting W = EW′; the transformation does not alter Algorithm 1 or its convergence analysis and can improve convergence rate in practice.Source: §B.1 · pp. 18–19
- The penalty starts at ρ0 = 0.1, and support stabilization triggers Algorithm 2 for 10 iterations on the fixed support.The reported initial regularization choice is λ2 = 0.01 Tr(X⊤X).Source: §B.1 · p. 19
- Comparisons include magnitude pruning, SparseGPT, Wanda, and DSnoT, with the latter three run using their authors’ default implementations.Magnitude pruning keeps the top k absolute dense-weight entries unchanged and zeros the rest.Source: §B.1 · p. 19
B.2.1 The importance of the ρ update scheme
ALPS's penalty update is reported to combine rapid convergence with high-quality solutions, unlike fixed-ρ ADMM settings. Its optimization takes longer than other algorithms, but its runtime is described as negligible relative to fine-tuning.
B.2.1 The importance of the ρ update scheme
- ALPS's ρ update achieves rapid convergence and high-quality solutions, while fixed ρ=3 is fast but poor and fixed ρ=0.3 fails to converge.Table 4 reports relative reconstruction error, and Table 5 reports support-change rate between iterations.Source: §B.2.1 · p. 19, Table 4 · p. 19, Table 5 · p. 19
- The support-change rate between consecutive iterations is used as a measure of algorithm convergence speed.The comparison is between ALPS and ADMM with fixed penalty parameters.Source: §B.2.1 · p. 19, Table 5 · p. 19
B.2.2 Runtime comparison
- ALPS's advanced optimization runs longer than other pruning algorithms, but its runtime is negligible compared with LoRA fine-tuning.Table 6 compares runtime for one-shot unstructured pruning of OPT and LLaMA models, including activation generation and pruning.Source: §B.2.2 · p. 20, Table 6 · p. 20
B.2.3 Comparison of ALPS and ADMM-Grad
ALPS jointly optimizes weights and sparsity patterns, achieving lower layerwise reconstruction loss than ADMM-Grad. With a fixed support, its PCG procedure beats ADMM-Grad in runtime and objective value.
B.2.3 Comparison of ALPS and ADMM-Grad
- ALPS achieves lower layerwise reconstruction loss than ADMM-Grad by jointly optimizing weights and sparsity patterns.ADMM-Grad selects its sparsity mask through iterative magnitude pruning before applying ADMM.Source: §B.2.3 · p. 20, Table 7 · p. 20
- Unlike ADMM-Grad, ALPS directly targets the ℓ0-constrained least-squares problem, optimizing weights and sparsity patterns together.ADMM-Grad first chooses a mask through iterative magnitude pruning.Source: §B.2.3 · p. 20
- PCG outperforms ADMM-Grad in computational time and objective value when solving problem (6) with a given support.Its time advantage comes from backsolving without explicitly computing matrix inverses.Source: §B.2.3 · p. 20, Table 8 · p. 20
B.2.4 Pruned model performance on MMLU benchmark
On LLaMA3-8B MMLU, ALPS outperforms the compared pruning methods, while performance degrades significantly at high sparsity. Broader evaluations cover OPT and LLaMA models across sparsity settings, language-modeling datasets, and zero-shot benchmarks.
B.2.4 Pruned model performance on MMLU benchmark
- ALPS outperforms MP, SparseGPT, Wanda, and DSnoT on mean MMLU accuracy for LLaMA3-8B.Table 9 presents performance at various sparsity levels; no exact accuracy values are given in the supplied passages.Source: §B.2.4 · p. 20, Table 9 · p. 21
- At high sparsity, MMLU performance degrades significantly, and the authors suggest fine-tuning or prompting to maintain practical performance.The evaluation compares MP, SparseGPT, Wanda, DSnoT, and ALPS on LLaMA3-8B.Source: §B.2.4 · pp. 20–21
B.2.5 Comprehensive model performance across sparsity levels
- Evaluations span OPT-1.3B–30B and LLaMA2/3 models, 40–90% and 2:4/4:8 sparsity, WikiText2, PTB, C4, PIQA, and ARC; LLaMA LAMBADA results are omitted.Each model, sparsity, and method combination is run five times, with mean and standard deviation reported.Source: §B.2.5 · p. 21, Table 10 · p. 22, Table 16 · p. 25