QG-MIL: Um Agregador Transformer com Mecanismo de Porta para Aprendizado de Múltiplas Instâncias Agnóstico de Domínio em Imagens Médicas

Resumo

Os agregadores de Aprendizado de Múltiplas Instâncias baseados em atenção em imagens médicas são propensos à concentração de atenção, gerando previsões excessivamente confiantes e instáveis. Apresentamos o QG-MIL, um agregador de transformador com portas que aborda esse problema por meio de quatro componentes arquitetônicos sinérgicos: pré-normalização baseada em RMSNorm, normalização QK por cabeça, controle granular da saída de atenção e módulos feed-forward estilo SwiGLU. Juntas, essas escolhas de design estabilizam o treinamento e distribuem a atenção de forma mais uniforme entre as instâncias, sem perdas auxiliares, mascaramento ou regularização em múltiplos estágios. Avaliamos o QG-MIL em seis benchmarks abrangendo patologia de lâmina inteira e hematologia em nível celular, cobrindo duas escalas fundamentalmente diferentes de MIL. As variantes do QG-MIL com melhor desempenho superam as principais referências em todos os seis benchmarks, com uma melhoria média de +6,1 pontos médios de macro F1. Sobreposições de atenção e análise de massa de atenção confirmam uma ponderação de instâncias mais distribuída. Estudos de ablação mostram que, embora componentes individuais possam igualar o modelo completo em conjuntos de dados específicos, o design do QG-MIL oferece o desempenho mais consistente entre domínios e a variância mais restrita em comparação com as referências selecionadas. Disponibilizamos uma implementação configurável para apoiar a reprodutibilidade em: 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