Enhancing One-shot Pruned Pre-trained Language Models through Sparse-Dense-Sparse Mechanism

TL;DR

This paper introduces SDS, a three-step Sparse-Dense-Sparse pruning framework that improves the performance of pruned pre-trained language models by optimizing weight distribution, outperforming existing methods at similar sparsity levels.

Contribution

The paper proposes a novel SDS pruning method that reconstructs a dense model before final pruning, enhancing performance of sparse PLMs without retraining.

Findings

SDS outperforms SparseGPT and Wanda at the same sparsity.

Reduces perplexity by 9.13 on Raw-Wikitext2.

Improves zero-shot accuracy by 2.05% on average.

Abstract

Pre-trained language models (PLMs) are engineered to be robust in contextual understanding and exhibit outstanding performance in various natural language processing tasks. However, their considerable size incurs significant computational and storage costs. Modern pruning strategies employ one-shot techniques to compress PLMs without the need for retraining on task-specific or otherwise general data; however, these approaches often lead to an indispensable reduction in performance. In this paper, we propose SDS, a Sparse-Dense-Sparse pruning framework to enhance the performance of the pruned PLMs from a weight distribution optimization perspective. We outline the pruning process in three steps. Initially, we prune less critical connections in the model using conventional one-shot pruning methods. Next, we reconstruct a dense model featuring a pruning-friendly weight distribution by reactivating pruned connections with sparse regularization. Finally, we perform a second pruning round, yielding a superior pruned model compared to the initial pruning. Experimental results demonstrate that SDS outperforms the state-of-the-art pruning techniques SparseGPT and Wanda under an identical sparsity configuration. For instance, SDS reduces perplexity by 9.13 on Raw-Wikitext2 and improves accuracy by an average of 2.05% across multiple zero-shot benchmarks for OPT-125M with 2:4 sparsity.