The Optimal BERT Surgeon: Scalable and Accurate Second-Order Pruning for Large Language Models (2203.07259v3)

Abstract: Transformer-based LLMs have become a key building block for natural language processing. While these models are extremely accurate, they can be too large and computationally intensive to run on standard deployments. A variety of compression methods, including distillation, quantization, structured and unstructured pruning are known to decrease model size and increase inference speed, with low accuracy loss. In this context, this paper's contributions are two-fold. We perform an in-depth study of the accuracy-compression trade-off for unstructured weight pruning of BERT models. We introduce Optimal BERT Surgeon (oBERT), an efficient and accurate weight pruning method based on approximate second-order information, which we show to yield state-of-the-art results in both stages of language tasks: pre-training and fine-tuning. Specifically, oBERT extends existing work on unstructured second-order pruning by allowing for pruning blocks of weights, and by being applicable at the BERT scale. Second, we investigate the impact of this pruning method when compounding compression approaches to obtain highly compressed but accurate models for deployment on edge devices. These models significantly push boundaries of the current state-of-the-art sparse BERT models with respect to all metrics: model size, inference speed and task accuracy. For example, relative to the dense BERT-base, we obtain 10x model size compression (in MB) with < 1% accuracy drop, 10x CPU-inference speedup with < 2% accuracy drop, and 29x CPU-inference speedup with < 7.5% accuracy drop. Our code, fully integrated with Transformers and SparseML, is available at https://github.com/neuralmagic/sparseml/tree/main/research/optimal_BERT_surgeon_oBERT.

The Optimal BERT Surgeon: Enhancing LLM Efficiency

The paper "The Optimal BERT Surgeon: Scalable and Accurate Second-Order Pruning for LLMs" addresses key challenges in the domain of LLMs, particularly focusing on BERT. The core contribution is the introduction of the Optimal BERT Surgeon (oBERT) method, which enhances model efficiency through scalable and precise second-order pruning techniques. This essay provides an overview of the paper's contributions, methodology, results, and implications for the future of AI model deployment.

Contributions and Methodology

The authors present oBERT as a method built on approximate second-order information, extending the Optimal Brain Surgeon (OBS) framework. The approach offers two primary advantages:

1. Scalable Application: It extends existing second-order pruning techniques to accommodate the scale of BERT models, allowing for both unstructured and block pruning.
2. Accurate Pruning: By employing second-order (curvature) pruning, oBERT retains high accuracy during both pre-training and fine-tuning phases, crucial for upstream and downstream tasks.

The paper also explores a compound compression strategy, integrating pruning with quantization. Such integration amplifies the method's applicability for deployment on edge devices, achieving significant model compression and inference speed improvements without substantial accuracy loss.

Key Results

The oBERT method sets state-of-the-art performance benchmarks, showcasing substantial advancements over prior approaches:

• Model Compression: The method achieves up to 10x compression in model size relative to dense BERT, with less than a 1% drop in accuracy.
• Inference Speed: The deployment on CPUs is notably enhanced, reaching up to 29x speedups with an acceptable accuracy trade-off (<7.5% drop).

Empirical evaluations across various benchmarks, such as SQuAD v1.1, MNLI, and QQP, substantiate these claims. The oBERT method consistently surpasses previous techniques like Movement Pruning (MvP) and Lottery Ticket-based approaches, demonstrating superior accuracy-sparsity trade-offs.

Implications

The implications of such advancements are multifold:

• Practical Deployment: By achieving high sparsity with minimal performance degradation, oBERT facilitates the deployment of LLMs on resource-constrained environments, opening up new application domains.
• Theoretical Insights: The results provide significant insights into the balance of model complexity, capacity, and efficiency, pushing the envelope in model compression strategies.

Future Directions

While the paper establishes a robust framework for second-order pruning, future research could focus on expanding this work to other types of LLMs and exploring the integration with advanced distillation techniques. Additionally, further investigation into adaptive methodologies that dynamically adjust pruning and quantization parameters during training is a promising avenue.

In conclusion, the paper makes a substantial contribution to the field of efficient LLMs, delivering a pragmatic and theoretically sound approach to model compression. The Optimal BERT Surgeon represents a critical step towards making powerful LLMs more accessible and sustainable.