Solving the Needle-in-a-Haystack Problem in Mammography Vision-Language Model with Differentiable Subset Sampling

1 Introduction

• Prior mammography VLMs focused mainly on multi-view learning, leaving high resolution and report homogeneity as additional challenges.Mammography VLM work grew from broader CLIP-style pretraining used across medical imaging modalities.Source: §1 · pp. 1–2
• Mammograms can exceed 12 megapixels, and small, sparse lesions risk being erased by downsampling; full-resolution CLIP training instead constrains batch size and may degrade performance.Most mammography datasets lack enough spatial annotations to train a reliable region-proposal network.Source: §1 · p. 2
• Screening-dominant mammography yields highly similar reports, violating CLIP’s assumption that captions in a batch are semantically distinct and causing similar image–report pairs to be repelled.In EMBED, 75% of exams are BI-RADS 1, with reports differing mainly in density and demographics.Source: §1 · p. 2
• TopKSigLIP introduces TopK-Patch to select likely lesion-containing high-resolution patches and Sup-sigmoid to use structured-data soft labels for image–report learning.The patch scores and locations also provide built-in localization, while sigmoid scoring treats pairs independently rather than forcing competition across the batch.Source: §1 · pp. 2–3
• The training pipeline pairs selected high-resolution image patches with text embeddings and jointly optimizes semantic-similarity and, when boxes exist, ROI supervision.The text tower encodes LLM-preprocessed radiology reports; ROI supervision uses Soft Dice when bounding-box annotations are available.Source: Figure 2 · p. 3
• TopKSigLIP achieves state-of-the-art performance on internal and external benchmarks across four prediction tasks in zero-shot and linear-probing settings.Its learned patch locations also outperform Grad-CAM for lesion detection.Source: §1 · p. 3

Mammography VLMs

• Mammography VLMs address views through exam-level aggregation, single-view modeling, local contrastive learning, or cross-view attention, but do not tackle high resolution or report homogeneity.The cited models include Mammo-CLIP, MaMA, and GLAM.Source: §2 · p. 3

Differentiable top-k learning

• Differentiable top-k sampling can use entropy-regularized optimal transport or perturbed optimization; this work adopts Gumbel-softmax built on weighted reservoir sampling.Source: §2 · p. 3

CLIP alternatives

• MedCLIP and UniCL use soft or hard targets to address semantically similar negatives, while prior efficiency approaches freeze the image encoder or replace softmax with sigmoid.The related-work passage identifies high-resolution memory costs in fully end-to-end trainable settings as unaddressed.Source: §2 · p. 3

3.1 Overview

• The model computes image–report similarities for all pairs in a minibatch and optimizes them to predict a label matrix.Vision and text towers produce L2-normalized embeddings; the text tower follows a standard transformer architecture.Source: §3.1 · p. 4

3.2 Vision tower

• The vision tower selects K high-resolution patches using scores from a downsampled image, encodes the patches, and aggregates their embeddings into one image embedding.The input combines the CC and MLO views of one breast.Source: §3.2 · p. 4

3.2.1 TopK-Patch: Differentiable Top-k Patch Selection

• TopK-Patch samples locations without replacement according to learned patch scores, using Gumbel-top-K reparameterization and a temperature-annealed softmax relaxation for gradient flow.The reparameterization yields exact Plackett–Luce samples; the relaxation is applied to the top-k operation.Source: §3.2.1 · p. 4

3.3.1 Sup-sigmoid loss

• Sup-sigmoid replaces identity targets with pairwise similarities derived from BI-RADS, density, cancer outcome, and lesion-subtype labels.The loss operates on predicted and ground-truth similarity matrices; sigmoid loss treats pairs independently.Source: §3.3.1 · p. 5

3.3.2 ROI supervision

• ROI supervision compares a mask aggregated from sampled patch locations with a binary mask generated from bounding-box annotations.ROI annotations are available for approximately 4% of EMBED exams and are used when available.Source: §3.3.2 · p. 5

3.3.3 Final objective

• The final objective adds a weighted Soft Dice loss to Sup-sigmoid, applying Dice only to the approximately 4% of samples with ROI labels.The scalar weight is λdice.Source: §3.3.3 · p. 5

4.1 Datasets

• TopKSigLIP is trained on EMBED cohorts 1–2, validated on cohort 8, and tested on cohorts 9–10; EMBED contains 364,000 US exams.Each exam includes multi-view images, tabular results, a radiology report, and ROI annotations for 4% of exams.Source: §4.1 · p. 5
• VinDr provides 5,000 Vietnamese exams with BI-RADS, density, and non-benign-finding ROI labels, using its predefined 4,000/1,000 train/test split.Its training set is used only for linear-probe training, with 20% held out for validation.Source: §4.1 · p. 5
• RSNA provides 11,913 US and Australian screening exams with cancer outcomes and BI-RADS scores, split 80/20 for train/test without ROI or finding labels.Its training set is used only for linear-probe training, with 20% held out for validation.Source: §4.1 · p. 5

4.2 Baselines

• The mammography baselines are Mammo-CLIP-B2/B5, MaMA, and GLAM; MII is the general-purpose medical VLM.Mammo-CLIP models use UPMC and VinDr training data, while MaMA and GLAM use EMBED cohorts 1 and 2.Source: §4.2 · p. 6
• For single-view Mammo-CLIP and MII models, multi-view embeddings are formed by averaging per-view outputs.The comparison follows each model’s proposed image and text preprocessing and published zero-shot prompts where available.Source: §4.2 · p. 6
• MII is excluded from RSNA evaluation because it was trained on RSNA and other non-breast radiology images.Source: §4.2 · p. 6

4.3 Implementation details

• Reports are processed with LlaMA-3 to extract mammography findings, excluding patient metadata and findings from non-mammography modalities.Only findings matching the laterality of the breast’s paired CC and MLO views are retained; pathology cancer outcomes are appended.Source: §4.3 · p. 6
• CC and MLO views per breast are resized to 1536 × 768 and vertically concatenated into a 1536 × 1536 image after background removal.When standard views are unavailable, other available views are substituted; the low-resolution input is downscaled by 3 to 512 × 512.Source: §4.3 · p. 6
• The weight-shared vision encoders are ImageNet-pretrained ConvNeXt-Tiny models, with the last convolutional block removed; the CLIP text encoder’s context length is extended to 170.The patch grid has 256 candidates, and K is set to 15 during training.Source: §4.3 · p. 6
• Training runs for 100,000 iterations with AdamW, learning rate 1 × 10−5, weight decay 0.05, and batch size 32.Training uses one RTX PRO 6000; λdice is 5, and the TopK-Patch temperature is annealed from 1 to 2/3.Source: §4.3 · p. 6

4.4 Evaluation protocol

• AUROC evaluates BI-RADS, density, finding-type, and cancer-outcome tasks in both zero-shot and linear-probing settings.BI-RADS categories 0 and 6 are excluded from the BI-RADS task.Source: §4.4 · p. 6
• Linear probing adds a single fully connected layer to a frozen image encoder and trains it separately for each task and dataset.Training uses class-weight-adjusted cross-entropy with AdamW for 20,000 iterations, at learning rate 4 × 10−3, weight decay 0.01, and batch size 32.Source: §4.4 · p. 6
• Localization is evaluated by pixel-level average precision and Pointing Game against ground-truth ROIs.TopKSigLIP uses patch-importance scores directly; baseline heatmaps use label-guided Oracle Grad-CAM in both evaluation settings.Source: §4.4 · p. 6

4.5.1 Accuracy

• +29.2 AUC points on VinDr and +17.8 AUC points on EMBED over prior best are the largest zero-shot gains, on finding detection and cancer prediction.TopKSigLIP achieves state-of-the-art performance across the reported tasks and datasets; k=15 covers about 10% of the image and is competitive, with diminishing returns at larger k.Source: §4.5.1 · p. 7, Table 1 · p. 7
• Under linear probing, the performance gap narrows, but TopKSigLIP retains clear advantages on finding detection and cancer prediction.The authors conjecture the gap reflects its lightweight ConvNeXt-Tiny encoder and batch size 32, compared with larger baseline encoders and batches above 100.Source: §4.5.1 · p. 7, Table 2 · p. 7

4.5.2 Localization

• TopKSigLIP’s patch-importance map outperforms Oracle Grad-CAM across baselines in zero-shot and linear-probe localization.TopKSigLIP’s linear-probe localization is omitted because its frozen vision tower produces the same heatmaps; Figure 3 shows the heatmaps alongside ground-truth ROI boxes and sampled top-15 patches.Source: §4.5.2 · p. 8, Table 3 · p. 8, Figure 3 · p. 8
• Baseline Grad-CAM heatmaps can produce false positives on high-intensity regions and clinically irrelevant anatomy, while TopKSigLIP’s false-positive patches tend to have plausible imaging features.The authors also observe that MaMA attends broadly across breast tissue rather than focusing on specific findings, and conjecture that global conditioning may explain its competitive linear-probe performance despite poor localization.Source: §4.5.2 · pp. 8–9

4.6 Ablation studies

• TopKSigLIP exceeds full-image encoding at 1536^2 on most tasks while using 64.2 GB and remaining trainable at batch size 32.Prediction accuracy and localization improve with resolution, but 1024^2 exceeds 77 GB and 1536^2 is out of memory at batch size 32; 1536^2 training uses multiple GPUs.Source: §4.6 · p. 9, Table 4 · p. 9
• Removing ROI supervision causes a large accuracy drop, indicating that it is essential for training TopK-Patch despite appearing in only about 4% of training samples.The ablation compares full-image encoding at four resolutions with TopKSigLIP across efficiency, prediction accuracy, and localization.Source: §4.6 · p. 9, Table 4 · p. 9
• With tabular supervision, sigmoid yields large gains, whereas supervised contrastive loss yields only a small benefit.The ablation holds other components fixed and compares contrastive versus sigmoid objectives with identity versus tabular-label supervision; the authors conclude sigmoid’s independent pair treatment is necessary to exploit non-identity targets.Source: §4.6 · p. 9, Table 5 · p. 9

• TopKSigLIP addresses mammography’s high-resolution inputs with sparse findings and homogeneous screening reports.Source: §5 · p. 9
• The model outperforms prior approaches on prediction accuracy and localization across internal and external benchmarks.Source: §5 · p. 9
• The paper argues that CLIP-style architectures and objectives can be modified to address domain-specific failure modes.Source: §5 · p. 9

5.0.1 Broader Impact

• The authors identify high-resolution whole-slide images and normal-skewed screening reports as adjacent domains where these methods may be useful starting points.Whole-slide images can exceed 100,000 × 100,000 pixels.Source: §5.0.1 · p. 14

5.1 Label Statistics

• The three datasets span diverse demographics and skewed label distributions; Table 6 reports breast-level label statistics, with counts roughly twice exam counts.EMBED includes 41% African American participants; VinDr is a Vietnamese cohort, and RSNA includes US and Australian cohorts.Source: §5.1 · pp. 14–15, Table 6 · p. 15

5.2 Report Parsing

• Report parsing uses Llama-3 and excludes patient metadata and ultrasound or MRI findings, while retaining laterality-specific findings and explicit negative semantics.Breast-level BI-RADS and biopsy cancer outcomes come from structured data; empty finding fields receive a default negative sentence.Source: §5.2 · pp. 15–16

5.3.1 Zero-shot evaluation

• Zero-shot evaluation uses published prompt templates where available and designs vocabulary-aligned prompts for tasks lacking supplied templates.Source: §5.3.1 · p. 16

5.3.2 Localization evaluation

• Localization evaluation gives baseline Grad-CAM ground-truth labels to form the target logit, then uses feature-map gradients to produce heatmaps.For linear probes, finding, cancer, and BI-RADS logits are combined using ground-truth labels or class indices.Source: §5.3.2 · pp. 16–17

5.4 Grad-CAM and NormGrad on linear probe

• Table 13 compares AP/PG localization for Grad-CAM and NormGrad without label supervision, with Oracle Grad-CAM shown for reference.The evaluation also varies which logit heads generate the heatmap.Source: §5.4 · p. 17, Table 13 · p. 20

5.5 Proof: Differentiable Top-k as Exact Sampling from the Plackett–Luce Distribution

• As temperature τ approaches zero, iterative softmax recovers hard top-K, and Gumbel-perturbed weights yield exact Plackett–Luce samples.Figures 4 and 5 show localization visualizations on random EMBED and VinDr samples.Source: §5.5 · pp. 18–20, Figure 4 · p. 21, Figure 5 · p. 22