A Robust Optimization Guided Pruning Framework for Vision and Large Language Models

Pruning is a common approach to reduce the memory footprint and inference cost of large vision and language models. As these architectures continue to scale, one-shot pruning methods - i.e. approaches that prune the network without any retraining - have become increasingly attractive. Many popular one-shot pruning methods (e.g., WoodFisher, CAP, SparseGPT, and ALPS) typically optimize a quadratic objective under sparsity constraints. However, in practice, this objective is affected by multiple sources of uncertainty, including noise in the calibration data and variability introduced by algorithmic updates. To address these issues, we introduce RobOP, a robust optimization framework that explicitly accounts for such uncertainties. RobOP is modular and flexible, and can be applied with any existing pruning method through simple modifications motivated by our theoretical framework. We demonstrate that by taking into account uncertainty, RobOP offers improvements over prior pruning approaches. Our framework applies tractably across a range of stylized uncertainty sets, enabling robust one-shot pruning at scale. Our code is available at https://github.com/mazumder-lab/RobOP.

Large vision and language models can achieve impressive results, but their size makes them expensive to store and run. Pruning offers a way to reduce these costs by removing parts of a model while aiming to preserve its performance. A pruned model may be retrained to recover lost accuracy, but for today’s large models this extra training can be prohibitively expensive. This makes one-shot pruning, where the model is compressed without retraining, an especially practical alternative. Unfortunately, many modern one-shot pruning methods rely on imperfect estimates to decide what to remove: these estimates may be computed from limited data or become outdated as the model is progressively pruned. We introduce RobOP, a framework for making one-shot pruning more reliable. Instead of assuming that the information used to guide pruning is perfectly accurate, RobOP explicitly accounts for possible errors in it. This leads to simple changes that can be added to existing one-shot pruning methods with little extra computational cost. We show that RobOP improves pruning performance for both vision transformers and large language models, with especially strong gains when only limited calibration data is available. By making one-shot pruning more reliable, RobOP can help make powerful AI models easier and cheaper to deploy.