ReAD: Reinforcement-Guided Capability Distillation for Large Language Models

TL;DR

ReAD introduces a reinforcement-guided framework for capability distillation in large language models, optimizing the transfer of interdependent abilities within a fixed token budget to enhance downstream performance.

Contribution

It proposes a novel capability distillation method that explicitly models interdependence and adaptively allocates resources, outperforming existing approaches.

Findings

ReAD improves downstream utility under the same token budget.

It reduces harmful capability spillover during distillation.

ReAD outperforms strong baselines in experiments.

Abstract

Capability distillation applies knowledge distillation to selected model capabilities, aiming to compress a large language model (LLM) into a smaller one while preserving the abilities needed for a downstream task. However, most existing methods treat capabilities as independent training targets and overlook how improving one capability can reshape the student's broader capability profile, especially when multiple abilities jointly determine task success. We study capability distillation under a fixed token budget and identify two consistent patterns: distillation induces systematic, budget-dependent cross-capability transfer, and additional budget often brings limited task-relevant gains while sometimes degrading other useful abilities. Building on these insights, we propose ReAD, a Reinforcement-guided cApability Distillation framework that explicitly accounts for capability interdependence. ReAD first infers task-essential capabilities, then generates capability-targeted supervision on the fly, and finally uses an uncertainty-aware contextual bandit to adaptively allocate the distillation budget based on expected utility gains. Extensive experiments show that ReAD improves downstream utility under the same token budget while reducing harmful spillover and wasted distillation effort compared to strong baselines. Our code is publicly available at https://github.com/LabRAI/ReAD.