Skip to content

CompoDistill: Attention Distillation for Compositional Reasoning in Multimodal LLMs

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; attention distillation from large to small MLLMs for compositional reasoning

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

CompoDistill is a knowledge distillation (KD) framework for Multimodal LLMs that identifies and corrects visual attention misalignment — the failure of student models to attend to the same image regions as their teacher — as the root cause of poor compositional reasoning (CR) transfer. It introduces two modules (VAT and TAF) with a three-stage training strategy that significantly improves CR performance while preserving VQA gains.

Problem

Existing KD methods for MLLMs successfully transfer visual recognition ability (VQA benchmarks) but leave visual perception ability — understanding spatial relationships, attributes, and compositional structure — essentially undistilled: student KD models match only SFT baseline on CR benchmarks like SugarCrepe and Winoground. The paper identifies the root cause as visual attention misalignment: in intermediate ("visual understanding") LLM layers (30–70% depth), the teacher–student cosine similarity of attention over visual tokens does not improve over the teacher–SFT baseline for CR tasks, even though it does for VQA tasks.

Method

Visual ATtention alignment (VAT) module. For each visual understanding layer \(l_s\) of the student, the attention sub-matrix over visual tokens is extracted as \(\tilde{A}^s_{l_s} = A_{l_s}[:, :N_v] \in \mathbb{R}^{(N_v+N_t)\times N_v}\). Rather than one-to-one layer matching, a Group Layer Matching strategy maps each student layer \(l_{s_j}\) to a sliding window of \(n = m - k + 1\) consecutive teacher layers \(G_j\). The attention distillation loss minimizes: $\(\mathcal{L}_\text{ADL} = \frac{1}{k}\sum_{j=1}^k \left(1 - \text{sim}\!\left(\bar{A}^t_j,\, \tilde{A}^s_{l_{s_j}}\right)\right), \quad \bar{A}^t_j = \frac{1}{n}\sum_{l \in G_j} \tilde{A}^t_l\)$

Teacher Adapter Fetch (TAF) module. Because the teacher's attention is tuned to its own vision–language feature space, directly imposing it on the student causes a feature space conflict. TAF reuses the teacher's frozen pretrained adapter \(P^t_{\psi_t}\) and prepends a lightweight trainable MLP \(P^s_{\psi_s}\) for dimensional alignment: $\(x_v^s = P^s_{\psi_s}\!\left(P^t_{\psi_t}(z_p)\right) \in \mathbb{R}^{N_v \times d_s}\)$

Three-stage training. Stage 1 (DPT): freeze vision encoder + LLM, train adapter with \(\mathcal{L}_\text{LM} + \mathcal{L}_\text{KL}\) to align visual–language space via TAF. Stage 2 (DFT): unfreeze student LLM and adapter, add \(\mathcal{L}_\text{ADL}\) to impose attention alignment via VAT. Stage 3 (SFT): fine-tune student only with \(\mathcal{L}_\text{LM}\) to consolidate transferred knowledge.

Key Contributions

  • First systematic demonstration that existing MLLM KD methods fail to distill visual perception (CR) due to visual attention misalignment in intermediate LLM layers, with empirical proof via layer-wise attention cosine similarity analysis.
  • VAT module with Group Layer Matching (one-to-many, sliding window) that outperforms both naive uniform sampling and adaptive matching baselines.
  • TAF module that bridges the teacher–student feature space gap, making attention distillation tractable for MLLMs.
  • Controlled intervention experiment (substituting student attention with averaged teacher attention at inference) confirming that attention alignment causally improves CR performance.
  • Demonstrated generalizability to a more advanced MLLM backbone.

Results

  • CompoDistill significantly outperforms state-of-the-art KD methods (MiniKD/Shu et al. 2024, LLaVA-KD/Cai et al. 2024, Xu et al. 2024) on CR benchmarks (SugarCrepe, SADE, BiVLC, Winoground) evaluated on 2B student models distilled from a 4B teacher.
  • Existing KD methods show CR performance on par with the SFT baseline despite outperforming it on VQA; CompoDistill breaks this pattern.
  • Attention substitution experiment on SugarCrepe (Swap, Replace, Add subtasks) shows consistent modest gains, confirming the causal link.
  • Positive correlation between teacher–student visual attention similarity and ground-truth answer probability demonstrated on 5,000 GQA test samples.
  • Competitive or improved VQA performance on GQA, TextVQA, VQAv2, VizWiz, and MME relative to KD baselines.

(Full benchmark tables are beyond the truncated text; specific accuracy numbers are not reported here.)

Limitations

  • Completely replacing student attention with teacher attention causes slight performance degradation (feature space mismatch), motivating TAF — but TAF introduces additional parameters and training complexity.
  • The group matching window size \(n\) is set in closed form (\(n = m - k + 1\)), which may not be optimal across all architecture depth ratios.
  • Analysis is conducted within the LLaVA architecture family; generalization to architectures with different adapter designs or tokenization strategies is not fully established.
  • Intervention experiments show only "modest" CR gains from attention substitution, suggesting attention alignment is necessary but perhaps not sufficient for full perception transfer.

Relevance to Vision-Language Models

This paper directly addresses a critical but previously overlooked failure mode in MLLM compression: KD methods optimize for token-level output mimicry and achieve VQA gains, but the richer spatial and relational reasoning encoded in teacher attention maps is not transferred. By formalizing the distinction between visual recognition and visual perception and grounding it in layer-wise attention analysis, CompoDistill provides both a diagnostic tool and a remedy applicable to any teacher–student MLLM pair. The finding that intermediate LLM layers are the locus of this gap is consistent with recent mechanistic interpretability work on MLLMs and has direct implications for future efficient-VLM design, particularly for tasks requiring fine-grained grounding, referring expression comprehension, and relation understanding.