DRPruning: Efficient Large Language Model Pruning through Distributionally Robust Optimization

TL;DR

DRPruning is a novel method for large language model pruning that dynamically adjusts data distribution during training to maintain balanced performance across diverse domains, improving efficiency and robustness.

Contribution

It introduces a distributionally robust optimization approach for structured pruning, enabling automatic adjustment of reference losses and data ratios for better multi-domain performance.

Findings

Outperforms comparable models in perplexity and downstream tasks

Demonstrates robustness across various domains and distribution shifts

Automatically determines optimal data ratios and reference losses

Abstract

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.