DRPruning: Efficient Large Language Model Pruning through Distributionally Robust Optimization
This addresses the issue of biased performance in pruned LLMs for applications requiring robust multi-domain or multi-tasking capabilities, representing an incremental improvement over existing pruning methods.
The paper tackles the problem of uneven performance degradation across domains in structured pruning of large language models by proposing DRPruning, which dynamically adjusts data distribution during training to restore balanced performance, achieving superior results in perplexity, downstream tasks, and instruction tuning compared to similarly sized models.
Large language models (LLMs) deliver impressive results but face challenges from increasing model sizes and computational costs. Structured pruning reduces model size and speeds up inference but often causes uneven degradation across domains, leading to biased performance. To address this, we propose DRPruning, a method that dynamically adjusts the data distribution during training to restore balanced performance across heterogeneous and multi-tasking data. Experiments in monolingual and multilingual settings show that DRPruning surpasses similarly sized models in both pruning and continued pretraining over perplexity, downstream tasks, and instruction tuning. Further analysis demonstrates the robustness of DRPruning towards various domains and distribution shifts. Furthermore, DRPruning can determine optimal reference losses and data ratios automatically, suggesting potential for broader applications. Code and scripts are available at https://github.com/hexuandeng/DRPruning.