Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DRPruning: Efficient Large Language Model Pruning through Distributionally Robust Optimization

Hexuan Deng, Wenxiang Jiao, Xuebo Liu, Jing Li, Min Zhang, Zhaopeng Tu

TL;DR

DRPruning addresses the problem of uneven domain degradation during structured pruning of large language models by integrating distributionally robust optimization (DRO) with pruning and continued pretraining. It dynamically adjusts the data distribution and uses a learned reference loss to steer training toward underperforming domains, resulting in faster convergence, better preservation of original capabilities, and improved downstream performance in both monolingual and multilingual settings. Across pruning, continued pretraining, and instruction-tuning, DRPruning demonstrates lower perplexity, higher task performance, and stronger domain robustness, with automatic estimation of optimal reference losses and data ratios. The approach has broad potential beyond pruning, offering a principled way to balance training across heterogeneous data and tasks in large-scale NLP systems.

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.

DRPruning: Efficient Large Language Model Pruning through Distributionally Robust Optimization

TL;DR

DRPruning addresses the problem of uneven domain degradation during structured pruning of large language models by integrating distributionally robust optimization (DRO) with pruning and continued pretraining. It dynamically adjusts the data distribution and uses a learned reference loss to steer training toward underperforming domains, resulting in faster convergence, better preservation of original capabilities, and improved downstream performance in both monolingual and multilingual settings. Across pruning, continued pretraining, and instruction-tuning, DRPruning demonstrates lower perplexity, higher task performance, and stronger domain robustness, with automatic estimation of optimal reference losses and data ratios. The approach has broad potential beyond pruning, offering a principled way to balance training across heterogeneous data and tasks in large-scale NLP systems.

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.

Paper Structure

This paper contains 65 sections, 6 equations, 7 figures, 20 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Structured Pruning
  4. Distributionally Robust Optimization
  5. Our Proposed DRPruning Method
  6. Distributionally Robust Pruning
  7. Integrate DRO into pruning and continued pretraining.
  8. Further improvement.
  9. Dynamic Loss Function
  10. Minimum performance estimation.
  11. Reference loss adjustment.
  12. Dynamic Potential Distribution
  13. Reference data ratio adjustment.
  14. Experiments
  15. Experimental Setup
  16. ...and 50 more sections

Figures (7)

  • Figure 1: Data proportion update procedure for DRPruning. The gray part represents the standard training process, the yellow part represents the normal process for DRO, and the blue part represents our newly added module.
  • Figure 2: The curve of PPL changes during pruning from 7B. Over the first 640 iterations (the vertical dash line), the model size is gradually reduced from 7B to the target size, which causes an initial increase in PPL.
  • Figure 3: Win rate during instruction tuning. DRPruning outperforms Sheared Llama and ReSheared.
  • Figure 4: Effectiveness of our method compared to constant scheduling and naïve DRO during the 1.3B continued pretraining. Left figure: PPL trends; Right: average performance across 15 downstream tasks.
  • Figure 5: The generation procedure for our sentence continuation task. The orange nodes represent data storage nodes, while the blue trapezoidal nodes represent data processing nodes.
  • ...and 2 more figures