CompoDistill: Attention Distillation for Compositional Reasoning in Multimodal LLMs

1. Problem Statement

Multimodal Large Language Models (MLLMs) have achieved remarkable success but suffer from high computational costs, necessitating the development of efficient smaller models via Knowledge Distillation (KD). While existing KD methods successfully transfer visual recognition abilities (e.g., object identification) to student models, they fail to effectively distill visual perception abilities.

• Visual Recognition vs. Visual Perception: Recognition involves identifying objects and basic attributes. Perception involves understanding complex relationships, spatial arrangements, and contextual interactions between objects (essential for Compositional Reasoning or CR).
• The Gap: Empirical analysis reveals that while KD improves performance on Visual Question Answering (VQA) benchmarks (recognition), it yields no significant improvement over standard Supervised Fine-Tuning (SFT) on CR benchmarks.
• Root Cause: The authors identify Visual Attention Misalignment as the primary culprit. Existing KD methods do not force the student model to attend to the same relevant image regions as the teacher model when processing text queries, leading to a failure in learning complex relational reasoning.

2. Methodology: CompoDistill

The authors propose CompoDistill, a novel KD framework designed to explicitly align the student's visual attention with the teacher's. The framework consists of two core modules and a three-stage training strategy.

A. Core Modules

1. Visual ATtention Alignment (VAT) Module:

   • Goal: Explicitly align the attention distributions of the student with the teacher over visual tokens within the "visual understanding layers" (intermediate layers of the LLM).
   • Mechanism: It minimizes the cosine distance between the student's attention matrix and the teacher's attention matrix.
   • Group Layer Matching: Since teachers and students often have different numbers of layers, the authors propose a Group Matching strategy. Instead of a naive one-to-one mapping, each student layer is matched with a group of consecutive teacher layers (one-to-many). This allows the student to capture a broader, more robust distribution of the teacher's attention knowledge.

2. Teacher Adapter Fetch (TAF) Module:

   • Goal: Resolve the feature space mismatch between the teacher's frozen adapter and the student's trainable adapter.
   • Problem: The teacher's attention mechanism is optimized for its own visual feature space. Directly applying this to a student with a different adapter creates a conflict.
   • Solution: The TAF module reuses the teacher's frozen, pre-trained adapter to project visual features into the language space. A lightweight, trainable MLP is then added only to align the dimensions with the student's hidden size. This ensures the student processes visual inputs through the "same lens" as the teacher, making the attention transfer effective.

B. Three-Stage Training Strategy

1. Distilled Pre-Training (DPT): Aligns the visual feature space with the language space. The student adapter is initialized via TAF (using the frozen teacher adapter) and optimized using Language Modeling (LLM) and KL-divergence losses.
2. Distilled Fine-Tuning (DFT): The core distillation phase. Both the student LLM and adapter are fine-tuned. The objective includes LLM loss, KL-divergence, and the Attention Distillation Loss (ADL) via the VAT module.
3. Supervised Fine-Tuning (SFT): A final stage using only the autoregressive language modeling loss to consolidate knowledge and strengthen instruction-following capabilities.

3. Key Contributions

1. Identification of the Bottleneck: The paper is the first to systematically identify and analyze that existing KD methods fail to distill visual perception due to visual attention misalignment between teacher and student models.
2. Novel Framework (CompoDistill): Introduces a framework that combines VAT (for attention alignment) and TAF (for feature space compatibility) to effectively transfer perception abilities.
3. Group Layer Matching: Proposes a simple yet effective strategy to handle architectural depth differences between teacher and student, outperforming simple uniform sampling or adaptive matching.
4. Comprehensive Analysis: Provides extensive ablation studies proving that intermediate layers are the critical "visual understanding" layers and that Cosine Similarity is the optimal loss function for attention distillation.

4. Experimental Results

The authors evaluated CompoDistill on various benchmarks using a 2B student model distilled from a 4B teacher (LLaVA-4B).

• Compositional Reasoning (CR): CompoDistill significantly outperformed existing KD methods (LLaVA-KD, LLaVA-MoD) and standard SFT models. It achieved performance levels competitive with 4B models, a feat previously unattained by 2B KD models.
  • Example: On SugarCrepe, CompoDistill achieved 82.9% accuracy vs. 75.3% for LLaVA-KD and 73.3% for SFT.
• Visual Question Answering (VQA): The method maintained strong performance on VQA tasks (e.g., GQA, TextVQA, MME), preserving visual recognition capabilities while enhancing perception.
• Relational Hallucination: The method significantly reduced relational hallucinations on R-Bench and Reefknot benchmarks, bringing the student's performance nearly on par with the teacher.
• Efficiency & Generalizability:
  • Achieved these results with only 1.2M training samples, demonstrating high data efficiency compared to methods requiring millions more.
  • Validated on different backbones (MobileLLaMA) and with larger teachers (7B), proving the approach is robust and generalizable.

5. Significance

• Paradigm Shift: Moves the focus of MLLM distillation from simple logit matching or recognition tasks to perception and reasoning, addressing a critical gap in current efficient MLLM research.
• Practical Impact: Enables the deployment of highly capable, small-scale MLLMs that can perform complex reasoning tasks (e.g., understanding object relationships) without the computational burden of massive models.
• Theoretical Insight: Establishes that attention alignment is the key mechanism for transferring visual perception, providing a new direction for future distillation research in multimodal domains.

In conclusion, CompoDistill demonstrates that by explicitly correcting visual attention misalignment and bridging feature space gaps, smaller models can inherit the complex reasoning capabilities of larger teachers, making efficient, high-performance multimodal AI more accessible.