From Point Estimates to Distributions: GMM Pooling for MIL in Preterm Birth Prediction

Resumo Técnico: De Estimativas Pontuais para Distribuições: GMM Pooling para MIL na Predição de Parto Prematuro

Formulação do Problema
O parto prematuro (PP), definido como o parto antes das 37 semanas de gestação, é uma causa significativa de morbidade e mortalidade neonatal. Embora as diretrizes clínicas recomendem a aquisição de múltiplas imagens de ultrassonografia transvaginal (USTV) por paciente para compensar a variabilidade dependente do operador na medição do comprimento cervical, as abordagens existentes de aprendizado de máquina para a predição de PP geralmente dependem de uma única imagem selecionada por paciente. Essa prática descarta a variabilidade informativa intra-exame (ex: diferenças de visualização e condições de aquisição) e limita o desempenho preditivo.

Para abordar isso, os autores formulam a predição de PP como um problema de Aprendizado de Instância Múltipla (MIL - Multiple Instance Learning). Neste framework, cada paciente corresponde a uma "sacola" (bag) de imagens de USTV de tamanho variável ($X = \{x_1, \dots, x_N\}$) associada a um único rótulo ao nível da sacola ($Y$) que indica o desfecho do parto (PP vs. termo). O desafio central é agregar esses números variáveis de características ao nível da instância em uma representação de dimensão fixa, preservando a variabilidade informativa dentro da sacola, em vez de colapsá-la em uma única estimativa pontual.

Metodologia: GMM Pooling
O artigo propõe um módulo de GMM pooling (agrupamento por Modelo de Mistura Gaussiana) para substituir os agregadores padrão de MIL (como max, mean ou attention pooling), que resumem uma sacola como uma única estimativa pontual no espaço de características. O método proposto modela explicitamente a distribuição de características dentro da sacola para capturar a heterogeneidade. A arquitetura consiste em três estágios:

1. Extração de Características: Um codificador compartilhado (implementado como um backbone U-Net para a tarefa de PP) mapeia cada imagem para um vetor de características da instância. O codificador é treinado conjuntamente para segmentação do canal cervical (como um objetivo auxiliar) e classificação ao nível da sacola para aprender representações anatomicamente significativas.
2. Estimativa de Parâmetros do GMM: Em vez de uma média simples, o método estima um GMM sobre as características da instância. Isso envolve:
   • Atribuições Suaves de Instância-para-Componente: Um MLP leve ($f_r$) gera uma distribuição de probabilidade ($r_{i,k}$) sobre $M$ componentes de mistura para cada instância.
   • Ponderação de Importância da Instância: Uma rede separada ($f_w$) calcula os pesos das instâncias ($w_i$) via softmax.
   • Combinação de Atenção: O peso de atenção final $a_{i,k}$ combina a importância da instância e a responsabilidade do componente ($a_{i,k} = w_i \cdot r_{i,k}$).
   • Aprendizado de Parâmetros: Usando esses pesos de atenção, os coeficientes de mistura ($\pi_k$), as médias ($\mu_k$) e as variâncias diagonais ($\sigma^2_k$) do GMM são estimados de ponta a ponta. Os autores optam por covariância diagonal para garantir estabilidade computacional e escalabilidade.
3. Incorporação da Distribuição (Distribution Embedding): Para produzir uma representação de comprimento fixo, o método avalia a log-densidade do GMM aprendido em um conjunto de $P$ vetores de sonda (probe vectors) aprendíveis. Os valores de log-densidade resultantes formam a incorporação final da sacola usada para a predição.

Principais Contribuições
Os autores identificam três contribuições primárias:

1. Módulo de GMM Pooling: Um novo mecanismo de pooling de MIL que preserva as distribuições de características informativas dentro da sacola, abordando o gargalo de informação inerente ao pooling de estimativa pontual.
2. Incorporação de Distribuição Diferenciável: Um método para construir uma incorporação de comprimento fixo via avaliações de log-densidade em vetores de sonda aprendíveis, permitindo que a representação de distribuição seja aprendida de ponta a ponta.
3. Avaliação Empírica Abrangente: Uma avaliação rigorosa em uma coorte privada de alto risco para PP e em um benchmark público de metástase de linfonodo, demonstrando que o MIL melhora em relação à predição ao nível da instância e que o GMM pooling oferece desempenho competitivo em tarefas de classificação e regressão.

Resultos Experimentais
O método foi avaliado em dois conjuntos de dados:

• Coorte Privada de PP: 182 pacientes grávidas de alto risco (44 casos de PP) com 1 a 43 imagens por paciente.
  • O MIL melhorou significativamente em relação aos modelos baseados em instância, elevando o PR-AUC de 0,44 para 0,56 (max pooling) e 0,56 (GMM pooling).
  • O GMM pooling alcançou um PR-AUC de $0,56 \pm 0,03$, comparável ao max ($0,57 \pm 0,05$) e attention pooling ($0,56 \pm 0,02$). No entanto, o GMM exibiu um desvio padrão menor entre as execuções, sugerindo maior estabilidade.
• Benchmark Público de Metástase de Linfonodo: 933 amostras de treinamento, 668 de validação e 736 de teste.
  • O GMM pooling alcançou desempenho de estado da arte com um F1-score de $0,91 \pm 0,01$ e ROC-AUC de $0,89 \pm 0,01$ para classificação.
  • Para regressão (porcentagem de pixels metastáticos), alcançou um Erro Médio Absoluto (MAE) de $0,18 \pm 0,02$, competitivo com os métodos de melhor desempenho.

Significância e Alegações
O artigo afirma que modelar distribuições de características via GMM pooling fornece uma alternativa robusta aos padrões de pooling de estimativa pontual. Enquanto os agregadores padrão colapsam a variabilidade dentro da sacola, o GMM pooling captura a estrutura de distribuição conjunta através das dimensões de características. Os autores observam que, embora a abordagem introduza hiperparâmetros específicos da distribuição (número de componentes, sondas, temperatura) que exigem ajuste dependente do conjunto de dados, os resultados em duas tarefas distintas de imagem médica sugerem que a modelagem de distribuição generaliza bem. O trabalho destaca que aproveitar todas as imagens disponíveis por paciente via MIL é crucial para a predição de PP, particularmente em coortes de alto risco onde a separabilidade de classes é baixa.