Mind the Heads: Topological Representation Alignment for Multimodal LLMs

Abstract

Key Contributions

Method

HeRA augments the standard language-modeling objective with a contrastive head-level alignment loss that pushes the multimodal representation of an image–caption pair \((I_i, x_i)\), computed at a small set of selected attention heads, toward its \(k\)-nearest neighbors in the latent space of a frozen vision teacher.

Figure 2. Overview of HeRA. Alongside the standard language modeling objective \(\mathcal{L}_{\text{LM}}\), HeRA employs a contrastive loss \(\mathcal{L}_{\text{HeRA}}\) to pull representations from selected LLM attention heads closer to their \(k\)-nearest neighbors (Top-\(k\)), computed in the latent space of a frozen teacher vision encoder.

Figure 3. The animation demonstrates the core mechanism of HeRA. Multimodal inputs from a batch are processed, and the visual data is passed through a frozen teacher vision encoder. This defines the local neighborhood structure in the teacher's latent space, serving as the target for HeRA's contrastive loss to align specific LLM attention heads.

Results

Across 9 LLMs · Vision-Centric Average

HeRA delivers consistent improvements on the demanding Vision category of the Cambrian benchmark suite, scaling across architectural families (Vicuna, LLama3, Qwen2.5, Qwen3) and model sizes (3B–14B), while preserving or improving General / Knowledge / OCR performance.

Comparison with Representation Alignment Methods

On Qwen3-8B, HeRA outperforms recent representation alignment strategies — including feature reconstruction (ROSS, JARVIS), middle-layer alignment (VIRAL), and global CKA-based alignment (CMAR) — with the highest gains on the demanding Vision category.

Figure 4. Comparison of MKNN scores of the Worst-5 and Top-5 heads after the second training stage performed with HeRA and our competitors (i.e. LLaVA, VIRAL, JARVIS, ROSS, CMAR).

Less Visual Hallucinations

Although mitigating hallucinations is not an explicit objective, HeRA consistently reduces hallucination rates on CHAIR-MSCOCO, AMBER (generative & discriminative), and HallusionBench, while preserving cognition/coverage scores. Aligning topology curbs over-reliance on linguistic priors.

Choice of the Vision Teacher

Self-supervised teachers work best: DINOv2 / DINOv3 deliver strong, consistent gains, while using SigLIP2 (the primary vision encoder) as teacher is mostly ineffective. Base-size DINOv2-B is already sufficient — scaling the teacher to 1B parameters gives no further benefit.

Figure 5. VQA results of HeRA with different teacher vision encoders.

Qualitative Examples

HeRA visibly improves visual grounding on fine-grained perception, counting, spatial reasoning, and dense scenes — reducing the model's tendency to fall back on linguistic priors.

Figure 6. Qualitative comparison of LLaVA, ROSS, and HeRA across Cambrian categories using Qwen3-8B and SigLIP2.

Figure 7. Failure cases of HeRA on VQA tasks.

Citation

@article{caffagni2026hera,
    title   = {{Mind the Heads: Topological Representation Alignment for Multimodal LLMs}},
    author  = {Caffagni, Davide and Compagnoni, Alberto and Melis, Federico and
                Sarto, Sara and Dovesi, Pier Luigi and Granroth-Wilding, Mark and
                Cornia, Marcella and Baraldi, Lorenzo},
    journal = {arXiv preprint},
    year    = {2026}
    }

Acknowledgments

This work has been supported by the EU Horizon project ELLIOT (No. 101214398), by the EuroHPC JU project MINERVA (GA No. 101182737), and by the PNRR project ITSERR (CUP B53C22001770006) funded by the EU — NextGenerationEU. We also acknowledge EuroHPC JU for awarding the project EHPC-AIF-2025SC04-225 access to LUMI at CSC, Finland.