ALPS: Improved Optimization for Highly Sparse One-Shot Pruning for Large Language Models

TL;DR

ALPS is an optimization-based pruning framework for large language models that significantly improves sparsity and performance, especially at high sparsity levels, by leveraging advanced optimization techniques and GPU acceleration.

Contribution

ALPS introduces a novel optimization-based approach for one-shot pruning of large language models, outperforming heuristic methods in achieving higher sparsity and better model performance.

Findings

Achieves 70% sparsity with 13% perplexity reduction on WikiText

Outperforms state-of-the-art methods in zero-shot benchmarks

Provides theoretical convergence guarantees for pruning

Abstract

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.