Lightweight Joint Optimization of General-Purpose Vision-Language Models and Retrievers for RAG-Based Medical Diagnosis

TL;DR

The paper introduces Jomed, a lightweight framework that jointly optimizes a general-purpose LVLM and a multimodal retriever for medical diagnosis without heavy domain pretraining. By training the reader with retrieved context and sequentially fine-tuning a dual-head retriever, the approach aligns retrieval with the LVLM's predictive distributions, improving performance on classification and VQA in low-data settings. A key finding is that inconsistent retrieval predictions are a major source of error, and joint optimization substantially mitigates these cases, though an oracle analysis reveals a remaining gap that current rerankers cannot fully close. The work demonstrates practical, data-efficient benefits of integrating retrieval and reasoning, with broad implications for cost-effective medical AI deployment and future research directions.

Abstract

Retrieving relevant visual and textual information from medical literature and hospital records can enhance diagnostic accuracy for clinical image interpretation. We develop a multimodal retrieval model jointly optimized with an LVLM for medical diagnosis, unlike standard RAG which doesn't backpropagate LVLM errors to the retriever. Using only general-purpose backbones with lightweight fine-tuning, our model achieves competitive results with medically-pretrained models on clinical classification and VQA tasks. In a novel analysis, we find that different top-retrieved images often yield different predictions for the same target, and that these cases are challenging for all models, even for non-retrieval models. Our joint retrieval optimization significantly improves these cases over standard RAG. However, oracle analysis reveals that while the correct diagnosis is frequently achievable using one of the top retrieved images, in practice there is a large performance gap from the oracle, and rerankers using frontier LVLMs do not close this gap -- leaving ample room for improvement by future methods. Code available at https://github.com/Nirmaz/JOMED.

Figures (3)

• Figure 1: We jointly optimize a multimodal retriever and a Large Vision-Language Model for medical tasks. We achieve competitive results without resource-intensive medical pre-training and significantly improve performance on challenging cases where different retrieved images lead to inconsistent retrieval prediction.
• Figure 2: joint optimization impacts inconsistent retrieval predictions. After joint optimization, retrieved images show greater alignment with query image labels. In Breast, a cancer-free ultrasound query (a) initially retrieved a lesion image (b), however, after joint optimization, the retrieved image is less directly related to the wrong label (c). In VQA-RAD, retrieval shifted from an unrelated medical condition (e) to an image depicting the condition of the query image.
• Figure 3: The two-phase Jomed training. The LVLM is first trained with a frozen retriever on augmented prompts containing retrieved image–caption/report pairs. With the LVLM frozen, the dual-head multimodal retriever is optimized using a KL divergence computed over a selected subset of the model’s logits. An LLM converts open-ended questions into closed form (during training only), which improves results.