💻 computer science

Mind the Heads: Topological Representation Alignment for Multimodal LLMs

This paper proposes HeRA, a novel method that improves Multimodal Large Language Models by enforcing topological representation alignment at the individual attention head level, revealing that targeting the least aligned heads yields significant performance gains and reduces visual hallucinations.

Original authors: Davide Caffagni, Alberto Compagnoni, Federico Melis, Sara Sarto, Pier Luigi Dovesi, Mark Granroth-Wilding, Marcella Cornia, Lorenzo Baraldi

Published 2026-06-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Davide Caffagni, Alberto Compagnoni, Federico Melis, Sara Sarto, Pier Luigi Dovesi, Mark Granroth-Wilding, Marcella Cornia, Lorenzo Baraldi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Robot to "See" with its Words

Imagine you have a brilliant robot that is a master storyteller but a terrible observer. It can write beautiful poems and answer complex history questions, but if you show it a picture of a cat sitting on a mat, it might confidently tell you the cat is flying because it has read so many stories about flying cats. This is a common problem with Multimodal Large Language Models (MLLMs): they rely too much on what they think they know (language) and not enough on what they actually see (vision).

To fix this, researchers usually try to "teach" the robot's brain to match the way a dedicated "vision teacher" (a specialized camera brain) sees the world. However, the old way of doing this was like trying to tune a whole orchestra by forcing every single instrument to play the exact same note at the exact same time. It's too rigid and often breaks the music (the robot's ability to speak and reason).

The New Solution: HeRA (Head-Wise Representation Alignment)

The authors of this paper propose a smarter, more surgical approach called HeRA. Instead of tuning the whole brain, they tune specific "musicians" inside the robot's brain.

1. The Brain is a Choir of Specialized Singers

Inside the robot's language brain (the LLM), there isn't just one big processor. There are many layers, and inside each layer, there are dozens of "attention heads." Think of these heads as individual singers in a choir.

Some singers are great at grammar.
Some are great at remembering facts.
Some are great at visual descriptions.
Some are just... bad at describing what they see.

2. The "Platonic" Idea: Keeping the Neighborhood

The paper is based on a theory called the Platonic Representation Hypothesis. Imagine the robot's brain is a map of a city.

In a good map, places that are similar (like two different types of dogs) should be located near each other.
In a bad map, a dog might be next to a toaster just because the robot got confused.

The goal is to make sure the robot's "visual map" matches the "language map" so that similar things stay neighbors. The researchers use a metric called MKNN (Mutual K-Nearest Neighbor) to check if the neighbors are staying in the right spots.

3. The Surprising Twist: Fix the Worst Singers

Here is the counter-intuitive magic of HeRA.

Old Method: Try to align the "best" singers (the ones already good at vision) with the teacher.
HeRA Method: The researchers found that the worst singers (the attention heads with the lowest alignment scores) are actually the ones that need the most help.

The Analogy: Imagine a team of runners. If you train the fastest runner, they might get a little faster, but the team's overall time won't change much. But if you take the slowest runner and give them a personal coach to help them catch up, the whole team improves significantly.

HeRA identifies the 5 "slowest" (least aligned) attention heads in the robot's brain and gives them a special training session to match the "vision teacher." It leaves the other heads alone so they don't forget how to speak or reason.

How It Works (The Training)

The Teacher: A frozen, super-smart vision encoder (like a DINOv2 camera brain) looks at an image and says, "This is a dog next to a ball."
The Student: The robot looks at the same image and text.
The Fix: The system checks which specific "singers" (attention heads) in the robot are getting the neighborhood wrong. It then uses a special "contrastive loss" (a math penalty) to gently push those specific singers to rearrange their internal map so that "dog" and "ball" stay close together, just like the teacher's map.

The Results: Better Vision, No Loss of Smarts

The paper tested this on many different robot models (from small 3-billion-parameter ones to massive 14-billion ones) and 18 different tests.

Vision Tasks: The robots got much better at hard visual tasks, like counting objects, understanding spatial relationships (is the cup on the table or under it?), and spotting specific details.
No "Brain Damage": Crucially, because they only tweaked the specific "vision-challenged" heads, the robots didn't lose their ability to speak, reason, or answer general knowledge questions. In fact, they often got better at those too.
Less Hallucination: The robots stopped making up things that weren't there. Because they were forced to respect the visual "neighborhood," they couldn't just guess based on what they read in books.

Summary

The paper introduces HeRA, a method that fixes multimodal AI by surgically training only the specific parts of the brain that are bad at seeing. Instead of forcing the whole brain to change, it identifies the "struggling" attention heads and helps them align with a vision teacher. This results in robots that see the world more accurately, make fewer mistakes, and don't lose their ability to talk and think.

Technical Summary: Head-Wise Representation Alignment (HeRA) for Multimodal LLMs

Problem Statement

Multimodal Large Language Models (MLLMs) have demonstrated significant capabilities in vision-language tasks, yet they continue to struggle with foundational visual reasoning challenges, such as object counting, spatial relationship understanding, and dense visual parsing. Current improvements in MLLMs are largely driven by scaling data and model size rather than principled interventions into internal mechanisms.

A primary limitation of existing representation alignment methods—which regularize MLLM internal representations to match those of an external vision encoder—is their reliance on aligning a fixed layer of the language backbone (e.g., the middle layer). This approach overlooks the fine-grained internal structure of Transformer models. Furthermore, rigid feature matching in MLLMs, which possess strong pre-trained language priors, can conflict with the language modeling objective, potentially leading to suboptimal performance or catastrophic interference.

Methodology: Head-Wise Representation Alignment (HeRA)

The authors propose HeRA, a method that enforces cross-modal alignment at the level of individual attention heads rather than coarse, fixed layers. The approach is grounded in the Platonic Representation Hypothesis (PRH), which posits that representations across different modalities share consistent local neighborhood structures (topology) even if their global geometries differ.

Core Components

Topological Alignment via MKNN:
The method utilizes the Mutual K-Nearest Neighbor (MKNN) metric to measure the agreement between the local neighborhood structures of textual and visual representations. High MKNN scores indicate that semantically similar inputs share similar neighbors in both the language and vision latent spaces.
Differentiable Contrastive Proxy:
Since the MKNN metric is non-differentiable (relying on discrete neighbor indices), HeRA introduces a contrastive objective as a differentiable proxy. For a given multimodal input, the model is trained to pull its representation closer to the representations of the $k$ -nearest neighbors (computed in a frozen teacher vision encoder's space) while pushing it away from other samples in the batch. This is implemented using a multi-target variant of the InfoNCE loss.
Head-Level Granularity:
Instead of aligning the entire pooled output of a layer, HeRA isolates the projected contribution of individual attention heads. For a multi-head attention layer, the output is decomposed into contributions from each head ( $h$ ) in layer ( $l$ ). The contrastive loss is applied independently to these specific head-level representations.
Strategic Head Selection:
A critical design choice is determining which heads to align. The authors compute MKNN scores for all heads using the language model before multimodal training (purely textual representations). Contrary to the intuition that one should align the "best" heads, experiments reveal that targeting the least aligned heads (the "Worst-5") yields the largest performance gains. This strategy strengthens misaligned components without disrupting the already well-aligned structures or the core language modeling task.

Key Contributions

Fine-Grained Alignment: Moving beyond fixed-layer alignment to operate at the individual attention head level, allowing for more atomic interventions in the model.
Topological Focus: Shifting from strict feature matching to preserving local topological structures (neighborhood relationships) across modalities, aligning with the PRH.
Counter-Intuitive Selection Strategy: Identifying that aligning the least aligned attention heads provides superior regularization and performance improvements compared to aligning already well-aligned heads or random subsets.
Regularization Against Hallucinations: Demonstrating that topological alignment naturally curbs the model's over-reliance on linguistic priors, thereby reducing visual hallucinations.

Experimental Results

The authors evaluated HeRA across multiple MLLMs (including Vicuna, LLaMA3, Qwen2.5, and Qwen3 families ranging from 3B to 14B parameters) and 18 benchmarks.

Vision-Centric Tasks: HeRA consistently improves performance on challenging vision-centric benchmarks (e.g., RWQA, MMVP, Blink, CVBench). For instance, Qwen3-8B achieved a +3.6 average improvement on vision tasks, and Qwen2.5-14B saw a +3.4 gain.
General Capabilities: Unlike previous methods that sometimes degrade general language performance, HeRA preserves or improves scores on General Knowledge, OCR, and Knowledge tasks, indicating successful mitigation of catastrophic interference.
Hallucination Reduction: Models trained with HeRA show reduced hallucination rates (measured by CHAIR and AMBER benchmarks) while maintaining or improving cognitive scores, suggesting a more faithful integration of visual evidence.
Ablation Findings:
- Aligning the "Worst-5" heads outperforms aligning the "Top-5" or random heads.
- Aligning too many heads (e.g., 10) leads to performance regression, confirming the need for selective alignment.
- Contrastive learning objectives outperform direct feature matching (cosine similarity) at the same granularity.
- Using DINO-based vision encoders as teachers yields better results than using the same encoder used for the primary vision input (SigLIP2).

Significance

The paper claims that HeRA offers a principled approach to enhancing MLLMs by intervening directly in the interaction between vision and language. By targeting specific attention heads and enforcing topological consistency rather than rigid feature matching, HeRA addresses the visual perception deficit in MLLMs without compromising their linguistic priors. The work suggests that the "least aligned" components of a multimodal model are the most critical targets for improvement, providing a new direction for regularization strategies in multimodal learning. The code is publicly released to facilitate further research.