QG-MIL: A Gated Transformer Aggregator for Domain-Agnostic Multiple Instance Learning in Medical Imaging

2026-06-18

AI Summary

This work addresses the issue of overconfident and unstable predictions in attention-based multiple instance learning (MIL) for medical imaging, which often arises from overly concentrated attention distributions. To mitigate this, the authors propose QG-MIL, a novel gated Transformer aggregator that uniquely integrates RMSNorm pre-normalization, per-head query-key normalization, fine-grained attention output gating, and SwiGLU feedforward modules. This architecture enhances attention uniformity and model generalization without relying on auxiliary losses or complex regularization strategies. Evaluated across six benchmarks spanning whole-slide histopathology and cellular-level hematology, QG-MIL achieves an average improvement of 6.1 macro F1 points over existing methods and demonstrates significantly lower prediction variance.

Abstract

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

Problem

• Multiple Instance Learning
• Attention Concentration
• Medical Imaging
• Transformer Aggregator
• Domain-Agnostic

Innovation

• gated transformer
• multiple instance learning
• attention stabilization
• domain-agnostic
• medical imaging