Reinforced Multi-Teacher Selection (RL-KD)

• The paper introduces a framework that uses reinforcement learning to dynamically select and weight multiple teacher models for improved knowledge distillation.
• It details adaptive techniques such as softmax weighting, binary gating, and guidance-on-demand to tailor teacher contributions per training instance.
• Empirical results demonstrate significant gains across diverse tasks, including up to 12.2% accuracy improvement in out-of-distribution scenarios.

Reinforced Multi-Teacher Selection (RL-KD), also referred to as Adaptive Multi-Guidance Policy Optimization (AMPO) in the context of language modeling, denotes a class of reinforcement learning (RL) and policy-gradient-based methods that dynamically leverage multiple teacher models for knowledge distillation. These frameworks are characterized by adaptive, context-sensitive selection or weighting of teacher guidance per training instance or RL rollout, optimizing student model performance in complex or heterogeneous scenarios. RL-KD unifies the paradigm of multi-teacher knowledge transfer with reinforcement-based teacher policy selection and is established across diverse domains including LLM reasoning, visual recognition, image forgery detection, and natural language processing.

1. Formal Framework and Distillation Objective

RL-KD settings consider a student policy $\pi_\theta$ with learnable parameters $\theta$ trained under the supervision of a collection of $M$ fixed teacher policies or networks $\{\pi_{\phi_j}\}_{j=1}^M$. At each step, the student receives a task input (e.g., prompt, image, sequence), for which:

• Each teacher produces teacher-specific outputs (e.g., trajectories, logits, feature embeddings)
• The student is updated via an RL objective or knowledge distillation loss, parameterized by dynamically selected or weighted guidance from the teacher pool.

The essential multi-teacher distillation objective takes the form

$$\mathcal{L}_{MTKD}(i) = \mathcal{H}\bigl(\sigma(\mathbf{z}^S_i), \mathbf{y}_i\bigr) + \alpha\sum_{m=1}^M w_{l,i}^m\,\mathcal{D}_{KL}\left(\sigma(\mathbf{z}^S_i)\,\|\, \sigma(\mathbf{z}_i^{T_m}) \right) + \beta\sum_{m=1}^M w_{f,i}^m\,\|\mathbf{f}_i^S - \mathbf{f}_i^{T_m}\|_2^2,$$

where $\mathbf{z}^{T_m}_i$ are teacher logits, $\mathbf{f}_i^{T_m}$ are feature embeddings, and $w_{l,i}, w_{f,i}$ are RL-selected weights (Yang et al., 22 Feb 2025, Yuan et al., 2020).

In RL-based reasoning, batch augmentation and reward normalization are intertwined. AMPO, for example, formulates a unified mixed-polices surrogate loss over both on-policy and adaptively injected off-policy teacher rollouts (Yuan et al., 2 Oct 2025).

2. Reinforcement Learning Formulation for Teacher Selection

Teacher selection is modeled as a sequential Markov decision process, where the RL agent (policy network) observes a state vector encapsulating the current student state, teacher outputs, and (optionally) history, then outputs actions in the form of teacher weights or binary teacher-selection signals. Typical policy parameterizations include:

• Continuous softmax weighting $\mathbf{w}_t =$ softmax$(g(s_t;\theta))$, as in (Yuan et al., 2020, Yang et al., 22 Feb 2025)
• Binary gating $a_{ik} \in \{0,1\}$ for teacher $k$ per input $i$ (Yu et al., 7 Apr 2025)
• Guidance-on-demand flag $I$ indicating whether off-policy teacher solutions should be incorporated (Yuan et al., 2 Oct 2025)

State representations stack teacher-student features (e.g., features, KL divergences, cross-entropy losses, similarity scores) and may be concatenated across all teachers.

Policy learning commonly uses policy-gradient methods (REINFORCE). The agent is rewarded on the improvement of the student, typically via negative distillation loss or post-update task metrics:

$$r_i^m = -\big[ \text{TaskLoss}(\text{Student}) + \text{KD Loss for teacher } m \big]$$

Policy parameters are updated by gradient ascent on expected returns.

3. Adaptive Guidance, Comprehension-Based Selection, and “Guidance-on-Demand”

Distinct RL-KD implementations introduce advanced mechanisms to distribute teacher guidance:

• Guidance-on-demand: Teacher trajectories are injected into the training batch only when student on-policy rollouts uniformly fail above a threshold, preserving self-discovery but expanding exploration when necessary (Yuan et al., 2 Oct 2025).
• Comprehension-based guidance selection: Among available teacher solutions, only those most probable (i.e., comprehensible) to the student, as measured by the student’s conditional likelihood of the ground-truth answer given the teacher’s chain of thought, are selected (Yuan et al., 2 Oct 2025).
• Dynamic per-input weighting: RL assigns different per-instance weights to teachers according to difficulty, student-teacher compatibility, or observed improvement (Yang et al., 22 Feb 2025, Yuan et al., 2020, Yu et al., 7 Apr 2025).
• Hybrid on-policy/off-policy objectives: On augmented batches with both student and teacher-sourced samples, losses are normalized and aggregated with appropriate weighting per source (e.g., sequence-level for off-policy, token-level for on-policy) (Yuan et al., 2 Oct 2025).

4. Empirical Evaluation and Benchmark Performance

Empirical studies across domains demonstrate the efficacy of RL-KD relative to static or heuristic teacher weighting:

• Mathematical Reasoning (LLMs) (Yuan et al., 2 Oct 2025): AMPO delivers $+4.3\%$ accuracy over the GRPO baseline on in-distribution math tasks, $+12.2\%$ on out-of-distribution tasks, with markedly improved Pass@k under limited teacher examples.
• Visual Recognition (Yang et al., 22 Feb 2025): MTKD-RL outperforms equal-weighted baselines by $+1.56\%$ (CIFAR-100, 4-teacher), and achieves superior accuracy on ImageNet, object detection, and segmentation benchmarks.
• NLP (GLUE) (Yuan et al., 2020): RL-KD achieves up to $1.5$ F1 improvement over static averaging and regular fine-tuning, with maximal gains on ambiguous/hard instances.
• Forgery Detection (Yu et al., 7 Apr 2025): Re-MTKD surpasses prior SOTAs by $+8.7\%$ AUC on multi-tamper detection, with ablations validating RL-based dynamic teacher selection.
• Ablation Analyses: Each RL-KD component—adaptive replacement, sequence-level aggregation, comprehension-based selection, hybrid loss—shows critical necessity. Removal or heuristic substitution results in consistent, significant performance degradation (Yuan et al., 2 Oct 2025, Yu et al., 7 Apr 2025).

5. Theoretical Foundations and Generalizations

RL-KD can be formalized within the Hidden Utility Bandit (HUB) and Active Teacher Selection (ATS) frameworks (Freedman et al., 2023), supporting POMDP-style optimization over teacher selection:

• Teachers are conceptualized as arms with hidden utility (accuracy) functions $u^*(\cdot)$ and costs $c_t$.
• Selecting teachers is cast as planning or bandit optimization, balancing exploitation of accurate teachers and exploration to estimate teacher reliability.
• ATS yields Bayes-regret bounds $\mathcal{O}(\log T)$, and strategies are extendable to batch or contextual RL-KD settings, with the potential for POMDP-based belief updates and meta-learned teacher policies.

A plausible implication is that future RL-KD schemes may integrate cost-sensitive, hierarchical, or contextual teacher selection by extending ATS/HUB formalism with explicit student-teacher feedback loops.

6. Algorithmic Structure and Practical Implementation

The standardized RL-KD/AMPO training routine comprises:

1. On-policy sampling: Each batch or prompt is processed, and G student rollouts are collected.
2. Eligibility checking: If all rollouts fail (reward $<\tau$), adaptive replacement is triggered.
3. Teacher solution selection: Comprehension-based scores $r_p$ are computed for candidate off-policy solutions; top-$k$ are chosen for batch augmentation.
4. Batch augmentation: An augmented batch $𝒢_{aug} = \{ G-k$ on-policy, $k$ off-policy$\}$ is constructed.
5. Unified reward normalization: Rewards and advantages are re-normalized globally.
6. Mixed surrogate loss: A joint loss $\mathcal{J}_{Mixed}$ over on- and off-policy samples is computed.
7. Policy update: $\theta$ is updated by gradient ascent; the RL teacher-selection policy is updated according to policy gradient or actor-critic, as applicable.

This approach is robust to a wide variety of domains, architectures (CNNs/ViTs for vision, BERT-like for NLP, Transformers for LLMs), and output modalities, given per-task reward and distillation formulations (Yuan et al., 2 Oct 2025, Yang et al., 22 Feb 2025, Yuan et al., 2020, Yu et al., 7 Apr 2025).

7. Limitations, Extensions, and Open Directions

RL-KD, while empirically impactful, incurs increased computational overhead due to RL controller optimization and on-the-fly policy evaluation (e.g., AMPO $32\times4$ GPU-hours vs. GRPO $22\times4$ (Yuan et al., 2 Oct 2025)). Additionally, policy-gradient methods are sensitive to reward design, variance, and exploration choices; actor-critic or bandit-based variants may stabilize learning in large teacher pools (Yu et al., 7 Apr 2025, Freedman et al., 2023).

Potential extensions include:

• Hierarchical teacher selection (selecting groups/subsets)
• Contextual RL-KD for distribution/domain shift
• Regret-minimizing query batching (bandit-style constraints)
• Meta-learned (adaptive) reward function design
• Soft/probabilistic teacher weighting (Gumbel-Softmax, differentiable gating)

These directions are directly motivated by observations from ATS/HUB and visual/NLP RL-KD studies, suggesting that principled, RL-driven selection among heterogeneous knowledge sources enables greater generalization, data efficiency, and performance robustness in student models (Freedman et al., 2023, Yuan et al., 2 Oct 2025, Yu et al., 7 Apr 2025, Yuan et al., 2020, Yang et al., 22 Feb 2025).