Skip to content

Mind the Heads: Topological Representation Alignment for Multimodal LLMs

🕒 Published (v1): 2026-06-22 00:00 UTC · Source: HuggingFace · link

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

HeRA (Head-Wise Representation Alignment) improves Multimodal LLMs by applying contrastive topological alignment at the level of individual attention heads rather than fixed layers, guided by the Mutual K-Nearest Neighbor (MKNN) metric. Counterintuitively, targeting the least aligned heads yields the largest gains. Evaluated across 9 LLM backbones (3B–14B) on 18 benchmarks, HeRA consistently improves vision-centric reasoning and reduces hallucinations.

Problem

Existing representation alignment methods for MLLMs align a fixed, hard-coded LLM layer (e.g., middle or final) toward an external vision encoder, ignoring the fine-grained specialization of individual attention heads. This coarse granularity can conflict with pre-existing language modeling priors, limiting gains or causing regressions in vision-centric tasks.

Method

HeRA operates within the LLaVA training framework and adds a head-wise contrastive loss alongside the standard language modeling objective.

Head selection: Before multimodal training, the MKNN alignment score (Eq. 3) is computed for every attention head of the base LLM against a frozen teacher vision encoder (DINOv2 ViT-L). The \(m=5\) least aligned heads \(H^{\text{worst}}_m\) are selected.

Head-level representation: Each head \(h\) in layer \(l\) contributes a projected representation \(M_{l,h}(I_i, x_i) = h_{l,h}(I_i, x_i) W_{O,h} \in \mathbb{R}^{d_G}\), exploiting the linear decomposition of the output projection matrix \(W_O\).

Contrastive objective: For each sample \(i\), the \(k\) nearest neighbors of image \(I_i\) in the teacher's visual space \(N_k^{V^t}(I_i)\) act as positive targets. A multi-target InfoNCE loss pulls the head's multimodal representation toward these neighbors:

\[\mathcal{L}_{\text{RA}}(I_i, x_i, M) = -\frac{1}{k}\sum_{j \in N_k^{V^t}(I_i)} \log \frac{\exp\!\left(\frac{M_{l,h}(I_i,x_i)^\top M_{l,h}(I_j,x_j)}{\tau}\right)}{\sum_{z \in B, z\neq i} \exp\!\left(\frac{M_{l,h}(I_i,x_i)^\top M_{l,h}(I_z,x_z)}{\tau}\right)}\]

The final objective is \(\mathcal{L} = \mathcal{L}_{\text{LM}} + \lambda \mathcal{L}_{\text{HeRA}}\) with \(\lambda = 0.01\), \(k = 10\), \(m = 5\).

Key Contributions

  • Head-level alignment granularity: First method to apply cross-modal representation alignment at the individual attention head level, bypassing the fixed-layer assumption.
  • MKNN-guided selection of least-aligned heads: Empirically demonstrates that targeting the bottom-5 MKNN heads outperforms top-5, random, and layer-level selection.
  • Differentiable topological proxy: Multi-target InfoNCE loss as a differentiable surrogate for the discrete MKNN metric, explicitly preserving local neighborhood structure across modalities.
  • Hallucination regularization: HeRA naturally reduces over-reliance on linguistic priors without explicit hallucination-specific training signal.
  • Scale and architecture generality: Validated on Vicuna, LLaMA3, Qwen2.5, Qwen3 from 3B to 14B parameters.

Results

VQA (Cambrian suite, Table 2) — HeRA vs. base LLaVA: - Qwen3-8B: +3.6 overall (55.9 → 59.5), MMVP +8.7, Blink +3.3 - Qwen2.5-14B: +3.4 overall (54.9 → 58.3), General +1.8, OCR +4.5 - Qwen2.5-3B: +2.4 overall (50.5 → 52.9), V* +6.8 - Consistent gains across all 9 backbones on vision-centric tasks; general/knowledge performance preserved or improved

Comparison vs. representation alignment baselines (Table 4, Qwen3-8B backbone): - HeRA: 59.5 overall vs. CMAR 58.7, ROSS 55.9, VIRAL 56.2, JARVIS 54.2 - HeRA +3.6 over base LLaVA; next best (CMAR) achieves +2.8

Hallucination (Table 3): - MSCOCO CHAIRi reduced on most backbones (e.g., LLaMA3-8B: 13.0 → 11.7, Qwen2.5-14B: 12.2 → 10.6) - AMBER HalRate reduced (Qwen2.5-3B: 30.1 → 27.6; Qwen3-8B: 28.9 → 27.9) - HallusionBench Hard accuracy: Qwen2.5-3B +5.0 (55.2 → 60.2), Qwen3-14B +3.8

Limitations

  • Evaluated exclusively within the LLaVA training framework; applicability to other MLLM pipelines (e.g., natively multimodal models, RLHF-tuned models) not demonstrated.
  • Head selection requires a pre-training MKNN pass over ~1,000 captioned samples; adds a setup step before multimodal training.
  • Ablations show performance degrades with \(m=10\) heads, but the mechanism for this conflict with the LM objective is not fully explained.
  • Hallucination results are mixed for some models (e.g., Qwen3-4B shows slight CHAIR regression; Vicuna-13B CHAIRs worsens), indicating inconsistent anti-hallucination benefit.
  • Only DINOv2 ViT-L is used as the teacher vision encoder; sensitivity to teacher choice is not explored.

Relevance to Vision-Language Models

HeRA directly addresses a core structural weakness in MLLM training: the mismatch between vision encoder representations and LLM internal representations, particularly at the attention-head level. By grounding head selection in the Platonic Representation Hypothesis and MKNN, it provides a principled diagnostic for understanding which model components are vision-misaligned — a tool broadly applicable to studying internal representations of any MLLM. The finding that least-aligned heads are the highest-leverage intervention points challenges the common heuristic of aligning middle layers, and has implications for future work on cross-modal representation learning and architectural analysis of VLMs. The method also operationalizes topological alignment (local neighborhood preservation) as a training objective, offering a cleaner alternative to feature-matching losses that conflict with language priors.