Dynamic expert pruning reduces memory use for multi-task language models
Dynamic Expert Pruning for Multi-Agent Systems
Machine LearningArtificial IntelligenceMultiagent Systems
Summary
Large language models often use special parts called experts to handle different tasks, but keeping all these experts active at once takes up a lot of memory. The authors found that using the specific instructions or prompts an agent receives is enough to predict which experts are needed for each task. They created a method called Dynamic Expert Pruning (DEP) that picks the right experts dynamically for each request, saving memory and improving accuracy compared to methods that prune experts in a fixed way. DEP works well across many types of tasks and can handle tasks it hasn’t seen before without retraining.
What this means in practice
- •For ai platform engineers: Reduce memory requirements of multi-task language model services by dynamically selecting which experts to activate per request based on task descriptions.
- •For cloud infrastructure teams: Optimize deployment footprint when hosting multi-agent language models for heterogeneous tasks by using dynamic expert pruning to save accelerator memory resources.
Authors
Jabin Koo, Soheil Abbasloo, Sungjae Lee, Jungseul Ok
Abstract
Mixture-of-Experts (MoE) architectures scale language models efficiently by activating only a few experts per token, but the saving is confined to computation: every expert must stay resident on the accelerator, so memory bounds where these models can be deployed. Expert pruning reduces this footprint, yet existing methods are static --- a single mask, calibrated offline, is applied to the model for every subsequent request. This assumption can fail when the workload is heterogeneous, most prominently in multi-agent systems, where one backbone serves many tasks and roles at once: our analysis shows that different tasks and roles recruit different experts, while static methods assign one fixed subset to all of them. We therefore propose Dynamic Expert Pruning (DEP), which rests on a finding we establish here: an agent's system and task prompts are by themselves sufficient to identify the experts that agent and its task require, since that text already describes what the agent will do. A lightweight predictor, trained once on workflow transcripts, turns those prompts into a specialized per-request mask in a single forward pass, with no per-configuration calibration. Across diverse tasks and roles, model scales, and MoE architectures, DEP achieves better overall accuracy than static pruning and merging baselines, and generalizes to workflows unseen in training without retraining. Its margin over those baselines is largest when few experts are retained, suggesting that the role specialization inherent to multi-agent systems permits sparser serving than static pruning allows.