July 01, 2026
While knowledge distillation (KD) is widely adopted for training lightweight models by leveraging supervision from larger teacher models, relying solely on output token distributions has proven insufficient for compressing Multimodal Large Language Models (MLLMs). Since output tokens are a byproduct of the model attending to visual inputs, prior works have explored explicitly distilling attention to provide a direct supervisory signal. While promising, the precise utility of which attention signals to distill remains under-explored. In this work, we challenge the conventional reliance on prompt-to-vision attention by revealing that downstream performance correlates strongly with response-to-vision attention similarity to the teacher, but negligibly with that of prompt-conditioned attention. Furthermore, we observe that attention distributions exhibit significant variance across individual tokens, indicating that a uniform distillation objective is suboptimal. To this end, we introduce Token-level Response-visual Attention Guidance (TRAG), a distillation objective that 1) shifts the focus to response-to-vision signals and 2) employs token-specific objectives by adaptively weighting the Kullback-Leibler divergence based on attention entropy, effectively guiding the student to mirror the teacher’s precise visual focus. Extensive experimental results on multiple benchmarks demonstrate that TRAG significantly outperforms prior distillation baselines. Our code is available at https://github.com/jhjangjh/TRAG.
Multimodal Large Language Models (MLLMs) [1]–[4] have achieved remarkable performance, yet the continuous scaling of their architectures and datasets entails prohibitive computational costs. This deployment bottleneck has intensified research into MLLM compression [5], [6], where knowledge distillation (KD) [7] has emerged as a predominant paradigm, building upon its established effectiveness in standard LLM frameworks [8]–[10]. However, standard distillation objectives designed for text-only LLMs rely solely on output distributions and thus do not explicitly transfer how visual evidence is grounded during response generation [11]–[13], even though output tokens are fundamentally a byproduct of the model attending to visual inputs.
To provide a more direct supervisory signal, prior works have explored explicitly distilling cross-modal attention [12], [14]. While this explicit guidance shows promise, identifying which attention signals to distill and at what granularity to apply supervision remains critically under-explored. Consequently, we find that the design choices made in existing frameworks often yield suboptimal supervision, failing to capture the dynamic evidence allocation that occurs during response decoding. In this work, we systematically characterize cross-modal attention distillation and its failure modes under existing design choices.
In detail, our contributions can be summarized as follows:
Response-to-Vision Attention as the Critical Signal: We show that downstream performance correlates strongly with the alignment of Response-to-Vision attention between the student and the teacher model, whereas matching Prompt-to-Vision attention, which has been the primary focus of prior works, provides negligible utility. This identifies the generative phase, rather than the initial perception phase, as the critical signal for multimodal distillation.
Exposing the Limitations of Uniform Objectives: We observe that cross-modal attention patterns exhibit significant variance across individual response tokens depending on their specific grounding roles, ranging from diffuse, scene-level context for function words to sharp, localized focus for visual entities. This high variance indicates that applying a uniform distillation objective across all tokens is suboptimal, necessitating an adaptive, token-wise supervision framework.
Motivated by these findings, we propose Token-level Response-visual Attention Guidance (TRAG) (see Figure 1). TRAG shifts the distillation focus entirely to Response-to-Vision signals and introduces an adaptive, token-specific objective. By modulating the Kullback-Leibler (KL) divergence based on the teacher attention entropy, TRAG balances coverage and precision for each decoding step. Experiments on diverse multimodal benchmarks demonstrate that TRAG significantly outperforms competitive distillation baselines, achieving superior attention fidelity and robust performance gains.
Prior KD approaches for MLLMs largely follow two trajectories. One focuses on improving modality-specific representation alignment and architectural adaptations. LLaVADI [11] analyzes key components for MLLM distillation, while LLaVA-KD [15] extends logit-based distillation to the vision modality and enforces structural consistency among visual tokens. Related compression efforts further improve compact MLLMs through visual-encoder mixtures [16] or MoE-based architectures [17]. While effective, these approaches do not explicitly supervise how visual evidence is selected during response generation, leaving visual grounding largely implicit.
Seeking more direct guidance on cross-modal processing, another trajectory explores attention-based distillation. Align-KD [14] transfers Prompt-to-Vision attention at a selected decoder layer, while CompoDistill [12] supervises Prompt-to-Vision attention maps. Concurrently, Align-TI [18] studies token-interaction distillation for multimodal alignment. However, these objectives mainly emphasize prompt-conditioned or globally defined interaction signals, overlooking token-wise dynamics during response generation. This gap motivates attention distillation that better aligns with response-time cross-modal processing, harmonizing with language modeling and logit distillation objectives.
Attention has evolved into a key supervisory signal for knowledge distillation, serving as a medium to transfer both spatial focus and complex inductive biases. Early methodologies focused on establishing correspondence between attention maps or saliency patterns to guide the student’s spatial perception toward task-relevant visual cues [19], [20]. Subsequent developments introduced more structured forms of transfer, such as cross-modal structural constraints and sparse connectivity to highlight salient relationships [21], [22]. More recently, attention supervision has been integrated with representation transfer to better preserve the teacher’s intermediate reasoning pathways and feature hierarchies [23]. These advancements establish attention as a robust mechanism for capturing the teacher’s internal prioritization, providing the technical groundwork for specialized attention-based strategies in the multimodal domain.
Standard MLLMs encode an input image into visual tokens \(\mathbf{v}=\{v_1,\dots,v_N\}\) via a vision encoder and an MLP-based projection module [1]. The model then processes the concatenated sequence \(\mathbf{c} = [\mathbf{v};\mathbf{x};\mathbf{y}]\), where \(\mathbf{x}=\{x_1,\dots,x_{L_P}\}\) denotes the prompt token sequence and \(\mathbf{y}=\{y_1,\dots,y_{L_R}\}\) denotes the autoregressive response token sequence. Letting \(L=L_P+L_R\), we index the elements \(c_i\) of \(\mathbf{c}\) such that indices \(i\in \{1,...,N\}\), \(i\in \{N+1,...,N+L_P\}\), and \(i \in \{N+L_P+1,...,N+L\}\) map to the visual, prompt, and response tokens, respectively.
The generative process is driven by a Transformer decoder, whose self-attention at each layer \(l\in\{1,\dots,M\}\) produces an attention matrix \(\mathbf{A}_l\in\mathbb{R}^{(N+L)\times(N+L)}\) with entries \(a_{l,i,j}\) [24]. Under causal masking, each query index \(i\) can attend only to keys \(j\le i\). In particular, any textual query \(i\in\{N+1,\dots,N+L\}\) can attend to visual keys \(j\in\{1,\dots,N\}\) through cross-modal attention.
For any index set \(\mathcal{I}\subseteq\{N+1,\dots,N+L\}\) of textual queries, we define the cross-modal attention submatrix from \(\mathcal{I}\) to visual keys as: \[\mathbf{A}_l(\mathcal{I}) := \mathbf{A}_l[\mathcal{I},\, 1\!:\!N] \in \mathbb{R}^{|\mathcal{I}|\times N}, \label{eq:submatrix95def}\tag{1}\] where \(\mathbf{A}_l(\mathcal{I})[i,j]=a_{l,i,j}\) for \(i\in\mathcal{I}\) and \(j\in\{1,\dots,N\}\). In this paper, we consider two instantiations of \(\mathcal{I}\): Prompt-to-Vision attention uses \(\mathcal{I}_P=\{N+1,\dots,N+L_P\}\), whereas Response-to-Vision attention uses \(\mathcal{I}_R=\{N+L_P+1,\dots,N+L\}\).
Prompt-to-Vision Attention (\(\mathbf{A}_{P \to V}\), Figure 2-(a)): For each prompt token in \(\mathbf{x}\) (i.e., \(c_i\), where \(i \in \mathcal{I}_P\)), the distribution \(\{a_{l,i,j}\}_{j=1}^{N}\) captures its attention over visual tokens during instruction encoding. This represents how strongly the model integrates visual evidence when interpreting the prompt token.
Response-to-Vision Attention (\(\mathbf{A}_{R \to V}\), Figure 2-(b)): For each response token in \(\mathbf{y}\) (i.e., \(c_i\), where \(i \in \mathcal{I}_R\)), the distribution \(\{a_{l,i,j}\}_{j=1}^{N}\) captures which visual tokens the model attends to when generating the next response token. This represents a token-level trace of how visual context influences the response distribution at position \(i\).
In prior works, the distillation objective for transferring knowledge from a teacher model \(\mathcal{T}\) to a student model \(\mathcal{S}\) is defined as a composite loss consisting of language modeling, uni-modal distillation, and cross-modal attention distillation:
\[\mathcal{L} = \mathcal{L}_{\text{LM}} + \mathcal{L}_{\text{uni-modal}} + \mathcal{L}_{\text{cross-modal}}. \label{eq:L95total}\tag{2}\]
The language modeling loss is the token-level cross entropy (negative log-likelihood) of the target response \(\mathbf{y}\) conditioned on \(\mathbf{v}\) and \(\mathbf{x}\): \[\mathcal{L}_{\text{LM}} = -\sum_{t=1}^{L_R} \log p^{\mathcal{S}}(y_t \mid \mathbf{v}, \mathbf{x}, y_{<t}). \label{eq:L95LM}\tag{3}\]
Uni-modal distillation aligns the teacher and student within each modality. This objective applies logit distillation for tokens in the response, alongside an auxiliary loss \(\mathcal{L}_{\text{aux}}\): \[\mathcal{L}_{\text{uni-modal}} = \sum_{t=1}^{L_R} D_{KL} \left( p^{\mathcal{T}}(\cdot \mid \mathbf{v}, \mathbf{x}, y_{<t}) \parallel p^{\mathcal{S}}(\cdot \mid \mathbf{v}, \mathbf{x}, y_{<t}) \right) + \mathcal{L}_{\text{aux}}. \label{eq:L95uni}\tag{4}\] The first term performs response-level logit distillation over the target response sequence \(\mathbf{y}=\{y_t\}_{t=1}^{L_R}\) [7] by minimizing the KL divergence \(D_{KL}(\cdot \parallel \cdot)\) between the teacher and student next-token distributions at each generation step. The auxiliary loss \(\mathcal{L}_{\text{aux}}\) introduces additional uni-modal supervision, typically defined on visual-side outputs or intra-visual relations. In LLaVA-KD [15], \(\mathcal{L}_{\text{aux}}\) distills vision-token logits and matches affinity matrices among vision tokens. Similarly, CompoDistill [12] distills intermediate-layer visual self-attention by aligning vision-to-vision attention maps with a cosine-based loss.2 Overall, \(\mathcal{L}_{\text{aux}}\) strengthens modality-specific alignment by providing additional supervision beyond response-level logit distillation.
Cross-modal distillation aligns the teacher and student by matching their cross-modal attention maps, where textual queries attend to visual tokens. In CompoDistill [12], this objective is instantiated with a group-based layer matching scheme to handle depth differences between the teacher and student:
\[\mathcal{L}_{\text{cross-modal}}^{\text{(CompoDistill)}} = \sum_{l \in \mathcal{M}_\mathcal{S}} \mathcal{D}_{\text{cos}}\!\left( \mathbf{A}^{\mathcal{T}}_{G_l}(\mathcal{I}),\; \mathbf{A}^{\mathcal{S}}_{l}(\mathcal{I}) \right), \label{eq:L95cross}\tag{5}\] where \(\mathcal{D}_{\text{cos}}(\cdot,\cdot)\) denotes cosine distance, \(\mathcal{M}_{\mathcal{S}}\) denotes a subset of student layers selected for supervision, and \(\mathcal{I}\) denotes the textual query indices used to extract \(\mathbf{A}_l(\mathcal{I})\) (Eq. 1 ). For each student layer \(l\), where \(G_l\) denotes the set of teacher layers assigned to \(l\), the grouped teacher attention is \[\mathbf{A}^{\mathcal{T}}_{G_l}(\mathcal{I}) := \frac{1}{|G_l|}\sum_{g\in G_l} \mathbf{A}^{\mathcal{T}}_{g}(\mathcal{I}). \label{eq:teacher95group95avg}\tag{6}\] In prior work, including Align-KD [14] and CompoDistill [12], \(\mathcal{I}\) is typically chosen as the prompt index range \(\{N+1,\dots,N+L_P\}\) (i.e., they set \(\mathcal{I}\) to \(\mathcal{I}_P\)), focusing supervision on Prompt-to-Vision attention.
Building on Section 3, we revisit two common but under-discussed choices in attention-based cross-modal knowledge distillation. First, prior work often supervises attention on prompt tokens by default. Second, they typically devote substantial design effort to deciding which layers to distill from and how to align layers across teachers and students with different depths. In this section, we ask whether prompt-based supervision is truly optimal, and whether layer selection is as critical as often assumed compared to the granularity of the attention signal.
If Prompt-to-Vision attention alignment to the teacher is an optimal distillation target, then Prompt-to-Vision attention similarity should consistently accompany a stronger student. To examine this, we compute teacher-student attention similarity and analyze its correlation with benchmark performance across multiple students. We measure teacher–student cross-modal attention similarity as cosine similarity under teacher forcing on an unseen dataset VIGC-InstData [25].
The results show that prompt-to-vision alignment does not reliably track model capability. As shown in Figures 3-(a) and (c), in Prompt-to-Vision attention (\(\mathbf{A}_{P\to V}\)), similarity to the teacher does not consistently correlate with downstream performance. Notably, the Supervised Fine-Tuning (SFT) baseline can achieve higher \(\mathbf{A}_{P\to V}\) similarity than distilled models such as LLaVA-KD and CompoDistill, while still underperforming them on downstream benchmarks. This indicates that aligning attention during prompt processing alone is not a sufficient proxy for the cross-modal behavior that distinguishes stronger models.
In contrast, Response-to-Vision alignment is substantially more predictive. Figures 3-(b) and (d) show a strong positive correlation between Response-to-Vision attention (\(\mathbf{A}_{R\to V}\)) similarity and performance across students. This suggests that the teacher’s effective visual grounding is expressed more clearly during response generation than during prompt processing, motivating attention supervision that targets the generative phase.
Section 4.1 suggests that the key transferable signal emerges during generation, where Response-to-Vision attention (\(\mathbf{A}_{R\to V}\)) tracks downstream performance. The next question is the granularity at which this signal should be distilled. Most attention distillation methods adopt layer-level granularity through layer selection or mapping. We ask whether this layer-level choice is necessary by analyzing \(\mathbf{A}_{R\to V}\) across both decoder depth and individual response tokens.
Focusing on intermediate depths (relative depth in \([0.3,0.6]\)), where cross-modal interactions are most pronounced [12], [26]–[29], Figure 4 shows that adjacent layers exhibit highly similar \(\mathbf{A}_{R\to V}\) patterns, indicating substantial redundancy along depth. In contrast, \(\mathbf{A}_{R\to V}\) varies sharply across response tokens. When generating function words (e.g., articles), the model exhibits weak visual grounding with uniform, blurry attention, whereas content tokens show more diverse patterns ranging from sharply localized evidence to diffuse, scene-level grounding. This suggests that the variation in cross-modal grounding is primarily token-dependent rather than layer-dependent, motivating token-wise supervision that adapts to each response token. Additional examples are provided in Appendix 12.
Motivated by our findings that Response-to-Vision attention alignment is a strong predictor of downstream task proficiency, and that cross-modal attention patterns vary substantially across response tokens (Section 4), we propose Token-level Response-visual Attention Guidance (TRAG).
Building on the standard training objective (Eq. 2 ), we take the CompoDistill formulation (Eq. 5 ) as a starting point for the cross-modal distillation term, \(\mathcal{L}_{\text{cross-modal}}\). TRAG reformulates this baseline into a response-token-wise objective that guides the student’s Response-to-Vision attention using the teacher’s attention as a reference for each target response token.
Motivated by the observation in Section 4.1, we shift cross-modal attention supervision from the prompt to the response, where visual grounding is expressed during generation. Accordingly, we replace the prompt index set \(\mathcal{I}_P\) with the response range \(\mathcal{I}_R=\{N+L_P\!+\!1,\dots,N+L\}\): \[\mathcal{L}_{\text{cross-modal}}^{\text{(response)}} = \sum_{l\in\mathcal{M}_\mathcal{S}}\mathcal{D}_{\cos}\!\left(\mathbf{A}^{\mathcal{T}}_{G_l}(\mathcal{I}_R),\mathbf{A}^{\mathcal{S}}_{l}(\mathcal{I}_R)\right). \label{eq:L95response}\tag{7}\] This localizes attention supervision to response generation directly. Next, we refine the granularity of the distillation signal within this response range.
Prior attention-based distillation methods typically define a single uniform objective over the supervised query set, comparing the teacher and student attention maps using a fixed discrepancy measure. However, Section 4.2 shows that token-level attention distributions are highly heterogeneous, suggesting that a single uniform constraint provides limited resolution to preserve per-token distributional intent. Moreover, Section 4.2 indicates substantial redundancy across intermediate layers, motivating aggregation across layers rather than sensitive layer selection.
Accordingly, we first aggregate cross-modal attention over intermediate layers to obtain a representative attention distribution for each response query, and then decompose supervision into token-level losses. Let \(\mathcal{M}_ \mathcal{S}\) and \(\mathcal{M}_\mathcal{T}\) denote the sets of intermediate layers (relative depth in \([0.3,0.6]\)) used for aggregation in the student and teacher, respectively. For each response query index \(i\in\mathcal{I}_R\), we denote the per-layer attention vectors over \(N\) visual tokens as \(\mathbf{a}^{\mathcal{S}}_{l,i}:=\{a^{\mathcal{S}}_{l,i,j}\}_{j=1}^{N}\) and \(\mathbf{a}^{\mathcal{T}}_{l,i}:=\{a^{\mathcal{T}}_{l,i,j}\}_{j=1}^{N}\), corresponding to the rows of \(\mathbf{A}^{\mathcal{S}}_{l}(\mathcal{I}_R)\) and \(\mathbf{A}^{\mathcal{T}}_{l}(\mathcal{I}_R)\) restricted to visual keys (Eq. 1 ). We then obtain layer-averaged attention vectors by averaging over intermediate layers \(\mathcal{M}_ \mathcal{S}\) and \(\mathcal{M}_ \mathcal{T}\): \[\bar{\mathbf{a}}^{\mathcal{S}}_{i} := \frac{1}{|\mathcal{M}_\mathcal{S}|}\sum_{l\in\mathcal{M}_\mathcal{S}} \mathbf{a}^{\mathcal{S}}_{l,i}, \qquad \bar{\mathbf{a}}^{\mathcal{T}}_{i} := \frac{1}{|\mathcal{M}_\mathcal{T}|}\sum_{l\in\mathcal{M}_\mathcal{T}} \mathbf{a}^{\mathcal{T}}_{l,i}. \label{eq:layer95agg95ai}\tag{8}\]
Since these vectors are obtained by restricting attention to visual keys, their total mass over \(\{1,\dots,N\}\) is not necessarily one. We therefore renormalize them on the vision span to form proper distributions: \[\tilde{\mathbf{a}}^{\mathcal{S}}_{i} := \frac{\bar{\mathbf{a}}^{\mathcal{S}}_{i}}{\sum_{j=1}^{N}\bar{a}^{\mathcal{S}}_{i,j}+\epsilon}, \qquad \tilde{\mathbf{a}}^{\mathcal{T}}_{i} := \frac{\bar{\mathbf{a}}^{\mathcal{T}}_{i}}{\sum_{j=1}^{N}\bar{a}^{\mathcal{T}}_{i,j}+\epsilon}, \label{eq:vision95renorm}\tag{9}\] where \(\epsilon>0\) is a small constant for numerical stability. Finally, we define the token-level cross-modal objective as \[\mathcal{L}_{\text{cross-modal}}^{\text{(token-level)}} = \sum_{i\in\mathcal{I}_R} \mathcal{D}_{\text{guidance}}\!\left( \tilde{\mathbf{a}}^{\mathcal{T}}_{i},\; \tilde{\mathbf{a}}^{\mathcal{S}}_{i} \right), \label{eq:L95tokenlevel}\tag{10}\] where \(\mathcal{D}_{\text{guidance}}(\cdot,\cdot)\) is defined in the following subsection in detail.
Having decomposed the cross-modal objective into per-token terms (Eq. 10 ), we now derive a token-dependent coefficient that modulates the guidance behavior. For each response query index \(i\in\mathcal{I}_R\), we compute the entropy of the teacher’s normalized layer-aggregated attention distribution \(\tilde{\mathbf{a}}_{i}^{\mathcal{T}}\) (Eq. 9 ):
\[H(\tilde{\mathbf{a}}_{i}^{\mathcal{T}}) = -\sum_{j=1}^{N}\tilde{a}^{\mathcal{T}}_{i,j}\log \tilde{a}^{\mathcal{T}}_{i,j}, \label{eq:entropy95calculation}\tag{11}\] where \(\tilde{a}^{\mathcal{T}}_{i,j}\) denotes the weight on the \(j\)-th visual token (\(j\in\{1,\dots,N\}\)) in \(\tilde{\mathbf{a}}_{i}^{\mathcal{T}}\).
We map \(H(\tilde{\mathbf{a}}_{i}^{\mathcal{T}})\) to a bounded coefficient \(\lambda_i\in[0,1]\) via min–max normalization with clipping: \[\lambda_i = \mathrm{clip}\!\left( \frac{H(\tilde{\mathbf{a}}_{i}^{\mathcal{T}})-H_{\min}}{H_{\max}-H_{\min}+\epsilon}, \,0,\,1 \right), \label{eq:lambda}\tag{12}\] where \(H_{\min}\) and \(H_{\max}\) are running lower and upper bounds updated online as exponential moving averages (momentum \(\beta=0.995\)) of the mini-batch minimum and maximum of \(H(\tilde{\mathbf{a}}_{i}^{\mathcal{T}})\), and \(\epsilon>0\) is a small constant for numerical stability. A larger \(\lambda_i\) corresponds to more diffuse attention, which we leverage to parameterize token-wise guidance in the following objective.
To achieve specialized supervision, we define the guidance operator \(\mathcal{D}_{\text{guidance}}\) as a weighted combination of forward and reverse Kullback-Leibler (KL) divergences. This formulation leverages the asymmetry of KL divergence. Forward KL, \(D_{KL}(\tilde{\mathbf{a}}_{i}^{\mathcal{T}}\parallel\tilde{\mathbf{a}}_{i}^{\mathcal{S}})\), strongly penalizes the student when it assigns low probability to regions where the teacher has non-negligible mass, thereby promoting coverage. In contrast, reverse KL, \(D_{KL}(\tilde{\mathbf{a}}_{i}^{\mathcal{S}}\parallel\tilde{\mathbf{a}}_{i}^{\mathcal{T}})\), penalizes probability mass placed outside the teacher’s high-density regions more aggressively, which encourages concentration on dominant modes [30], [31].
Using the normalized coefficient \(\lambda_i\) defined in Eq. 12 , we combine the two asymmetric KL terms into a token-wise guidance divergence per response token: \[\mathcal{D}_{\text{guidance}}(\tilde{\mathbf{a}}_{i}^{\mathcal{T}}, \tilde{\mathbf{a}}_{i}^{\mathcal{S}}) = \lambda_i D_{KL}(\tilde{\mathbf{a}}_{i}^{\mathcal{T}} \parallel \tilde{\mathbf{a}}_{i}^{\mathcal{S}}) + (1 - \lambda_i) D_{KL}(\tilde{\mathbf{a}}_{i}^{\mathcal{S}} \parallel \tilde{\mathbf{a}}_{i}^{\mathcal{T}}). \label{eq:L95wkl}\tag{13}\] Substituting Eq. 13 into Eq. 10 yields the final cross-modal objective (Figure 1): \[\mathcal{L}_{\text{cross-modal}}^{(\text{TRAG})} = \sum_{i\in\mathcal{I}_R} \left[ \lambda_i\, D_{KL}(\tilde{\mathbf{a}}_{i}^{\mathcal{T}} \parallel \tilde{\mathbf{a}}_{i}^{\mathcal{S}}) + (1-\lambda_i)\, D_{KL}(\tilde{\mathbf{a}}_{i}^{\mathcal{S}} \parallel \tilde{\mathbf{a}}_{i}^{\mathcal{T}}) \right]. \label{eq:trag95cross95modal95final}\tag{14}\] This adaptive regulation overcomes the resolution limits of uniform objectives, enabling the student to capture visual evidence from localized to diffuse patterns.
We adopt a three-stage training schedule (DPT–SFT–DFT), consistent with prior MLLM distillation works [12], [15]. DPT trains the projection module on image–caption data with the language backbone frozen, SFT improves instruction following with the standard language modeling objective, and DFT re-applies teacher-guided distillation on the instruction data. The objective is stage-dependent: in DPT and DFT, we optimize \[\mathcal{L}_{\text{TRAG}} = \mathcal{L}_{\text{LM}} + \mathcal{L}_{\text{uni-modal}} + \mathcal{L}_{\text{cross-modal}}^{(\text{TRAG})}, \label{eq:trag95loss}\tag{15}\] whereas SFT minimizes only \(\mathcal{L}_{\text{LM}}\). In Eq. 15 , \(\mathcal{L}_{\text{uni-modal}}\) follows the uni-modal distillation components of LLaVA-KD [15], and \(\mathcal{L}_{\text{cross-modal}}^{(\text{TRAG})}\) is defined in Eq. 14 . Detailed formulations are provided in Appendix 14.
Both the teacher and student models use the same vision encoder, SigLIP-SO400M/14 at 384px resolution [32], and their LLM backbones are drawn from the Qwen1.5 and Qwen2.5 families. A two-layer MLP with GELU nonlinearity serves as the projection module to map visual features into the language embedding space [33].
We use LLaVA-Pretrain-558K for the DPT stage and LLaVA-Instruct-665K for the SFT and DFT stages [34]. Throughout all stages, the teacher model is frozen. For the student, we train only the projection module in the DPT stage, keeping the vision encoder and LLM frozen. In the SFT and DFT stages, we freeze the vision encoder and fully finetune the projection module and the LLM. All models are trained on 2 NVIDIA A100 GPUs with DeepSpeed ZeRO-2. Detailed hyperparameters are provided in the Appendix 8.
We evaluate our models on two benchmark suites covering general Visual Question Answering (VQA) and compositional reasoning (CR). The general VQA suite consists of eight benchmarks, including GQA [35], ScienceQA [36], TextVQA [37], MME [38], MMBench in English and Chinese [39], POPE [40], and MMMU [41]. We further evaluate on three CR benchmarks: SugarCrepe [42], BiVLC [43], and Winoground [44]. Detailed evaluation metrics and setup are described in the Appendix 9.
| Method | LLM | # Samples | Visual Question Answering Benchmarks | \(Avg_6\) | \(Avg_8\) | |||||||||||
| GQA | ScienceQA | TextVQA | MME | MMB | \(\text{MMB}^{\text{CN}}\) | POPE | MMMU | |||||||||
| Bunny | Phi2-2.7B | 2.6 M | 62.5 | 70.9 | 56.7 | 74.4 | 68.6 | 37.2 | - | 38.2 | - | - | ||||
| Imp-3B | Phi2-2.7B | 1.5 M | 63.5 | 72.8 | 59.8 | - | 72.9 | 46.7 | - | - | - | - | ||||
| MobileVLM | MobileLLaMA-2.7B | 1.2 M | 59.0 | 61.0 | 47.5 | 64.4 | 59.6 | - | 84.9 | - | 62.7 | - | ||||
| MoE-LLaVA | Phi2-2.7B | 2.2 M | 62.6 | 70.3 | 57.0 | - | 68.0 | - | 85.7 | - | - | - | ||||
| MiniCPM-V | MiniCPM-2.4B | 570 M | 51.5 | 74.4 | 56.6 | 68.9 | 64.0 | 62.7 | 79.5 | - | 65.8 | - | ||||
| TinyLLaVA | Phi2-2.7B | 1.2 M | 62.0 | 69.1 | 59.1 | 73.2 | 66.9 | - | 86.4 | 38.4 | 69.4 | - | ||||
| \(^\dag\) | Qwen1.5-4B | 1.2 M | 62.7 | 70.2 | 60.0 | 70.5 | 67.5 | 67.2 | 86.9 | 38.6 | 69.6 | 65.4 | ||||
| \(^\dag\) | Qwen2.5-3B | 1.2 M | 63.0 | 76.1 | 60.1 | 71.6 | 73.1 | 70.4 | 87.8 | 41.3 | 71.9 | 67.9 | ||||
| LLaVADI | MobileLLaMA-2.7B | 1.2 M | 61.4 | 64.1 | 50.7 | 68.8 | 62.5 | - | 86.7 | - | 65.7 | - | ||||
| Imp-2B | Qwen1.5-1.8B | 1.5 M | 61.9 | 66.1 | 54.5 | 65.2 | 63.8 | 61.3 | 86.7 | - | 66.3 | - | ||||
| Bunny-2B | Qwen1.5-1.8B | 2.6 M | 59.6 | 64.6 | 53.2 | 65.0 | 59.1 | 58.5 | 85.8 | - | 64.5 | - | ||||
| Mini-Gemini-2B | Gemma-2B | 2.7 M | 60.7 | 63.1 | 56.2 | 67.0 | 59.8 | 51.3 | 85.6 | 31.7 | 65.4 | - | ||||
| MoE-LLaVA-2B | Qwen1.5-1.8B | 2.2 M | 61.5 | 63.1 | 48.0 | 64.6 | 59.7 | 57.3 | 87.0 | - | 63.9 | - | ||||
| \(^\dag\) | Qwen1.5-1.8B | 1.2 M | 60.9 | 65.2 | 47.7 | 61.3 | 57.8 | 56.2 | 83.4 | 34.8 | 62.7 | 58.4 | ||||
| \(^\dag\) | Qwen2.5-1.5B | 1.2 M | 61.7 | 70.9 | 58.4 | 70.4 | 70.6 | 64.0 | 86.0 | 39.2 | 69.7 | 65.1 | ||||
| \(^\dag\) | Qwen1.5-1.8B | 1.2 M | 60.9 | 64.6 | 53.3 | 66.3 | 65.8 | 63.1 | 86.5 | 32.7 | 66.2 | 61.6 | ||||
| LLaVADI | MobileLLaMA-1.4B | 1.2 M | 55.4 | 56.0 | 45.3 | 58.9 | 55.0 | - | 84.7 | - | 59.2 | - | ||||
| LLaVA-MoD | Qwen2-1.5B | 5 M | 58.8 | 69.2 | 59.9 | 69.2 | 68.9 | 64.4 | 87.2 | - | 68.8 | - | ||||
| Align-KD | MobileLLaMA-1.4B | 3.6 M | 60.1 | 67.7 | 53.1 | 65.1 | 57.5 | - | 87.0 | - | 65.0 | - | ||||
| \(^\dag\) | Qwen1.5-1.8B | 1.2 M | 61.2 | 66.5 | 53.5 | 67.0 | 64.5 | 63.0 | 85.5 | 34.1 | 66.4 | 61.9 | ||||
| LLaVA-KD | Qwen1.5-1.8B | 1.2 M | 62.3 | 64.7 | 53.4 | 69.1 | 64.0 | 63.7 | 86.3 | 33.6 | 66.6 | 62.1 | ||||
| LLaVA-KD | Qwen2.5-1.5B | 1.2 M | 62.5 | 71.6 | 59.7 | 70.0 | 71.0 | 66.6 | 86.7 | 35.8 | 70.2 | 65.4 | ||||
| Qwen1.5-1.8B | 1.2 M | 61.7 | 66.9 | 55.6 | 67.0 | 66.6 | 63.3 | 86.8 | 35.2 | 67.4 | 62.9 | |||||
| Qwen2.5-1.5B | 1.2 M | 62.6 | 73.2 | 60.0 | 71.2 | 70.8 | 69.3 | 87.2 | 38.7 | 70.8 | 66.6 | |||||
| \(^\dag\) | Qwen1.5-0.5B | 1.2 M | 56.7 | 60.9 | 46.6 | 59.4 | 55.2 | 51.8 | 83.9 | 31.9 | 60.4 | 55.8 | ||||
| \(^\dag\) | Qwen2.5-0.5B | 1.2 M | 58.6 | 61.2 | 47.9 | 62.3 | 59.3 | 52.4 | 85.1 | 30.4 | 62.4 | 57.1 | ||||
| \(^\dag\) | Qwen1.5-0.5B | 1.2 M | 56.8 | 60.4 | 46.1 | 58.9 | 54.7 | 49.7 | 84.9 | 32.3 | 60.3 | 55.5 | ||||
| LLaVA-MoD | Qwen2-0.5B | 5 M | 56.6 | 61.1 | 57.1 | 67.0 | 58.7 | 54.1 | - | - | - | - | ||||
| \(^\dag\) | Qwen1.5-0.5B | 1.2 M | 59.1 | 61.1 | 48.2 | 63.8 | 59.6 | 54.9 | 85.7 | 34.0 | 62.9 | 58.3 | ||||
| LLaVA-KD | Qwen1.5-0.5B | 1.2 M | 59.6 | 60.6 | 49.9 | 64.5 | 60.1 | 55.5 | 85.9 | 30.2 | 63.4 | 58.2 | ||||
| LLaVA-KD | Qwen2.5-0.5B | 1.2 M | 59.8 | 60.6 | 52.0 | 64.7 | 61.3 | 57.0 | 86.4 | 28.3 | 64.1 | 58.7 | ||||
| Qwen1.5-0.5B | 1.2 M | 59.5 | 62.1 | 48.9 | 65.2 | 60.9 | 54.6 | 86.5 | 30.9 | 63.8 | 58.5 | |||||
| Qwen2.5-0.5B | 1.2 M | 60.2 | 62.7 | 52.0 | 68.1 | 61.7 | 57.7 | 87.6 | 31.3 | 65.3 | 60.2 | |||||
Table 1 demonstrates that TRAG consistently improves performance over the PT-SFT baseline across all tested backbones and student scales. The performance gains are most significant in the \(\le\) 0.5B regime. For the Qwen2.5-0.5B student, TRAG increases \(Avg_6\) from 62.4 to 65.3 and \(Avg_8\) from 57.1 to 60.2 compared to the PT-SFT baseline. This result also surpasses LLaVA-KD on the same backbone by 1.2 points in \(Avg_6\) and 1.5 points in \(Avg_8\). Similar robust improvements are maintained for the Qwen1.5-0.5B backbone, where TRAG achieves the highest scores in both average metrics.
In the \(\le\) 2B regime, TRAG continues to exhibit superior performance. Using the Qwen2.5-1.5B backbone, TRAG achieves 70.8 in \(Avg_6\) and 66.6 in \(Avg_8\), outperforming both the TinyLLaVA baseline and LLaVA-KD. Notably, TRAG surpasses LLaVA-MoD (68.8 in \(Avg_6\)) despite using 76% fewer training samples (1.2M vs. 5M).
Compared to CompoDistill, which relies on uniform prompt-conditioned attention supervision, TRAG achieves stronger overall results by transitioning to a decomposed token-level objective targeting response-conditioned attention, e.g., improving \(Avg_8\) from 61.9 to 62.9 on Qwen1.5-1.8B. This shift indicates that grounding cues manifested during decoding are particularly informative for distillation when supervised at token-level granularity. By leveraging this signal in the response phase, TRAG provides a more targeted supervision pathway, yielding robust improvements in downstream VQA across backbones.
Table ¿tbl:table:compo95result? reports reasoning performance on SugarCrepe, BiVLC, and Winoground. TRAG substantially strengthens compositional robustness, particularly in low-capacity regimes. For Qwen1.5-0.5B and 1.8B, TRAG improves the average scores to 70.3 and 79.8, exceeding the strongest prior baseline, CompoDistill, by +1.2 and +3.0 points, respectively. Similar gains are observed in the Qwen2.5 family, where the 1.5B student reaches an average of 85.0, narrowing the gap with the 3B teacher to 0.9 points. These results highlight that token-wise guidance during decoding enhances the student’s ability to ground compositional judgments in the right visual evidence, leading to robust gains on compositional reasoning benchmarks.
We further evaluate TRAG on the MobileLLaMA backbone, where it improves the 1.4B student over the PT-SFT baseline by \(+3.0\) VQA Avg and \(+7.3\) CR Avg. Detailed results are provided in Appendix 13.
3.5pt
c|c|c|*3wc1.8cm|c & & & &
& & & SugarCrepe & BiVLC & Winoground &
&TinyLLaVA &4B &87.3 &93.2 &70.1 &83.5
[1pt/1pt] &TinyLLaVA & & 74.8 & 83.8 & 62.6 & 73.7
&KD & & 76.1 & 85.3 & 61.5 & 74.3
&LLaVA-KD & & 76.6 & 85.5 & 60.9 & 74.3
&CompoDistill & & 81.7 & 86.5 & 62.2 & 76.8
& & &83.1 &89.1 &67.2 &79.8
[1pt/1pt] &TinyLLaVA & & 52.8 & 56.5 & 50.2 & 53.2
&KD & & 51.0 & 50.1 & 52.1 & 51.0
&LLaVA-KD & & 66.9 & 74.5 & 51.7 & 64.4
&CompoDistill & & 72.1 & 78.6 & 56.6 & 69.1
& & &73.4 &80.2 &57.3 &70.3
&TinyLLaVA &3B &88.1 &94.6 &75.0 &85.9
[1pt/1pt] &TinyLLaVA & &87.9 &93.3 &70.3 &83.8
& & &87.3 &93.8 &73.9 &85.0
[1pt/1pt] &TinyLLaVA & &75.1 &85.1 &57.0 &72.4
& & &78.9 &87.1 &59.6 &75.2
We evaluate attention fidelity on VIGC-InstData [25] using layer-wise teacher–student \((A_{R\rightarrow V})\) cosine similarity during autoregressive generation (without teacher forcing), following the similarity definition in Section 4.1 (Figure 5). At the 30-60% layer depth, where cross-modal grounding is most intensive [12], [26], TRAG achieves an attention fidelity of approximately 0.72, significantly surpassing the SFT baseline (0.18) as well as LLaVA-KD (0.32) and CompoDistill (0.41). This elevated fidelity indicates that by supervising response-conditioned attention at the token level, TRAG successfully replicates the teacher’s grounding logic throughout the generative process. Notably, this internal consistency directly correlates with the performance gains observed in Tables 1 and ¿tbl:table:compo95result?. These results confirm that maintaining high fidelity to the teacher’s internal attention patterns is fundamental to achieving superior visual understanding and compositional robustness in low-capacity models.
We report the main ablations on query-token selection and attention matching below. Additional ablations on relative-depth selection and response-supervision granularity are provided in Appendix 13.
| Query Token Type | VQA \(Avg\) | CR \(Avg\) |
|---|---|---|
| Prompt | 58.6 | 69.2 |
| Prompt + Response | 58.4 | 71.2 |
| Response (Ours) | 60.2 | 75.2 |
r0.5
4pt
Table 2 compares query token choices for forming visual-attention targets. Supervising only prompt tokens yields lower averages (VQA: 58.6, CR: 69.2) than using response tokens (VQA: 60.2, CR: 75.2). Using both prompt and response tokens does not improve VQA (58.4 vs.) and provides only a modest gain on CR (71.2 vs.), remaining substantially below response-only supervision. Overall, the best performance is achieved when supervision is concentrated on response tokens, suggesting that the most informative cross-modal signal emerges during response generation.
| Attention Matching | VQA \(Avg\) | CR \(Avg\) |
|---|---|---|
| MSE | 58.2 | 74.4 |
| Cosine Similarity | 58.7 | 72.5 |
| Forward KL | 57.6 | 73.9 |
| Reverse KL | 58.1 | 73.4 |
| JSD | 57.9 | 71.6 |
| Weighted Sum KL (Ours) | 60.2 | 75.2 |
r0.5
4pt
Table 3 compares alternative discrepancy measures for matching teacher–student response-to-vision attention. Among the fixed objectives, cosine similarity yields the strongest VQA average (58.7), while MSE performs best on CR (74.4), indicating that no single fixed measure dominates across benchmarks. The one-sided KL variants (forward or reverse) and the symmetric JSD objective also fall short of the best fixed baselines on at least one benchmark. In contrast, our entropy-aware weighted combination of forward and reverse KL achieves the best results on both VQA (60.2) and CR (75.2), improving over the strongest fixed baselines by +1.5 and +0.8, respectively. Overall, these gains suggest that effective attention distillation benefits from token-adaptive guidance rather than a globally fixed matching rule.
| Teacher | Student | VQA \(Avg\) | CR \(Avg\) |
|---|---|---|---|
| Qwen1.5-0.5B (SFT) | 55.8 | 53.2 | |
| Qwen1.5-1.8B | Qwen1.5-0.5B | 57.3 | 63.0 |
| Qwen1.5-4B | 58.5 | 70.3 | |
| Qwen2.5-0.5B (SFT) | 57.1 | 72.4 | |
| Qwen2.5-1.5B | Qwen2.5-0.5B | 58.3 | 74.3 |
| Qwen2.5-3B | 60.2 | 75.2 | |
r0.5
Table 4 examines the impact of teacher capacity on distillation performance for a fixed 0.5B student model. Across both Qwen1.5 and Qwen2.5 backbones, TRAG consistently improves over the SFT baseline, with larger teachers yielding greater gains, especially on compositional reasoning. For Qwen1.5, CR Avg rises from 53.2 to 63.0 and 70.3 as the teacher scales from SFT to 1.8B and 4B; for Qwen2.5, it similarly increases from 72.4 to 74.3 and 75.2 as the teacher scales from SFT to 1.5B and 3B. These results indicate that TRAG effectively captures and transfers the enhanced grounding capabilities of larger teachers, demonstrating favorable scalability with respect to teacher model size.
We presented TRAG, a response-token-wise attention guidance objective for distilling MLLMs. By shifting cross-modal supervision from Prompt-to-Vision to Response-to-Vision attention, TRAG directly targets visual grounding during generation. It further decomposes supervision into token-level objectives and adaptively balances forward and reverse KL divergences according to the teacher’s attention entropy. Integrated into a standard MLLM distillation pipeline, TRAG consistently improves performance on general VQA and compositional reasoning benchmarks while also improving teacher-student attention fidelity.
This work was supported by Samsung Electronics Co., Ltd (IO251223-14934-01) and the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.RS-2022-II220184, Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics, and No.RS-2021-II211381, Development of Causal AI through Video Understanding and Reinforcement Learning, and Its Applications to Real Environments).
Supplementary Material for
Token-level Response-visual Attention Guidance
for Multimodal LLMs Knowledge Distillation
Contents
Figure 6:
.
This section provides the comprehensive training configurations and hyperparameters for the three-stage training pipeline of TRAG.
As detailed in Table 5, we utilize the SigLIP-SO400M/14 vision encoder with an input image resolution of 384\(\times\)384, yielding a vision sequence length of 729. The maximum sequence length for the LLM is set to 2048. The vision encoder remains frozen across all three training stages.
During the Distilled Pre-Training (DPT) stage, only the projection module is trained. We use the LLaVA-Pretrain-558K dataset for this stage, utilizing a learning rate of 1e-3 and a global batch size of 256. In the subsequent Supervised Fine-Tuning (SFT) and Distilled Fine-Tuning (DFT) stages, both the projection module and the LLM backbone are fully finetuned. Both stages use the LLaVA-Instruct-665K dataset and employ a reduced learning rate of 2e-5 with a global batch size of 128.
Across all stages, training is conducted for 1.0 epoch using the AdamW optimizer (\(\beta_1 = 0.9, \beta_2 = 0.99\)). The learning rate follows a cosine decay schedule with a warm-up ratio of 0.03, and weight decay is set to 0.0. All models are trained using BFloat16 numerical precision and DeepSpeed ZeRO-2 optimization.
| Configuration | Distilled | Supervised | Distilled |
| Pre-Training | Fine-Tuning | Fine-Tuning | |
| Vision Encoder | |||
| Projection Module | |||
| LLM Backbone | |||
| Training data | LLaVA-Pretrain-558K | LLaVA-Instruct-665K | LLaVA-Instruct-665K |
| Vision Encoder | SigLIP-SO400M/14@384 | ||
| Image resolution | 384 \(\times\) 384 | ||
| Vision encoder sequence length | 729 | ||
| LLM sequence length | 2048 | ||
| Optimizer | AdamW | ||
| Optimizer hyperparameters | \(\beta_1 = 0.9,\;\beta_2 = 0.99\) | ||
| Learning rate | \(1\text{e}{-3}\) | \(2\text{e}{-5}\) | \(2\text{e}{-5}\) |
| Learning rate scheduler | Cosine decay | ||
| Weight decay | 0.0 | ||
| Training epochs | 1.0 | ||
| Warm-up ratio | 0.03 | ||
| Global batch size | 256 | 128 | 128 |
| Numerical precision | BFloat16 | ||
| Model parallelism | ZeRO-2 | ||
-1pt 0pt 7pt
For the general VQA benchmarks, we strictly adhere to the evaluation configurations established by LLaVA-1.5 [34]. In contrast, compositional reasoning benchmarks are not directly amenable to standard VQA-style evaluation, as they are formulated as image-caption pairs rather than explicit question-answer pairs. We therefore adopt the multiple-choice evaluation protocol originally designed to assess GPT-4V in SugarCrepe [42], and apply the same procedure to BiVLC [43] and Winoground [44].
Specifically, each instance in these compositional reasoning benchmarks comprises one or two images paired with a positive (matching) caption and a negative (mismatching) caption. As illustrated in Figure 7, the model is prompted to select the index of the caption that best describes the visual content. As part of this adopted protocol, to mitigate positional bias, we evaluate each instance twice by swapping the order of the candidate captions (e.g., placing the positive caption as option (1) in the ‘Positive-First Prompt’ setting, and as option (2) in the ‘Negative-First Prompt’ setting). The model’s response is considered correct only if it successfully outputs the exact index corresponding to the positive caption.
We report benchmark-specific metrics following the standard evaluation protocol for each respective dataset. For GQA [35], TextVQA [37], and MMMU [41], accuracy serves as the primary evaluation metric. For ScienceQA [36], we report accuracy exclusively on the image-context subset to focus on multimodal evaluation. For MME [38], we report the official benchmark score. For MMBench and \(\text{MMBench}^{\text{CN}}\) [39], we report the accuracy scores obtained directly from the official submission server. For POPE [40], we report the F1-score, aligning with prior works such as LLaVA-MoD [17].
For SugarCrepe [42], BiVLC [43], and Winoground [44], we evaluate the accuracy under both the positive-first and negative-first prompt configurations. We report the average of these two accuracies for each predefined split within the benchmarks.
In this section, we revisit the attention distillation objective employed in CompoDistill [12] and reinterpret its formulation by decomposing it into uni-modal and cross-modal components.
Let \(\mathbf{A}_l \in \mathbb{R}^{(N+L)\times(N+L)}\) denote the self-attention matrix at decoder layer \(l\), where the initial \(N\) indices correspond to visual tokens and the subsequent \(L\) indices correspond to textual tokens. In CompoDistill, supervision is applied to a specific attention submatrix where the queries consist of the sequence preceding the response generation (i.e., visual and prompt tokens) and the keys are restricted to the visual tokens. Formally, for each supervised layer \(l\), the target attention block is defined as:
\[\mathbf{A}_l[1\!:\!N+L_P,\;1\!:\!N]. \label{eq:compodistill95block}\tag{16}\]
Following their grouped layer matching scheme, let \(\mathcal{M}_{\mathcal{S}}\) denote the set of supervised student layers, and \(G_l\) denote the assigned group of teacher layers for a given student layer \(l \in \mathcal{M}_{\mathcal{S}}\). The aggregated teacher attention block is computed as:
\[\mathbf{A}^{\mathcal{T}}_{G_l}[1\!:\!N+L_P,\;1\!:\!N] = \frac{1}{|G_l|} \sum_{g\in G_l} \mathbf{A}^{\mathcal{T}}_{g}[1\!:\!N+L_P,\;1\!:\!N].\]
Using this notation, the original attention distillation objective in CompoDistill can be formulated as: \[\mathcal{L}^{(\text{CompoDistill})}_{\text{attn}} = \sum_{l\in\mathcal{M}_{\mathcal{S}}} \mathcal{D}_{\cos} \!\left( \mathbf{A}^{\mathcal{T}}_{G_l}[1\!:\!N+L_P,\;1\!:\!N], \mathbf{A}^{\mathcal{S}}_{l}[1\!:\!N+L_P,\;1\!:\!N] \right), \label{eq:compodistill95attn95appendix}\tag{17}\] where \(\mathcal{D}_{\cos}(\cdot,\cdot)\) denotes the cosine distance.
This supervised attention block can be partitioned based on the modality of the queries. Specifically, the query index range \(1:N+L_P\) encompasses both the visual token span (\(1:N\)) and the prompt token span (\(N+1:N+L_P\)). Consequently, the attention block can be decomposed into two distinct components:
\[\begin{align} \mathbf{A}_l[1\!:\!N+L_P,\;1\!:\!N] = \begin{bmatrix} \mathbf{A}_l[1\!:\!N,\;1\!:\!N] \\ \mathbf{A}_l[N+1\!:\!N+L_P,\;1\!:\!N] \end{bmatrix}. \end{align}\]
The upper submatrix, \(\mathbf{A}_l[1:N, 1:N]\), represents attention from visual queries to visual keys. As both queries and keys operate within the visual modality, this component functions strictly as a uni-modal (vision-to-vision) alignment term.
The lower submatrix, \(\mathbf{A}_l[N+1:N+L_P, 1:N]\), captures attention from prompt queries to visual keys. Following the notation established in the main paper, this corresponds precisely to the Prompt-to-Vision attention submatrix:
\[\mathbf{A}_l(\mathcal{I}_P) := \mathbf{A}_l[\mathcal{I}_P,\;1\!:\!N], \qquad \mathcal{I}_P=\{N+1,\dots,N+L_P\}.\]
Because this submatrix governs the interaction between textual queries and visual keys, it serves as the cross-modal attention alignment term.
Therefore, while CompoDistill aggregates attention supervision into a single block, this formulation inadvertently conflates two distinct functional roles: uni-modal visual self-attention and cross-modal Prompt-to-Vision attention. In this work, our primary focus lies in the cross-modal component, as it explicitly dictates how textual tokens attend to and extract visual evidence. Building upon this decomposition, TRAG explicitly isolates the cross-modal interaction and shifts the supervisory signal from Prompt-to-Vision to Response-to-Vision attention, effectively providing a more targeted distillation objective.
We first report empirical computational costs in Table 6. All measurements are conducted on 2 NVIDIA A100 80GB GPUs under the same training configuration. Compared with the non-attention-based LLaVA-KD baseline, TRAG incurs additional overhead due to attention extraction and token-wise adaptive KL computation. Specifically, TRAG requires 52.8 hours in total and 67.3 GB peak memory per GPU, corresponding to an overhead of +3.0 hours and +17.2 GB/GPU over LLaVA-KD. However, compared with the attention-based CompoDistill baseline, TRAG reduces both training time and memory usage by aggregating attention across depth before applying the adaptive KL objective, saving 2.4 hours and 3.9 GB/GPU.
| Method | Peak Memory (GB/GPU) \(\downarrow\) | PT/DPT (h) \(\downarrow\) | SFT (h) \(\downarrow\) | DFT (h) \(\downarrow\) | Total (h) \(\downarrow\) | ||
| PT-SFT | 36.7 | 4.7 | 6.9 | - | 11.6 | ||
| LLaVA-KD [15] | 50.1 | 14.9 | 6.8 | 28.1 | 49.8 | ||
| CompoDistill [12] | 71.2 | 12.5 | 7.3 | 35.4 | 55.2 | ||
| (Ours) | 67.3 | 15.5 | 6.8 | 30.5 | 52.8 | ||
5pt
We further analyze the asymptotic computational overhead of TRAG by decomposing the distillation objectives. Let \(N\) denote the number of visual tokens, \(L_P\) the prompt length, \(L_R\) the response length, \(M_S = |\mathcal{M}_{\mathcal{S}}|\) the number of intermediate student layers selected for supervision, \(M_T = |\mathcal{M}_{\mathcal{T}}|\) the number of teacher layers used for aggregation in TRAG, and \(D_\mathbf{v}\) the dimensionality of the vision logits. For CompoDistill, let \(G_l\) denote the teacher layer group paired with student layer \(l \in \mathcal{M}_{\mathcal{S}}\). All referenced methods share a foundational base cost comprising the standard language modeling loss and response-level logit distillation.
First, we evaluate the uni-modal distillation overhead. CompoDistill [12] relies on intra-modal visual attention distillation (\(\boldsymbol{A}_{V \rightarrow V}\)). This requires forming a teacher-side reference attention map for each supervised student layer by aggregating the corresponding teacher group \(G_l\), followed by cosine distance computation between two \(N \times N\) attention matrices. Accordingly, the resulting uni-modal complexity is \[\mathcal{O}\!\left(\left(M_S + \sum_{l \in \mathcal{M}_{\mathcal{S}}} |G_l|\right) \cdot N^2\right),\] where the term \(M_S\) corresponds to the repeated cosine evaluations on the student side, and the term \(\sum_{l \in \mathcal{M}_{\mathcal{S}}} |G_l|\) accounts for teacher-group reference formation.
In contrast, TRAG adopts the uni-modal supervision strategy of LLaVA-KD [15]. Crucially, the logit matching \(\mathcal{O}(N \cdot D_\mathbf{v})\) and visual affinity matching \(\mathcal{O}(N^2)\) are computed only once at the final output representation layer. Thus, the total uni-modal complexity for TRAG is \(\mathcal{O}(N^2 + N \cdot D_\mathbf{v})\), effectively eliminating the depth-wise multiplier from the uni-modal distance computation.
Second, we analyze the cross-modal distillation overhead. CompoDistill applies attention supervision from prompt queries to visual keys (\(\boldsymbol{A}_{P \rightarrow V}\)). For each supervised student layer, it first forms a teacher-side prompt-to-vision reference by aggregating the corresponding teacher group \(G_l\) and then computes a cosine distance loss over an \(L_P \times N\) submatrix. Accordingly, the cross-modal complexity becomes \[\mathcal{O}\!\left(\left(M_S + \sum_{l \in \mathcal{M}_{\mathcal{S}}} |G_l|\right)\cdot L_P \cdot N\right).\]
In contrast, TRAG structurally decouples layer aggregation from the distance computation for its Response-to-Vision (\(\boldsymbol{A}_{R \rightarrow V}\)) objective. First, TRAG aggregates the student and teacher attention distributions across their respective target layer sets via simple element-wise additions, incurring \[\mathcal{O}\!\left((M_S + M_T)\cdot L_R \cdot N\right).\]
Then, an Adaptive Kullback-Leibler (KL) objective is applied at the token level, where the teacher attention entropy determines the weighting between the forward and reverse KL terms for each response token. Crucially, this per-token evaluation is performed exactly once on the single depth-aggregated representation. Therefore, despite the granular token-level formulation, the actual adaptive distance computation complexity of TRAG is \(\mathcal{O}(L_R \cdot N)\) after aggregation.
In summary, CompoDistill scales both uni-modal and cross-modal distance computations with the supervised depth through repeated student-side matching and teacher-group reference formation, resulting in \(\mathcal{O}((M_S + \sum_{l \in \mathcal{M}_{\mathcal{S}}} |G_l|)\cdot N^2)\) for uni-modal distillation and \(\mathcal{O}((M_S + \sum_{l \in \mathcal{M}_{\mathcal{S}}} |G_l|)\cdot L_P \cdot N)\) for cross-modal distillation. In contrast, TRAG restricts the heavy uni-modal distance computation to a single output-level evaluation and confines the depth-dependent part of the cross-modal branch to simple aggregation operations. The remaining Adaptive KL evaluation is performed only once on the aggregated Response-to-Vision representation. Table 7 clearly illustrates how TRAG removes the depth multiplier from the heavy divergence evaluation itself and limits depth-dependent computation to aggregation.
| Method | Uni-modal Distillation | Cross-modal Distillation | ||
|---|---|---|---|---|
| 2-5 | Complexity | Depth | Complexity | Objective Granularity |
| LLaVA-KD | \(\mathcal{O}(N^2 + N \cdot D_\mathbf{v})\) | \(1\) (Output) | - | - |
| CompoDistill | \(\mathcal{O}((M_S + \sum_l |G_l|)\cdot N^2)\) | Layer-wise | \(\mathcal{O}((M_S + \sum_l |G_l|)\cdot L_P \cdot N)\) | Block-level Cosine |
| \(\mathcal{O}(N^2 + N \cdot D_\mathbf{v})\) | \(1\) (Output) | \(\mathcal{O}((M_S + M_T)\cdot L_R \cdot N)\) | Aggregation | |
| \(+\mathcal{O}(L_R \cdot N)\) | + Per-token Adaptive KL | |||
0pt
This section provides additional examples (Figures 8, [fig:obs295bird] and [fig:obs295baseball]) to support the attention granularity analysis discussed in the main paper. To systematically examine this phenomenon, we focus on relative depth, a \([0, 1]\) normalized depth coordinate representing the sequential position of decoder layers. Following the observation in the main paper that cross-modal interactions are most pronounced at intermediate depths, we analyze the relative depth range of \([0.3, 0.6]\). The yellow bounding boxes in the figures highlight this intermediate region, where Response-to-Vision attention is most prominent.
Within this intermediate region, cross-modal attention patterns exhibit significant redundancy across adjacent decoder layers. However, when examined at the token level, these distributions reveal highly dynamic and granular grounding behaviors that strictly depend on the semantic role of the generated token.
The following examples demonstrate that response tokens require varying degrees of spatial focus. While some tokens necessitate sharp, highly localized attention on specific anatomical features or objects, others demand a diffuse, scene-level understanding. These visualizations reinforce the necessity of token-level adaptive supervision, as aggregating attention across the entire response sequence would inevitably blur these distinct spatial signals.
To examine whether TRAG generalizes beyond the Qwen backbone family, we additionally evaluate it on MobileLLaMA-based MLLMs. As shown in Table 8, TRAG substantially improves the MobileLLaMA-1.4B student over its PT-SFT baseline, with gains of \(+3.0\) in VQA Avg and \(+7.3\) in CR Avg. This confirms that the proposed response-time attention guidance is not restricted to a specific LLM backbone.
| Setting | VQA \(Avg\) | CR \(Avg\) |
|---|---|---|
| MobileLLaMA-2.7B (PT-SFT) | 59.7 | 57.5 |
| MobileLLaMA-1.4B (PT-SFT) | 55.2 | 50.0 |
| 2.7B \(\rightarrow\) 1.4B | 58.2 | 57.3 |
6pt
Table 9 studies the effect of the relative depth range used for attention aggregation. The shifted intermediate range \([0.4, 0.7]\) achieves performance comparable to the default \([0.3, 0.6]\) range, whereas early and late ranges degrade both VQA and compositional reasoning performance. These results suggest that the benefit of TRAG comes from supervising intermediate cross-modal grounding behavior rather than from a model-specific layer choice.
| Relative Depth Range | VQA \(Avg\) | CR \(Avg\) |
|---|---|---|
| Early \([0.0, 0.3]\) | 58.7 | 69.2 |
| Late \([0.6, 1.0]\) | 59.2 | 73.2 |
| Intermediate \([0.4, 0.7]\) (Shifted) | 60.3 | 75.0 |
| Intermediate \([0.3, 0.6]\) (Default) | 60.2 | 75.2 |
6pt
Table 10 studies the granularity of response-token supervision. Aggregating attention targets over the entire response sequence yields lower performance than applying per-token response supervision. This indicates that token-wise supervision better preserves the step-specific grounding signals required during response generation.
| Response Token Usage | VQA \(Avg\) | CR \(Avg\) |
|---|---|---|
| Aggregation | 57.7 | 70.4 |
| Token-Level (Ours) | 60.2 | 75.2 |
6pt
This section provides a comprehensive breakdown of the objective function optimized during the distillation stages of TRAG. While recent work has explored perception-decomposed confidence rewards for vision-language reasoning [45], TRAG uses teacher-guided distillation to supervise response-time visual grounding. The overall training loss is formulated as the sum of three distinct components:
\[\mathcal{L}_{\text{TRAG}} = \mathcal{L}_{\text{LM}} + \mathcal{L}_{\text{uni\text{-}modal}} + \mathcal{L}^{(\text{TRAG})}_{\text{cross\text{-}modal}}. \label{eq:appendix95total95loss}\tag{18}\]
We apply the standard autoregressive cross-entropy loss over the target response sequence to maintain the foundational generation capabilities of the student model:
\[\mathcal{L}_{\text{LM}} = - \sum_{t=1}^{L_R} \log p^{\mathcal{S}} \!\left( y_t \mid \mathbf{v}, \mathbf{x}, \mathbf{y}_{<t} \right), \label{eq:appendix95lm95loss}\tag{19}\] where \(\mathbf{v}\), \(\mathbf{x}\), and \(\mathbf{y}\) denote the visual tokens, prompt tokens, and target response tokens, respectively.
The uni-modal distillation term aligns the teacher and student models within their respective modalities. It combines a response-level logit matching term with an auxiliary vision-modality loss, denoted as \(\mathcal{L}_{\text{aux}}\):
\[\mathcal{L}_{\text{uni-modal}} = \sum_{t=1}^{L_R} D_{KL} \left( p^{\mathcal{T}}(\cdot \mid \mathbf{v}, \mathbf{x}, y_{<t}) \parallel p^{\mathcal{S}}(\cdot \mid \mathbf{v}, \mathbf{x}, y_{<t}) \right) + \mathcal{L}_{\text{aux}}. \label{eq:appendix95unimodal95main}\tag{20}\] The first term minimizes the Kullback-Leibler (KL) divergence between the teacher and student next-token probability distributions at each decoding step.
To ensure stable training and facilitate a fair comparison, the auxiliary loss \(\mathcal{L}_{\text{aux}}\) strictly follows the configuration established in LLaVA-KD [15]. It comprises vision-token logit matching and visual affinity matching: \[\mathcal{L}_{\text{aux}} = \lambda_\alpha \sum_{i=1}^{N} D_{KL}\!\left( p^\mathcal{T}_{\mathrm{vision},i} \,\|\, p^\mathcal{S}_{\mathrm{vision},i} \right) + \lambda_{\beta} \mathcal{D}_{\cos} \!\left( \mathbf{R}^{\mathcal{T}}_{\text{vision}}, \mathbf{R}^{\mathcal{S}}_{\text{vision}} \right). \label{eq:appendix95laux}\tag{21}\] Here, \(p^\mathcal{T}_{\mathrm{vision},i}\) and \(p^\mathcal{S}_{\mathrm{vision},i}\) denote the teacher and student output distributions for the \(i\)-th visual token, respectively, while \(\mathbf{R}^{\mathcal{T}}_{\text{vision}}\) and \(\mathbf{R}^{\mathcal{S}}_{\text{vision}}\) represent their corresponding visual affinity matrices. \(\mathcal{D}_{\cos}\) denotes the cosine distance. We retain the original hyperparameter weighting from LLaVA-KD, setting \(\lambda_{\alpha} = 1.0\) and \(\lambda_{\beta} = 0.5\).
The core contribution of our distillation framework is encapsulated in the cross-modal attention alignment term. As formulated in the main paper, the TRAG objective dynamically balances the forward and reverse KL divergences based on the token-specific entropy of the teacher’s attention distribution: \[\mathcal{L}^{(\text{TRAG})}_{\text{cross-modal}} = \sum_{i \in \mathcal{I}_R} \left[ \lambda_i D_{KL} \!\left( \tilde{\mathbf{a}}^{\mathcal{T}}_i \,\|\, \tilde{\mathbf{a}}^{\mathcal{S}}_i \right) + (1-\lambda_i) D_{KL} \!\left( \tilde{\mathbf{a}}^{\mathcal{S}}_i \,\|\, \tilde{\mathbf{a}}^{\mathcal{T}}_i \right) \right]. \label{eq:appendix95crossmodal}\tag{22}\]
While TRAG effectively enhances cross-modal grounding, we identify two primary limitations in the current framework.
First, the proposed distillation objective inherently assumes vocabulary consistency between the teacher and the student. Because the token-wise attention alignment and the response-level logit supervision are defined over aligned target response positions under teacher forcing, the current framework cannot be directly applied to teacher-student pairs with mismatched tokenizers or entirely disjoint vocabularies.
Second, although our method focuses on the response phase, the cross-modal supervision is computed under a standard teacher forcing paradigm using fixed ground-truth target sequences. Consequently, the distilled attention patterns are not explicitly optimized to dynamically adapt to the student’s own autoregressive generation trajectory, which may introduce a discrepancy between training-time supervision and inference-time behavior.
Despite this theoretical limitation, the main paper shows that TRAG maintains high attention fidelity during autoregressive generation without teacher forcing, reaching approximately 0.72 cosine similarity at intermediate layers. While such high fidelity demonstrates that our approach meaningfully mitigates the train-inference mismatch, an inherent structural discrepancy inevitably persists.
Extending TRAG to an on-policy distillation setting would therefore be a natural direction for future work [13]. By defining supervision along the student’s actual decoding trajectory and dynamically querying the teacher’s attention for student-generated tokens, such a formulation could provide a more faithful approximation of the cross-modal grounding behavior required at inference time.
To ensure the full reproducibility of our proposed framework, TRAG, we provide comprehensive methodological details throughout the main paper. Furthermore, exhaustive training configurations and hyperparameters are specified in Appendix 8, and the exact mathematical formulations of all objective functions are detailed in Appendix 14.
In this section, we present a qualitative comparison of inference samples generated by our student model, TRAG (Qwen1.5-0.5B), the teacher model (Qwen1.5-4B), and competitive baselines (CompoDistill and LLaVA-KD, both Qwen1.5-0.5B). Figures 9–12 illustrate that TRAG more closely matches the teacher’s response behavior and fine-grained visual grounding at inference time. These examples support the effectiveness of our response-phase attention distillation in transferring the teacher’s response characteristics, allowing the student to move beyond generalized descriptions and capture salient localized features that are often overlooked by the baselines.