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

RobOP is the official implementation of a paper accepted by ICML 2026, proposing a robust optimization-based model pruning framework that significantly reduces computational overhead while maintaining model performance through uncertainty sets and robust optimization techniques. This framework addresses the core dilemma of traditional pruning methods—performance degradation and insufficient robustness when reducing computational load—and is applicable to both vision models and large language models (LLMs).

Large language models and vision models have massive parameters and high computational costs, hindering practical deployment. Model pruning is an effective compression technique, but traditional methods are based on heuristic rules or magnitude thresholds, lacking systematic consideration of uncertain factors during pruning. This leads to insufficient robustness and unstable performance of pruned models on out-of-distribution data or adversarial samples, making it difficult to meet reliability requirements in production environments.

RobOP (Robust Optimization Guided Pruning Framework) introduces robust optimization theory, with its core being the min-max paradigm: optimizing worst-case performance during pruning. It includes two variants: RobOP-ALPS (adapted for Adaptive Layer-wise Pruning Strategy) and RobOP-CAP (adapted for Channel Attention Pruning). Key mechanisms:

1. Uncertainty set modeling (Baseline, CTE, Trace, E sets), providing theoretical guarantees for performance lower bounds;
2. Alternating optimization strategy (outer layer searches for pruning masks, inner layer solves for worst-case adversarial perturbations);
3. Compatibility with existing pruning methods, plug-and-play.

RobOP performs excellently in multiple benchmark tests:

• Large Language Models: On Llama3.1-8B, RobOP-ALPS maintains over 90% of the original performance, reduces parameters by more than 40%, and has less performance degradation on out-of-distribution tests than traditional methods;
• Vision Models: On DeiT-Small, RobOP-CAP has an ImageNet Top-1 accuracy loss ≤2% and a 1.8x improvement in inference speed;
• Uncertainty Set Comparison: The Trace set is optimal for LLMs, while the CTE set is better for vision tasks.

RobOP provides a solution that balances efficiency and reliability for deploying AI models on resource-constrained devices, especially suitable for fields with high robustness requirements such as autonomous driving and medical diagnosis. Its open-source implementation is based on PyTorch, compatible with Hugging Face Transformers, has a concise command-line interface, supports flexible configuration of uncertainty sets and pruning strategies, and is easy to integrate quickly.

Limitations:

1. The additional computational overhead of robust optimization is significant for large-scale models;
2. The selection of uncertainty sets requires domain knowledge, and automated mechanisms need to be explored. Future Directions:

• Develop more efficient uncertainty quantification methods;
• Explore joint optimization with compression techniques such as quantization and distillation;
• Extend to multimodal models.

RobOP is an important step in the evolution of model pruning toward robustness orientation. By introducing robust optimization theory to improve the reliability of compressed models, it provides a new theoretical perspective for subsequent research. For engineers and researchers working on model deployment optimization, RobOP is a powerful tool worth exploring in depth.