Source-linked AI summary
QG-MIL: A Gated Transformer Aggregator for Domain-Agnostic Multiple Instance Learning in Medical Imaging
Luca Zedda, Davide Antonio Mura, Cecilia Di Ruberto, Maurizio Atzori, Muhammed Furkan Dasdelen, Carsten Marr, Andrea Loddo
TL;DR
Attention-based MIL can concentrate attention on a few instances, causing overconfident and unstable predictions. QG-MIL uses a gated transformer aggregator to address this architecturally, outperforming leading baselines across six pathology and hematology benchmarks by an average of +6.1 macro F1 points.
Problem
Attention-based MIL can collapse attention onto a few instances, yielding overconfident predictions and limiting generalization across cohorts and imaging sources.
Method
QG-MIL is a domain-agnostic gated transformer aggregation module that mitigates attention concentration through architectural design without auxiliary losses or multi-stage training.
Results
+6.1 macro F1 points average improvement over leading baselines across six pathology and hematology benchmarks.
Takeaways & Limitations
QG-MIL provides more consistent cross-domain performance alongside smoother, more distributed attention weighting across medical imaging MIL tasks.
Takeaways & Limitations
The full gated architecture can introduce variance in small patient cohorts, while fine-grained gating and SwiGLU increase training memory and time.
Abstract
from arXiv · showhide
Attention-based Multiple Instance Learning aggregators in medical imaging are prone to attention concentration, producing overconfident and unstable predictions. We introduce QG-MIL, a gated transformer aggregator that addresses this through four synergistic architectural components: RMSNorm-based pre-normalization, per-head QK normalization, fine-grained attention output gating, and SwiGLU-style feed-forward modules. Together, these design choices stabilize training and distribute attention more uniformly across instances without auxiliary losses, masking, or multi-stage regularization. We evaluate QG-MIL across six benchmarks spanning whole-slide pathology and cell-level hematology, covering two fundamentally different MIL scales. The best-performing QG-MIL variants outperform leading baselines on all six benchmarks, with an average improvement of +6.1 mean macro F1 points. Attention overlays and attention mass analysis confirm more distributed instance weighting. Ablation studies show that while individual components can match the full model on specific datasets, the QG-MIL design provides the most consistent cross-domain performance and tightest variance when compared to selected baselines. We release a configurable implementation to support reproducibility at: https://github.com/unica-visual-intelligence-lab/QG-MIL
1 Introduction
Attention-based MIL uses instance weights for prediction and interpretability but can suffer from attention concentration. QG-MIL addresses this through architectural changes alone and is evaluated across six pathology and hematology benchmarks with ablations.
- Motivation: Attention-based MIL assigns instance weights that drive predictions and proxy instance importance, but attention concentration is a documented failure mode.This limits the reliability of attention as an interpretability signal.
- Method: QG-MIL is a gated transformer aggregation module that mitigates attention concentration without auxiliary losses or multi-stage training.It is designed as a drop-in replacement for standard MIL pipelines.
- Evaluation: QG-MIL is evaluated on six benchmarks spanning pathology and hematology, assessing predictive performance, attention distribution, and localization quality.The evaluation includes the QG-MIL model and its ablations.
- Ablations: Ablation studies isolate each design choice and examine how cohort size and model depth affect ablation performance.These analyses assess the contribution of individual architectural choices and experimental conditions.
2 Methodology
QG-MIL projects bag instances into a shared latent space, refines them with stacked gated transformer blocks, and aggregates them using gated attention pooling for bag classification. Its methodology includes per-head Q/K normalization, attention-output gating, pre-norm residual blocks with SwiGLU feed-forward networks, and ablation variants evaluating these design choices.
- Architecture: QG-MIL projects instances into a shared latent space, processes them with L gated transformer blocks, and produces a bag representation for classification.The method denotes input, attention, gating, and output projections within this pipeline.
- Transformer blocks: Per-head queries and keys are normalized before scaled dot-product multihead attention to improve stability.Queries and keys are reshaped into H heads with per-head dimension d_h, satisfying D = H d_h.
- Transformer blocks: Attention outputs are gated before projection, using either headwise scalar gates or elementwise feature gates.Headwise gates modulate all features uniformly, whereas elementwise gates independently modulate features within each head.
- Transformer blocks: Each block uses a pre-norm residual layout with a SwiGLU feed-forward network, and stacking L blocks yields refined instance embeddings.The refined embeddings are subsequently passed to gated attention pooling.
- Pooling and ablations: Gated attention pooling aggregates the processed instances into a bag representation, which a classification head uses to predict the bag label.Evaluated ablations include Elementwise, noGate, noQKnorm, LayerNorm, and Light variants.
3 Experiments and Results
Across pathology and hematology benchmarks, QG-MIL’s performance depends on cohort size and depth, while its gating mechanisms produce more distributed attention than ABMIL. The framework achieves strong diagnostic and prognostic results, including 93.3% macro F1 on Lung, 57.5% on Prostate, and 69.5% on APL classification.
- Experimental protocol: A patient-stratified 20% hold-out test set and five-fold cross-validation on the remaining 80% support final evaluation through mean-probability ensembling.Five independent fold-specific models are trained and ensembled on the fixed test set.
- General architectural and performance findings: QG-MIL’s optimal complexity is cohort-dependent: low-patient datasets favor relaxed configurations, whereas larger cohorts benefit from deeper aggregation variants.Stringent gating and Q/K normalization can increase variance in smaller cohorts, while deeper variants improve performance when sufficient data are available.
- Pathology benchmarks on diagnosis: 93.3% macro F1 is achieved on the Lung benchmark, exceeding previously reported classification accuracies of 81.6% and 81.0%.The low-patient Lung dataset favors relaxed QG-MIL configurations.
- Pathology benchmarks on prognosis: 57.5% overall average macro F1 is achieved by QG-MIL Layernorm on the challenging Prostate prognostic task, compared with 54.1% for baseline methods.The standard Transformer peaks at 56.9% macro F1, while QG-MIL Layernorm achieves the highest overall average.
- Hematology benchmarks on diagnosis: 69.5% macro F1 is achieved for APL versus non-APL classification with the deep configuration, while QG-MIL outperforms hematology baselines by an average of over 5% macro F1.This follows the depth-to-regularization trend observed in the Lung dataset.
- Attention distribution: QG-MIL produces more uniform attention than ABMIL, with Gini 0.16±0.07 and entropy 0.99±0.01 versus 0.62 ± 0.13 and 0.90 ± 0.05.On the Breast test set, QG-MIL maintains approximately 15% mean top-10% attention mass, indicating reduced attention concentration.
4 Conclusion and Limitations
QG-MIL is presented as a domain-agnostic gated transformer aggregator that mitigates attention sinks without auxiliary losses, improving performance and attention distributions across six benchmarks. The study identifies adaptive gating, cohort-size-dependent complexity, multimodal evaluation, and experiment emissions as future-work or reporting considerations.
- Conclusion: +6.1 macro F1 points was the average improvement over leading baselines across six pathology and hematology benchmarks.QG-MIL also yielded a smoother, more clinically plausible attention distribution.
- Limitations and Future Work: Future work will explore adaptive gating that dynamically scales complexity based on cohort size.The authors also plan to expand evaluations to multimodal clinical pipelines to further validate generalizability.
- Limitations and Reporting: 7.78 kg CO2eq was the estimated total emissions from experiments run on in-house NVIDIA A100 PCIe 80 GB infrastructure.This reports the computational emissions associated with the experiments.