Technical Summary: ReAD - Reinforcement-Guided Capability Distillation for Large Language Models

1. Problem Statement

The paper addresses the challenge of capability distillation for Large Language Models (LLMs), a specialized form of knowledge distillation (KD) where a small student model is trained to imitate a large teacher model's performance on specific capabilities (e.g., reasoning, coding, math) to reduce serving costs.

Existing methods typically treat capabilities as independent training targets, assuming that optimizing one capability does not negatively impact others. However, the authors identify a critical gap: under a fixed token budget, improving a single capability often triggers systematic, budget-dependent cross-capability transfer, where gains in the target capability come at the expense of other useful abilities. This leads to two specific inefficiencies:

1. Diminishing Returns: Additional tokens allocated to a single target yield progressively smaller marginal gains.
2. Harmful Spillover: Updates intended for one capability degrade performance on non-target capabilities required for downstream task success.

The core problem is formulated as Problem 2.3: How to allocate a fixed distillation budget $B$ across a set of interdependent capabilities $C$ to maximize downstream task utility $U_\tau$, while minimizing negative side effects (spillover) on other capabilities.

2. Methodology: ReAD Framework

The authors propose ReAD (Reinforcement-guided cApability Distillation), a framework that explicitly models capability interdependence and adaptively allocates the token budget. ReAD operates through three main stages:

2.1 Identifying Task-Essential Capabilities

ReAD first constructs a task requirement vector $r_\tau \in \Delta^{|C|}$ for a given downstream task $\tau$.

• Mechanism: A lightweight identifier $g_\phi$ (a Transformer encoder followed by an MLP) maps a task description card to $r_\tau$.
• Training: The identifier is pretrained offline using "local supervision signals." By performing small, sparse interventions on a student model and measuring the resulting changes in capability profiles ($\Delta s$) and task utility ($\Delta U$), the model learns to predict which capabilities are beneficial for a specific task.
• Role: $r_\tau$ serves as a prior for allocation, prioritizing capabilities that improve task utility and treating drops in these specific dimensions as "harmful spillover."

2.2 On-the-Fly Capability-Targeted Data Generation

Instead of using a static dataset, ReAD generates supervision data dynamically during the distillation process.

• Templates: For each capability $c$, a template library defines instruction scaffolds, slots, and output constraints.
• Difficulty Control: A deterministic difficulty score $d_c(x)$ is computed based on factors like reasoning hops or code complexity. The system employs a curriculum that shifts from easy to hard examples as training progresses.
• Sampling: At each decision step, a capability label is sampled based on the current allocation policy, a prompt is instantiated, and the teacher model provides a completion for distillation.

2.3 Uncertainty-Aware Contextual Bandit for Allocation

ReAD frames the budget allocation as a sequential decision problem solved via a contextual bandit.

• Context ($x_t$): Includes the task requirement vector $r_\tau$, the current student's capability profile $s_{probe}(S_t)$, remaining budget $b_t$, and recent allocation history.
• Action Space: A finite set of allocation vectors $A(\tau)$, including perturbations of the previous allocation and sparse actions focused on top-$k$ capabilities.
• Proxy Reward ($\hat{R}_t$): After a distillation step, the system calculates:
  $$\hat{R}_t = r_\tau^\top \Delta s_{probe} - \beta \cdot \text{Spill}_t - \lambda \cdot \text{cost}_t$$
  Where $\text{Spill}_t$ penalizes regressions on task-essential capabilities.
• Policy Update: An ensemble of reward regressors (MLPs) estimates the expected reward and uncertainty ($\sigma$) for each candidate action. The next allocation is selected using an Upper Confidence Bound (UCB) rule:
  $$w_{t+1} = \arg \max_{w} \left( \mu(x_{t+1}, w) + \kappa \sigma(x_{t+1}, w) \right)$$
  This allows ReAD to balance exploitation (choosing high-reward allocations) with exploration (reducing uncertainty), shifting budget away from saturated or high-spillover capabilities as training progresses.

3. Key Contributions

1. Empirical Study of Cross-Capability Interactions: The authors provide a controlled study demonstrating that distilling for a single capability induces systematic, budget-dependent redistribution of performance across other capabilities, often causing negative transfer that worsens as the budget increases.
2. Budgeted Capability Distillation Formulation: They formalize the problem as allocating a fixed budget across interdependent capabilities to maximize utility while controlling side effects, moving beyond the assumption of independent capabilities.
3. Theoretical Analysis: The paper provides a local theoretical account showing that cross-capability transfer arises from non-orthogonal update directions in shared representations and that diminishing returns make continued allocation to a single capability wasteful.
4. ReAD Framework: The proposal of a reinforcement-guided framework that infers task requirements, generates targeted data on-the-fly, and uses an uncertainty-aware bandit to adaptively allocate resources.

4. Experimental Results

Experiments were conducted distilling Llama-3.3-70B into Llama-3.1-8B and Qwen2.5-72B into Qwen2.5-14B under token budgets of 20M and 150M.

• Utility Gains (RQ1): ReAD consistently outperformed six strong baselines (combining SFT/KD objectives with Resp/CoT/Logit signals) across all eight evaluated capabilities (General, Steerability, Reasoning, Math, Code, Tool Use, LCU, Multilingual). It achieved the strongest average capability profile, with significant improvements on bottleneck capabilities like Steerability and Math.
• Budget Efficiency (RQ2): ReAD demonstrated superior gain-spillover trade-offs. While single-capability KD achieved high target gains, it caused significant off-target degradation. ReAD maintained strong on-target gains while substantially reducing negative transfer.
• Adaptive Scheduling (RQ3): Comparisons against static mixtures, greedy heuristics, and grid-searched allocations showed that ReAD's adaptive bandit scheduling outperformed all alternatives, particularly at larger budgets (150M) where saturation and spillover are more pronounced.
• Generalization: ReAD showed positive transfer to held-out downstream tasks (XSTest) without requiring new capability heads, improving both safe-refusal and unsafe-refusal metrics compared to baselines.

5. Significance and Claims

The paper claims that ReAD represents a shift from treating capability distillation as a collection of independent tasks to a joint optimization problem that accounts for the intrinsic interdependence of LLM capabilities.

• Efficiency: By explicitly modeling diminishing returns and cross-capability interference, ReAD reduces "budget waste," ensuring that limited training tokens are spent on updates that provide genuine task-relevant utility rather than causing harmful side effects.
• Robustness: The framework is shown to be robust across different teacher-student pairs (Llama and Qwen families) and does not rely on a specific model architecture.
• Practicality: The method uses lightweight probes and a reusable requirement identifier, making it computationally feasible for deployment in resource-constrained MLaaS settings.

The authors conclude that effective capability distillation under fixed budgets requires reasoning about task-relevant capabilities, cross-capability interactions, and dynamic budget allocation simultaneously, a gap that ReAD fills.