ReAD: Reinforcement-Guided Capability Distillation for Large Language Models

2026-05-11

AI Summary

This work addresses the limitations of existing large language model distillation methods, which often neglect the interdependencies and trade-offs among diverse capabilities under fixed computational budgets, leading to marginal performance gains or even degradation of other useful abilities. To overcome this, the paper introduces ReAD, a novel framework that, for the first time, explicitly models systematic capability transfer and conflicts. ReAD integrates reinforcement learning with an uncertainty-aware contextual bandit mechanism to dynamically identify task-critical capabilities and adaptively allocate distillation resources. By optimizing token budget allocation based on expected utility, ReAD significantly enhances downstream task performance under identical computational constraints while mitigating harmful capability interference and inefficient distillation, consistently outperforming multiple strong baselines.

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.

Problem

• capability distillation
• large language models
• token budget
• capability interdependence
• knowledge distillation

Innovation

• capability distillation
• reinforcement learning
• contextual bandit
• token budget
• large language models