Technical Summary: Hide to See: Reasoning-prefix Masking for Visual-anchored Thinking in VLM Distillation

Problem Statement

Recent "think-answer" approaches in Vision-Language Models (VLMs), such as Qwen3-VL-Thinking, significantly enhance reasoning performance by generating intermediate reasoning steps before the final answer. However, these long reasoning traces incur substantial computational costs, limiting their deployment in resource-constrained scenarios. Knowledge distillation (KD) offers a solution by transferring capabilities from large teacher models to compact student models.

A critical challenge in distilling think-answer VLMs is visual forgetting. As the student model generates a long reasoning trace, it tends to rely heavily on the accumulated textual reasoning prefixes (the previously generated tokens) to predict subsequent tokens. These prefixes often contain sufficient information to mimic the teacher's trajectory without requiring the student to re-examine the visual input. Consequently, the student learns to ignore visual evidence, leading to a degradation in visual-anchored reasoning capabilities. Existing distillation methods that transfer visual attention maps or token relations often fail to address this specific tendency to shortcut via textual cues.

Methodology: Masking-KD

The authors propose Masking-KD, a novel distillation framework designed to force the student model to anchor its reasoning in visual evidence by deliberately obscuring the most influential textual cues.

Core Mechanism

Instead of using the standard causal mask (which allows the student to see all previous tokens), Masking-KD employs a salient reasoning-prefix mask. This mask blocks access to both future tokens (standard causality) and specific salient reasoning prefixes that the student would otherwise rely on to imitate the teacher. By removing these high-influence textual cues, the student is compelled to utilize visual evidence to compensate for the missing information.

The framework consists of two adaptive components:

1. Token-wise Salient Reasoning-Prefix Masking:

   • Identification: Before the main distillation step, an auxiliary forward pass of the student model (under a standard causal mask) is performed. This generates a response-to-response attention map ($A^{resp}$), which reveals which previous tokens the student attends to most heavily when predicting the current token.
   • Selection: For each next-token prediction, the method identifies the top-$\rho$ most salient prefixes based on their attention weights. These are the tokens providing the strongest "shortcut" to the answer.
   • Masking: These identified salient prefixes are masked (set to $-\infty$ in the attention logits) during the actual distillation training step, preventing the student from attending to them.

2. Self-Paced Masking Budget Scheduling:

   • The difficulty of distilling each token varies. Tokens that are easily imitated (low discrepancy between teacher and student) provide weak learning signals if the student can simply copy the prefix.
   • The method calculates the token-wise reverse KL divergence between the teacher and student distributions during the auxiliary forward pass.
   • A cumulative ratio threshold ($\rho_n$) is adaptively assigned to each token based on this difficulty. Tokens that are easily imitated (low KL divergence) receive a higher masking budget (more prefixes are masked) to force the student to rely on visual cues. Tokens that are harder to learn receive less masking to ensure stable training.

Training Objective

The student is trained to minimize the reverse KL divergence between its predictive distribution (under the masked context) and the teacher's distribution (under the full context). The loss function is defined as:
$$\mathcal{L}_{Distill} = \frac{1}{N} \sum_{n=1}^{N} \sum_{y \in \mathcal{V}} p_s(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \tilde{\mathbf{M}}) \log \frac{p_s(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \tilde{\mathbf{M}})}{p_t(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \mathbf{M})}$$
Where $\tilde{\mathbf{M}}$ is the salient reasoning-prefix mask and $\mathbf{M}$ is the standard causal mask.

Key Contributions

1. Masking-KD Framework: A new distillation paradigm that enhances visual-anchored thinking by masking salient reasoning prefixes, forcing the student to utilize visual evidence rather than relying on textual shortcuts.
2. Token-wise Adaptive Masking: A mechanism that dynamically identifies and masks the most influential reasoning prefixes for each specific token prediction, acknowledging that the critical context varies across decoding steps.
3. Self-Paced Budget Scheduling: An adaptive strategy that adjusts the masking intensity based on the distillation difficulty (KL divergence) of each token, amplifying learning signals for easily imitated tokens.
4. Empirical Validation: Extensive experiments demonstrating that this approach outperforms existing VLM distillation methods and self-distillation baselines.

Experimental Results

The authors evaluated Masking-KD on multimodal reasoning benchmarks, including Geometry-3K, MathVista, We-Math, MMK12, MathVerse, LogicVista, and MMMU-Pro.

• Performance Gains: Masking-KD achieved state-of-the-art results across all model sizes (2B, 4B, and 8B). Notably, the distilled 4B model surpassed the undistilled 8B teacher, and the 2B model outperformed the undistilled 4B teacher.
• Comparison with Baselines: The method significantly outperformed standard response-level distillation, as well as methods that explicitly transfer visual patterns (e.g., LLaVA-KD, CompoDistill, Align-TI).
• Self-Distillation: The framework was effective in self-distillation settings (where the student serves as its own teacher), outperforming other self-distillation methods like OPSD.
• Visual Attention Analysis: Qualitative and quantitative analyses confirmed that Masking-KD models maintain higher visual attention ratios throughout the generation process compared to baselines, effectively mitigating the "visual forgetting" phenomenon. The models were observed to activate relevant image regions more strongly when textual cues were masked.

Significance and Claims

The paper claims that the primary bottleneck in distilling think-answer VLMs is the student's tendency to exploit textual shortcuts provided by accumulated reasoning prefixes, thereby neglecting visual inputs.

The significance of Masking-KD lies in its shift from explicitly transferring the teacher's visual knowledge (e.g., via attention maps) to guiding the student to develop its own visual-anchored reasoning process. By "hiding" the most obvious textual clues, the method forces the student to "see" the image to solve the problem. The authors conclude that this approach is more effective for VLM reasoning distillation than previous methods that focus solely on mimicking internal visual patterns, as it directly addresses the root cause of visual forgetting in long-chain reasoning tasks.