Learning Generalizable 3D Medical Image Representations from Mask-Guided Self-Supervision

1. O Problema

A área de inteligência artificial médica enfrenta uma lacuna significativa em comparação com a visão computacional geral e o processamento de linguagem natural: a falta de modelos fundamentais (foundation models) para imagens médicas 3D.

• Limitações dos Métodos Atuais: Os métodos de auto-supervisão (SSL) existentes para imagens médicas geralmente dependem de objetivos de reconstrução de baixo nível (como Masked Autoencoders - MAE) ou objetivos contrastivos. Essas abordagens falham em capturar a semântica anatômica crítica para a análise médica, focando apenas em texturas ou características globais da imagem.
• Custo de Anotação: A abordagem supervisionada tradicional exige anotações pixel a pixel feitas por especialistas, o que é extremamente caro, demorado e limitado a taxonomias pré-definidas. Isso impede a escalabilidade para grandes conjuntos de dados e a generalização para estruturas anatômicas ou patologias não vistas anteriormente.
• Complexidade Médica: Diferente de imagens naturais, onde o fundo é dominante, quase todos os voxels em uma imagem médica têm significado clínico. Além disso, a precisão espacial é vital (ex: a localização exata de uma hemorragia ou o tamanho de um tumor alteram o tratamento), algo que métodos genéricos de SSL muitas vezes ignoram.

2. Metodologia: MASS

Os autores propõem o MASS (MAsk-guided Self-Supervised learning), um novo framework de auto-supervisão que utiliza a segmentação em contexto (in-context segmentation) como tarefa pretext, guiada por máscaras geradas automaticamente.

A. Geração de Máscaras sem Anotação (Annotation-Free Mask Generation)

O núcleo da inovação é a eliminação da necessidade de anotações humanas.

• Uso de SAM2: O framework utiliza o modelo Segment Anything Model 2 (SAM2), treinado em imagens naturais, para gerar máscaras agnósticas de classe (class-agnostic) em volumes médicos 3D.
• Pipeline:
  1. Pré-processamento: Conversão de volumes médicos (CT, MRI, PET) para 3 canais de entrada (usando janelas de intensidade complementares para CT ou normalização por quantis para MRI/PET) para compatibilidade com o SAM2.
  2. Amostragem e Geração 2D: O SAM2 é aplicado em fatias 2D amostradas do volume, gerando centenas de propostas de máscaras baseadas em limites de intensidade e coerência morfológica.
  3. Propagação 3D: As máscaras 2D são propagadas através do volume 3D usando a capacidade de previsão de vídeo do SAM2, criando máscaras volumétricas.
• Qualidade: Embora essas máscaras sejam ruidosas e não tenham rótulos semânticos (não dizem "o que" é o órgão, apenas "onde" está uma região), elas capturam estruturas anatômicas e patológicas significativas.

B. Aprendizado de Auto-supervisão Guiado por Máscaras

O MASS formula o pré-treinamento como uma tarefa de segmentação em contexto (inspirada no trabalho Iris e UniverSeg):

• Entrada: Para cada passo de treinamento, são amostradas duas visões aumentadas de um volume e suas máscaras correspondentes: uma visão de referência $(x_s, y_s)$ e uma visão de consulta $(x_q, y_q)$.
• Mecanismo:
  1. O modelo codifica a imagem de referência e extrai um embedding de tarefa ($T$) que captura "o que" deve ser segmentado com base na máscara de referência.
  2. Este embedding guia o decodificador para prever a máscara correspondente na imagem de consulta ($x_q$).
• Objetivo de Aprendizado: O modelo é forçado a aprender a invariância semântica. Como as máscaras são as mesmas em ambas as visões (mesmo com transformações espaciais e de aparência), o modelo deve aprender a identidade fundamental da estrutura anatômica (forma, contexto espacial, relações) em vez de apenas correspondência de pixels ou textura. Isso permite que o modelo aprenda conceitos médicos gerais sem rótulos explícitos.

3. Principais Contribuições

1. Novo Framework de Auto-supervisão: O MASS é a primeira abordagem a usar segmentação em contexto com máscaras geradas automaticamente para criar representações médicas 3D generalizáveis, eliminando a dependência de anotações de especialistas.
2. Escalabilidade sem Anotação: O método escala de pré-treinamento em pequenos conjuntos de dados (20-200 exames) até grandes conjuntos multimodais (5.000 volumes de CT, MRI e PET), demonstrando melhoria consistente com o aumento da diversidade de dados.
3. Capacidade de Few-Shot (Poucos Exemplos): Modelos pré-treinados com MASS podem realizar segmentação em contexto com apenas 1 a 10 exemplos de referência, sem ajuste fino (finetuning). Com um ajuste fino mínimo (20-40% dos dados rotulados), eles igualam ou superam a supervisão total.
4. Generalização para Tarefas Não Vistas: O encoder congelado do MASS performou excepcionalmente bem em tarefas de classificação de patologias nunca vistas durante o pré-treinamento, igualando o desempenho de treinamento supervisionado completo com milhares de amostras.

4. Resultados Experimentais

Os experimentos foram conduzidos em diversos conjuntos de dados (BCV, AMOS, KiTS, BraTS, etc.) e modalidades.

• Segmentação Few-Shot:
  • O MASS-IC (inferência em contexto sem ajuste) superou todos os métodos de base (Scratch, MAE, SimCLR, Model Genesis) em configurações de 1-shot, alcançando desempenho competitivo sem nenhuma atualização de parâmetros.
  • O MASS-FT (com ajuste fino em poucos exemplos) atingiu o estado da arte (SOTA). Em alguns cenários (ex: BCV 1-shot), superou a melhor base em mais de 20 pontos de Dice.
  • Com apenas 20-40% dos dados de treinamento rotulados, o MASS igualou o desempenho da supervisão total.
• Pré-treinamento em Grande Escala:
  • Ao treinar em 5.000 volumes multimodais, o MASS superou métodos de auto-supervisão em grande escala (como OpenMind com 114k imagens) e métodos supervisionados (SuPreM, Iris) em tarefas de transferência, especialmente em dados fora da distribuição (out-of-distribution).
  • Destaque para a transferência entre modalidades (CT para MRI e PET), onde o MASS demonstrou robustez superior.
• Classificação (Encoder Congelado):
  • Em tarefas de classificação de patologias não vistas (ex: hemorragia intracraniana no RSNA ICH), o encoder do MASS congelado superou o treinamento do zero e igualou o desempenho de modelos supervisionados massivos, mesmo com apenas 5% dos dados de classificação disponíveis para ajuste do classificador.
• Análise de Falhas: O modelo teve desempenho inferior em segmentação de tumores sem ajuste fino (devido à alta variabilidade de localização entre pacientes), mas o ajuste fino corrigiu isso rapidamente, provando que as representações aprendidas contêm os recursos necessários.

5. Significado e Impacto

O trabalho MASS representa um avanço crucial para a criação de Modelos Fundamentais em Imagem Médica:

• Democratização: Permite o pré-treinamento em escala massiva sem o gargalo de anotações de especialistas, tornando viável o desenvolvimento de modelos fundamentais para instituições que não possuem recursos para anotação manual em larga escala.
• Generalização Real: Ao aprender a partir de máscaras agnósticas de classe, o modelo não fica preso a taxonomias pré-definidas, permitindo a adaptação a novas estruturas anatômicas e patologias que não foram vistas durante o treinamento.
• Eficiência de Dados: Demonstra que a diversidade de dados (anatomia e modalidade) é mais importante do que a quantidade bruta de dados, e que a supervisão fraca (máscaras automáticas) é suficiente para aprender representações ricas e transferíveis.
• Caminho Futuro: Abre a porta para modelos que podem ser adaptados rapidamente a novos cenários clínicos com apenas alguns exemplos, reduzindo drasticamente o tempo e o custo de implementação de IA na medicina.

Em resumo, o MASS prova que é possível aprender representações semânticas profundas e generalizáveis em imagens médicas 3D utilizando apenas dados não rotulados e ferramentas de segmentação automática, superando as limitações dos métodos atuais de auto-supervisão e supervisionados.