Skip to content

RAU: Reference-based Anatomical Understanding with Vision-Language Models

Conference: ECCV2026
Paper: ECCV paper page
Authors: Yiwei Li, Yikang Liu, Jiaqi Guo, Lin Zhao, Zheyuan Zhang, Xiao Chen, Boris Mailhe, Ankush Mukherjee, Terrence Chen, Shanhui Sun
Affiliations: United Imaging Intelligence; University of Georgia; Northwestern University
Area: Medical Imaging
Keywords: anatomical understanding, reference-conditioned segmentation, spatial reasoning, SAM2, reinforcement learning

Identity check: the author list in acceptance record 3449 matches the paper's first page. The PDF uses "Vision Language Models" and the conference title uses "Vision-Language Models"; only the hyphen differs. This note retains the conference title.

TL;DR

RAU trains Qwen2.5-VL-7B to compare anatomical spatial relationships between reference and target images, then connects segmentation tokens to SAM2 reference memory as soft prompts, achieving Dice scores of 0.7503 on CAMUS and 0.7010 on LERA when segmentation training uses only RAOS and Arcade.

Background & Motivation

Recognizing anatomy does not always mean naming an organ. A particular coronary segment, a locally similar tissue region, or a corresponding region across examinations may have little distinctive texture, making position and neighboring structures more informative. Atlas registration can propagate labels, but thin vessels, motion artifacts, incomplete views, and anatomical changes can destabilize it; directly training a segmentor requires expensive expert annotations.

A vision-language model (VLM) can use medical knowledge and follow instructions, yet knowing an organ's typical appearance is different from finding the region corresponding to a reference annotation. RAU therefore explicitly supplies an annotated reference and asks the model to infer correspondence through relative spatial relationships. The task is reference-conditioned identification, localization, and segmentation, not general medical question answering or segmentation without annotations.

The output representation also matters: a discrete region index only identifies a match, whereas a rectangle poorly represents curved, branching, or overlapping vessels. Core Idea: use reference-guided spatial reasoning to select the intended region, then use SAM2 with reference-mask memory to delineate its boundary, connecting these capabilities through task training and reinforcement learning.

Method

Overall Architecture

The inputs are an unlabeled target image, a reference image with region annotations, and an instruction specifying the queried region; outputs can be a region index, bounding boxes, or a pixel mask. The paper explores VQA, box localization, and segmentation in sequence. The final segmentor initializes from RL-VQA weights and does not require running the box-localization model first.

DINOv2 feature retrieval selects the reference. For segmentation, its region annotations support the VLM's spatial comparison, while the reference image and corresponding mask enter SAM2 memory. An MLP projects the hidden state of a generated <Seg> token into a soft spatial prompt used in memory reading and mask decoding, without explicit target-image point or box prompts.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["Target image and instruction"] --> Retrieve["Reference Retrieval"]
    Bank["Annotated reference bank"] --> Retrieve
    Retrieve --> Reason["Reference-guided Reasoning"]
    Reason -->|box variant| Box["Token Regression and<br/>Global Matching"]
    Box --> Labels["Boxes and region labels"]
    Reason -->|segmentation token| Fuse["Soft-prompt and<br/>Reference-memory Fusion"]
    Retrieve -->|reference image and mask| Fuse
    Fuse --> Mask["Target anatomical mask"]

Key Designs

1. Reference Retrieval: provide a comparable anatomical layout

DINOv2 extracts features from the target and reference-bank images, and cosine similarity selects the best-matching annotated reference. The VLM thus receives a pair of images suitable for region correspondence, rather than an isolated medical image. Retrieval selects a reference; it neither directly registers the reference mask to the target nor guarantees visibility of every small structure.

The reference serves two roles: overlaid region annotations tell the VLM what the instruction refers to, while memory encoded from the reference image and mask conditions SAM2's detailed segmentation. This condition is essential, so cross-dataset generalization should not be read as requiring no annotated references from the new dataset. Reference-bank construction is deferred to supplementary material absent from the cache; bank size, sampling, and patient-level splitting cannot be verified here.

2. Reference-guided Reasoning: learn correspondence instead of memorizing organ appearance

The VQA task assigns candidate indices to reference regions and asks the model to identify the match using relationships between the target region and surrounding structures. The prompt explicitly discourages matching by color and emphasizes spatial relations. SFT first adapts Qwen2.5-VL-7B to paired images and constrained answers, after which GRPO rewards correct indices and valid formatting.

This stage uses only RAOS-CT. Table 1 labels the training data RAOS (650k), which must not be interpreted directly as 650k patients. Training redirects the model from producing plausible medical descriptions toward making verifiable region decisions for a particular image pair. The authors use chain-of-thought examples to support spatial reasoning, but the later reference-layout shuffling ablation offers more direct evidence.

The VQA initialization also supplies the starting point for the final segmentor. RL-VQA is not a separate inference-time judge: it denotes the same VLM after task-specific training. Subsequent segmentation learning retains the reference-conditioned input while changing the output interface and optimization signal.

3. Token Regression and Global Matching: avoid generating precise coordinates as text

The box variant does not autoregressively emit numerical coordinates. It generates one or more <Seg> tokens and regresses boxes from their embeddings with a lightweight MLP. All label-associated boxes are predicted together and jointly matched to reference-derived label prototypes, reducing duplicate assignments and spatial conflicts caused by independent category decisions.

The paper formulates this assignment as optimal transport: construct a spatial or semantic cost matrix between predicted boxes and label prototypes, then minimize total transport cost under uniform marginal constraints. Unmatched or low-confidence regions trigger a fallback that regenerates candidate boxes. The cost construction, prototypes, and fallback thresholds are deferred to the supplement, so they cannot be replaced with an assumed implementation.

This variant tests whether explicit geometric correspondence is sufficient. It improves over RL-VQA on RAOS labeling accuracy but performs substantially worse on Arcade vessels: rectangular regions poorly preserve the identity of thin, branching structures. There is also an unresolved wording conflict: Section 3.2's method predicts target-image regions, whereas its experimental settings describe corresponding boxes in the reference image.

4. Soft-prompt and Reference-memory Fusion: pass region selection to a pixel decoder

The segmentation branch still emits <Seg>, but its MLP projects the hidden state into SAM2's memory space instead of regressing four box parameters. Reference images and masks are encoded as memory slots; projected queries read these slots, and the fused representation enters the SAM2 decoder. The VLM conditions which corresponding structure to attend to, while SAM2 uses target-image features and reference memory to recover boundaries.

The following reconstructs the dot-product softmax and weighted sum explicitly described around Equation (7). The cached equation is typographically damaged; this notation restores that stated mechanism without adding another component:

\[ \alpha_{ij}=\frac{\exp(\mathbf{q}_i^\top\mathbf{m}_j)}{\sum_k\exp(\mathbf{q}_i^\top\mathbf{m}_k)},\qquad \mathbf{z}_i=\sum_j\alpha_{ij}\mathbf{m}_j. \]

Here \(\mathbf{q}_i\) is a projected segmentation-token query, \(\mathbf{m}_j\) a reference-memory slot, and \(\mathbf{z}_i\) the fused representation. The reference is therefore not merely an additional image in the prompt: it also enters the pixel-prediction pathway. Relative to a baseline that only supplies reference memory to SAM2, RAU adds target-selection conditioning trained through reference-guided reasoning.

A Worked Example

Consider a coronary segment in Arcade: the reference identifies the intended segment, while the target X-ray may contain several nearby branches with similar appearance. DINOv2 selects a reference, and the VLM forms a region condition from the relative positions of these branches, not from shared display colors.

The box variant compresses candidate regions into rectangles and performs global assignment, potentially including neighboring branches. The final segmentor instead sends the <Seg> hidden state into reference-memory fusion and has SAM2 directly predict the desired segment mask. This illustrates the task flow rather than a new quantitative experiment, and it does not imply guaranteed vascular connectivity.

Loss & Training

VQA uses SFT followed by GRPO, with a task-level reward combining index accuracy and formatting validity. Table 1 compares 3 and 5 SFT epochs and 800, 1600, and 2400 GRPO steps; these budgets should not be treated as compute-matched conditions.

Box training first freezes the VLM to optimize the MLP, then fine-tunes with GRPO. Equation (5) summarizes the reward using detection AP and format validity, whereas Figure 2 also shows thinking-format, JSON-format, and box-IoU-above-0.5 rewards. The task-level equation is therefore not a complete implementation specification.

Segmentation initializes from RL-VQA. During SFT, segmentation-token embeddings, the MLP, and the SAM2 decoder are trained while other VLM weights remain frozen. Figure 2 gives the supervised loss as:

\[ \mathcal{L}_{\mathrm{seg}}=\alpha\mathcal{L}_{\mathrm{BCE}}+\beta\mathcal{L}_{\mathrm{Dice}}. \]

GRPO subsequently optimizes the VLM with a segmentation-quality reward related to Dice and BCE. The prose describes unfreezing the VLM, whereas Figure 2 shows LoRA. The exact trainable parameter scope, loss-to-reward transformation, numerical weights, and compute costs cannot be determined from this main-text cache; neither reward signs nor training hyperparameters should be invented.

Key Experimental Results

Main Results

Segmentation training uses RAOS-CT and Arcade-X-Ray; cardiac ultrasound from CAMUS and bone X-ray from LERA are out-of-distribution tests. The following is selected from original Table 3. Dice is a dimensionless overlap score, higher is better, and gains are direct score differences rather than relative percentages.

Dataset Setting MedSAM2-Memory-Ref-SFT Dice RAU Dice Gain
Arcade ID, vascular X-ray 0.2435 0.6754 +0.4319
RAOS ID, CT 0.2965 0.7151 +0.4186
CAMUS OOD, cardiac ultrasound 0.4290 0.7503 +0.3213
LERA OOD, bone X-ray 0.2603 0.7010 +0.4407

Even medically specialized SAM2 with reference memory and SFT struggles to select the correct region, suggesting that contour decoding alone is insufficient. Table 3 also reports a metric named gIoU, but the current main text does not clearly define its calculation or aggregation. This note does not assume it means either generalized IoU or global IoU.

Original Table 2 measures best labeling accuracy, not Dice; the percentages below preserve that distinction.

Dataset RL-VQA RL-Bbox VLM+SAM2
RAOS 74.68% 78.16% 89.38%
Arcade 64.37% 41.09% 81.62%
LERA 54.57% 55.92% 61.87%
CAMUS 88.33% 55.06% 95.41%

Table 1's RAOS-only VQA reaches 70.68% after 2400 GRPO steps, unlike Table 2's 74.68%; the two tables must not be treated as the same checkpoint. Table 1's Mixture pools OOD datasets and should not be counted again as an independent dataset in a simple cross-domain average.

Ablation Study

The following selects dimensionless Dice scores from original Table 5. Reference shuffling is reported only for ID data. Vanilla and SFT variants retain the reference-conditioned pipeline while changing VLM weights or training; GT Bbox is an oracle-localized SAM2 baseline prompted with ground-truth boxes.

Config Arcade RAOS CAMUS LERA
RAU, VLM-Vanilla 0.1072 0.1968 0.2644 0.1231
RAU, VLM-SFT 0.4370 0.5093 0.5506 0.5041
SAM2, GT Bbox 0.2168 0.2972 0.4480 0.3216
Upper/lower reference halves shuffled 0.0966 0.2088 Not reported Not reported
RAU, full model 0.6754 0.7151 0.7503 0.7010

Key Findings

  • Relative to VLM-SFT, full RAU raises CAMUS Dice from 0.5506 to 0.7503, a difference of 0.1997; the training recipe's benefit is not confined to ID data.
  • Shuffling the reference layout reduces Arcade Dice from 0.6754 to 0.0966, supporting reliance on spatial correspondence rather than merely which categories appear. However, this artificial perturbation also causes distribution shift and does not independently prove human-like reasoning.
  • Section 4 initially cites Table 2 in error; the initialization ablation is Table 4, where LERA labeling accuracy is 33.95% with SFT-VQA initialization and 61.87% with RL-VQA initialization.

Highlights & Insights

  • The reference both specifies the query and supplies segmentation memory. These roles connect correct structure selection with accurate delineation, targeting correspondence more directly than simply increasing medical pretraining data.
  • The negative box-localization result is informative. More explicit geometric output is not necessarily better for vessels; the output representation must suit the target's shape and topology.
  • A transferable idea is to train verifiable region correspondence before fine-grained decoding. Applying it to other thin structures still requires suitable reference annotations and cross-domain evaluation.

Limitations & Future Work

  • The authors propose structural hierarchies or organ-tree priors, adaptive memory, and temporal and 3D extensions. Current results do not establish performance on volumetric data or dynamic clinical workflows.
  • The cache lacks supplementary material, preventing verification of reference banks, data splits, reward weights, fallback rules, and costs. The main text also contains ambiguities about box-coordinate direction, trainable parameters, and checkpoints across tables.
  • OOD transfer remains conditional on annotated references. The tables report best labeling accuracy without confidence intervals or repeated-run variation verifiable here, so they do not establish clinical readiness or statistical significance.
  • vs SAM2 / MedSAM2: these baselines use reference images and masks as memory; RAU adds reference-guided VLM conditioning to select the corresponding structure. The comparison concerns this reference-conditioned setting, not every medical segmentation task.
  • vs LISA / SegZero: all connect language reasoning to segmentation, but RAU emphasizes anatomical correspondence between paired images and reference memory. Table 3 compares LISA-SFT and SegZero-SFT, not the upper bounds of all their training recipes.
  • vs atlas registration: registration estimates a spatial transform and propagates labels; RAU learns reference-conditioned region correspondence and decoding. It does not output an interpretable deformation field, so the approaches have different reliability constraints.

Rating

  • Novelty: 4/5. The combination of reference-guided spatial reasoning and memory-based segmentation targets a specific task.
  • Experimental Thoroughness: 4/5. ID, OOD, initialization, oracle-box, and reference-perturbation comparisons are included, but missing supplementary details limit reproducibility assessment.
  • Writing Quality: 3/5. The task progression is clear; table references, box direction, and training configuration need clarification.
  • Value: 4/5. Useful research on anatomy correspondence with limited annotations, not evidence of clinical deployability.