On-Policy Attention Distillation (OPAD)

• The paper introduces OPAD, which aligns internal attention distributions between teacher and student models to mitigate exposure bias and enhance structural learning signals.
• It employs a composite loss framework that integrates reinforcement learning, knowledge distillation, and Jensen–Shannon divergence-based attention alignment for precise supervision.
• Empirical results demonstrate significant performance boosts in multimodal VQA and reinforcement learning, with improvements scaling with input density and policy pool size.

On-Policy Attention Distillation (OPAD) is a class of techniques in machine learning and reinforcement learning that utilizes attention-based mechanisms for mutual or teacher-student knowledge transfer, focusing specifically on aligning internal attention or feature representations rather than, or in addition to, output distributions. OPAD has been applied in both large language modeling and reinforcement learning policy distillation contexts, targeting improvements in reasoning, perceptual grounding, and sample efficiency by providing finer-grained structural learning signals during post-training or joint policy optimization (Li et al., 4 Feb 2026, Yu et al., 2024).

1. Formal Foundation and Theoretical Objectives

In the context of multimodal LLMs (MLLMs), On-Policy Attention Distillation operates by matching a student model's internal attention distributions to those of a fixed, higher-capacity teacher model for the same prompt and self-sampled generation path. Formally, the student policy $\pi_\theta$ receives a multimodal prompt (e.g., text plus image or video) and autoregressively generates a token sequence $x_{P+1},\ldots,x_T$. At each generation step $t$, the Transformer’s top layer of $\pi_\theta$ outputs attention weights $\alpha_t \in \mathbb{R}^{t-1}$, forming a per-token attention policy $$p_\theta^t(i) = \alpha_{t,i} / \sum_{j=1}^{t-1} \alpha_{t,j}$$ for $i = 1,\ldots,t-1$.

On-policy refers to the property that the teacher’s attention distributions $p_\phi^t$ are computed by running the teacher $\pi_\phi$ on the same token paths sampled by the student, mitigating exposure bias and ensuring alignment in hidden representation spaces.

In reinforcement learning policy distillation, OPAD refers to the online mutual distillation of $N$ policies $\{\pi_i\}_{i=1}^N$ via attention-based aggregation. Each agent, observing state $s_t$, computes feature vectors $h_i(s_t)$ and logits or Q-values $o_i(s_t)$. At every step, each agent leverages a Decision-Attention mechanism to dynamically aggregate information from other agents, forming a "soft teacher" from group-derived targets based on weighted peer outputs and features (Yu et al., 2024).

2. Loss Functions and Optimization Objectives

The total OPAD objective in MLLMs combines three loss terms:

$$L_{\text{total}}(\theta) = L_{\text{RL}}(\theta) + \mu \cdot L_{\text{GKD}}(\theta) + \gamma_{\text{attn}} \cdot L_{\text{AttnDistill}}(\theta)$$

• $$L_{\text{RL}}(\theta) = -\mathbb{E}_{\tau \sim \pi_\theta} [\sum_t A_\tau \cdot \log \pi_\theta(a_t|s_t)]$$, where $A_\tau$ denotes (possibly group-relative) sequence advantage.
• $$L_{\text{GKD}}(\theta) = \mathbb{E}_{\tau}[\sum_t \mathrm{KL}(\pi_\theta(\cdot|s_t) \| \pi_\phi(\cdot|s_t))]$$, aligning output distributions (typically reversed KL or cross entropy).
• $$L_{\text{AttnDistill}}(\theta) = \mathbb{E}_{\tau}[\sum_{t=P+1}^T \mathrm{JSD}(p_\theta^t \, \| \, p_\phi^t)]$$, where $\mathrm{JSD}$ denotes Jensen-Shannon divergence.

In policy distillation, each policy $\pi_i$ minimizes:

$$\mathcal{L}^i = L_{\mathrm{RL}}(\pi_i) + \lambda_{\mathrm{dec}} \, \mathbb{E}_{s_t}\left[ D_{\mathrm{KL}}(\sigma(O^{\mathrm{tea}}_i)\|\sigma(o_i)) \right] + \lambda_{\mathrm{feat}} \, \mathbb{E}_{s_t}\|H^{\mathrm{tea}}_i - h_i(s_t)\|_2^2$$

Here, $O^{\mathrm{tea}}_i$ and $H^{\mathrm{tea}}_i$ are group-aggregated decision and feature targets, $\sigma(\cdot)$ is the (possibly temperature-scaled) softmax, and both decision-level and feature-level distillation terms are essential for performance (Yu et al., 2024).

3. Algorithmic Procedures and Integration

The OPAD algorithm for MLLMs proceeds as:

1. Initialize the student model $\pi_\theta$, with the teacher $\pi_\phi$ pre-trained and fixed; freeze the visual encoder.
2. Iterate:
   • Sample multimodal prompt $s_0$.
   • Generate rollout $\tau = \{s_t, a_t\}_{t=1}^T$ using $\pi_\theta$.
   • Compute rewards and advantages for $\tau$.
   • For each generation step $t$:
     • Extract student attention logits to compute $p_\theta^t$ (softmax over logits).
     • Run teacher on the same prefix to extract and normalize $p_\phi^t$.
     • Calculate token-level RL, KD, and attention distillation losses.
   • Aggregate per-step losses into $L_{\text{total}}(\theta)$.
   • Backpropagate and update $\theta$.

For reinforcement learning, OPAD additionally requires all $N$ student policies to interact with the environment in parallel, exchange intermediate representations, compute attention weights via a cross-attention mechanism, and aggregate peer outputs accordingly. No additional architectural modifications are necessary beyond enabling access to intermediate features/logits and implementing the Decision-Attention module (Yu et al., 2024).

4. Empirical Performance and Ablation Studies

Empirical results in MLLMs demonstrate that incorporating attention distillation (with $\gamma_{\mathrm{attn}} > 0$) yields improvements over pure knowledge distillation in 7 out of 8 image VQA tasks (up to +3.6 points on V*, +1.8 on MuirBench) and in 5 out of 7 long-video QA tasks (up to +4.4 on NExTQA, +2.6 on Video-MME). OPAD provides especially robust gains at higher image resolutions and video frame counts, with performance improvements scaling as input density increases (e.g., +6.3 points at 2048 tokens vs. +1.6 at 512 tokens) (Li et al., 4 Feb 2026).

For online policy distillation, OPAD delivers significant gains over independent learning, with +47% (Breakout, PPO), +51% (BeamRider, PPO), +31% (SpaceInvaders, DQN), and +41% (Breakout, DQN). Ablation of the attention module in favor of uniform aggregation causes marked drops in efficiency (+35% to +72% attention gain), and removing decision-level distillation severely degrades performance (–93.8% on SpaceInvaders, DQN) compared to only moderate drops for removing feature-level loss (Yu et al., 2024).

5. Theoretical and Practical Distinctions from Conventional Knowledge Distillation

Traditional knowledge distillation aligns student and teacher output token distributions, which becomes a weak learning signal when teacher and student have divergent hidden representations or peaky output distributions. OPAD, by supervising the latent attention or feature selection process, provides a dense, per-token or per-state signal, guiding the student to emulate not just outputs but the evidence selection or "where to look" strategies of the teacher.

In MLLMs, OPAD thereby addresses exposure bias in internal representations, helping the student model generalize the teacher's information allocation even in novel generation contexts (Li et al., 4 Feb 2026).

In policy distillation, OPAD's Decision-Attention based aggregation dynamically modulates peer contributions, averting premature policy homogenization and effectively capturing diverse knowledge across the policy pool. This yields not only faster convergence but improved final performance and scalability with increasing agent count (Yu et al., 2024).

6. Practical Recommendations and Limitations

Optimal performance with OPAD in MLLMs is typically attained with $\gamma_{\mathrm{attn}} \approx 0.5$ and tuning $\mu$ between $0.5$ and $5$. Freezing the visual encoder and ensuring strict on-policy trajectory matching for the teacher are crucial for stable results. In policy distillation, both decision-level (KL-based) and feature-level (MSE-based) losses are required, with decision-level distillation exhibiting the most pronounced effect. Increasing the policy pool size enhances representation richness, but practical considerations (e.g., capped attention head dimensions) are recommended for scaling.

A notable finding is the efficacy of OPAD even in “zero-think” settings, where the student produces only final answers, evidencing that attention alignment alone can drive major portions of performance increase (Li et al., 4 Feb 2026). This suggests broader applicability in scenarios that favor process-level over output-level supervision.

7. Impact and Future Directions

OPAD establishes attention distributions and inter-policy feature sharing as explicit first-class objects of supervision for both multimodal generative models and reinforcement learning agents. The dense structural signals provided by attention-based alignment support improved grounding, cross-modal perception, and more sample-efficient policy optimization. By mitigating exposure bias and enabling process-aware transfer from high-capacity teachers or peer groups, OPAD represents a principled alternative to traditional output-focused distillation strategies. Ongoing directions include scaling OPAD to larger agent pools, adapting the methodology to continuous control via distributional matching, and exploring applications in domains beyond vision-language and classical RL benchmarks (Li et al., 4 Feb 2026, Yu et al., 2024).