A Broader Look into MLLM's Retrieval Heads

Locating relevant evidence in long inputs is a core capability of Multimodal Large Language Models. Prior studies attribute this capability to retrieval heads, a sparse set of attention heads that transfer information from the context to the output. However, these analyses are confined to a limited set of modalities and the corresponding models. Moreover, how retrieval heads can be leveraged to enhance the model calls for further exploration. In this paper, we ask whether this mechanism expands together with the input modalities of the model. Taking audio-visual input as a newly added modality, we construct paired inputs in which the same answer is either spoken or shown on screen. With the evidence content fixed, retrieval heads remain sparse in the new modality and are shared to a large extent across modalities. Across 18 language, vision-language, and omni-modal checkpoints, omni-modal models largely preserve the retrieval heads of their backbones. Ablating these heads eliminates nearly all audio-visual retrieval. The heads shared across modalities make up fewer than 2\% of all heads, yet ablating them also degrades document understanding, spatial perception, and visual mathematics. Building on these findings, we turn retrieval heads into a target of optimization by training their attention with a ranking objective. This raises the page Recall@1 of Qwen3-VL-8B on MMDocIR by 6.03 points and improves document retrieval across three model families. Functional attention heads can thus serve both as a lens on multimodal models and as a target for improving them.

한 줄 요약

다중모달 대형 언어 모델(MLLM)의 retrieval heads가 새로운 모달(음성-비주얼)로 확장되며, 이를 직접 최적화함으로써 문서 검색 성능을 향상시킨다.

핵심 기여도

18개의 언어, 비주얼-언어, 옴니모달 체크포인트에서 retrieval heads가 모달 확장 시에도 유지됨을 확인.
모달 간 공유되는 retrieval heads는 전체 attention head의 2% 미만이지만, 제거 시 문서 이해, 공간 인지, 시각 수학 성능이 저하됨.
retrieval heads의 attention을 ranking objective로 직접 최적화하여 Qwen3-VL-8B의 MMDocIR에서 page Recall@1을 6.03 포인트 향상.
3개 모델 가족에서 문서 검색 성능 개선을 일반화함.

핵심 아이디어

기존 연구는 retrieval heads가 텍스트 및 이미지 모델에서 정보를 추출하는 핵심 메커니즘이라는 점을 밝혔으나, 새로운 모달(예: 음성-비주얼)로 확장될 때의 역할은 명확하지 않았다. 본 연구는 retrieval heads가 새로운 모달로 확장될 때에도 유지되며, 모달 간 공유된다는 점을 발견하였다. 이는 retrieval heads가 모달 독립적인 핵심 메커니즘임을 시사한다. 또한, retrieval heads의 attention을 ranking objective로 직접 최적화함으로써 문서 검색 성능을 향상시킬 수 있음을 보여준다. 이는 retrieval heads가 해석 도구이자 최적화 대상이 될 수 있음을 의미한다.

기술적 접근법

**paired audio-visual inputs**: 동일한 정보가 음성 또는 화면 텍스트로 제공되는 입력을 구성하여 retrieval heads의 작동을 분석.
**head ablation**: retrieval heads 제거 실험을 통해 기능적 역할을 검증.
**ranking objective**: retrieval heads의 attention을 문서 검색 성능을 향상시키기 위해 ranking loss로 최적화.
**모델**: Qwen3-VL-8B, 다른 두 모델 가족 포함.
**데이터셋**: MMDocIR, 10개 문서 도메인, 4개 out-of-domain 벤치마크.

주요 결과

Qwen3-VL-8B의 MMDocIR에서 page Recall@1이 ranking training 후 6.03 포인트 향상.
retrieval heads 제거 시 audio-visual retrieval이 거의 사라짐.
모달 간 공유 retrieval heads(2% 미만) 제거 시 문서 이해, 공간 인지, 시각 수학 성능 저하.
3개 모델 가족에서 문서 검색 성능 개선.
layout retrieval도 layout supervision 없이 향상됨.

의의 및 한계

본 연구는 retrieval heads가 모달 확장 시에도 유지되며, 모델 이해와 개선에 동시에 기여할 수 있음을 보여준다. retrieval heads의 기능적 역할을 밝히고, 이를 최적화하여 문서 검색 성능을 향상시킨 점은 학술적·실용적 가치가 있다. 그러나 retrieval heads의 작동 원리에 대한 이론적 해석은 부족하며, 다양한 모델 아키텍처에 대한 일반화 가능성도 추가 연구가 필요하다. 또한, retrieval heads가 단순히 attention 메커니즘에 의존하는 한계도 존재한다.

실용적 활용

본 연구는 문서 검색, 다중모달 정보 추출, 비주얼 QA 등에서 retrieval heads를 활용한 모델 최적화에 적용 가능하다. 특히, 랭킹 기반 학습을 통해 retrieval heads의 attention을 직접 훈련함으로써, 다양한 모델 가족에서 문서 이해 및 검색 성능을 향상시킬 수 있다.