QG-MIL : un agrégateur Transformer à portes pour l'apprentissage multi-instances agnostique au domaine en imagerie médicale

Résumé

Les agrégateurs d'apprentissage multi-instance basés sur l'attention en imagerie médicale sont sujets à une concentration de l'attention, produisant des prédictions trop confiantes et instables. Nous présentons QG-MIL, un agrégateur à transformateur avec portes qui résout ce problème grâce à quatre composants architecturaux synergiques : la pré-normalisation basée sur RMSNorm, la normalisation QK par tête, le contrôle de sortie d'attention à grain fin, et les modules feed-forward de type SwiGLU. Ensemble, ces choix de conception stabilisent l'entraînement et répartissent l'attention plus uniformément entre les instances, sans pertes auxiliaires, masquage ni régularisation multi-étapes. Nous évaluons QG-MIL sur six référentiels couvrant la pathologie sur lames entières et l'hématologie au niveau cellulaire, représentant deux échelles MIL fondamentalement différentes. Les variantes les plus performantes de QG-MIL surpassent les références de premier plan sur l'ensemble des six référentiels, avec une amélioration moyenne de +6,1 points F1 macro moyens. Les superpositions d'attention et l'analyse de la masse d'attention confirment une pondération des instances plus répartie. Les études d'ablation montrent que, bien que des composants individuels puissent égaler le modèle complet sur des ensembles de données spécifiques, la conception QG-MIL offre la performance inter-domaines la plus cohérente et la variance la plus faible par rapport aux références sélectionnées. Nous publions une implémentation configurable pour soutenir la reproductibilité à l'adresse : https://github.com/unica-visual-intelligence-lab/QG-MIL.

English

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