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

1. El Problema

El aprendizaje auto-supervisado (SSL) a gran escala ha transformado la visión por computadora y el procesamiento de lenguaje natural mediante modelos fundacionales (como GPT o DINO) que aprenden representaciones de propósito general a partir de datos no etiquetados. Sin embargo, en el campo de la imagen médica 3D, falta un enfoque equivalente.

Los métodos existentes presentan limitaciones críticas:

• Métodos de reconstrucción o contraste: Las técnicas actuales (como MAE o SimCLR) se basan en la reconstrucción de píxeles o objetivos de contraste de bajo nivel. Estas fallan en capturar la semántica anatómica crítica para el análisis médico, ya que en las imágenes médicas, casi todos los vóxeles tienen significado clínico (densidad ósea, textura de tejidos, patrones vasculares), a diferencia de las imágenes naturales donde el fondo es mayoritario.
• Dependencia de anotaciones: Los enfoques supervisados requieren anotaciones pixel a pixel realizadas por expertos, lo cual es costoso, lento y limita la escalabilidad a nuevos tipos de anatomía o patologías no predefinidas.
• Falta de generalización: Los modelos preentrenados con taxonomías fijas no pueden adaptarse fácilmente a estructuras anatómicas novedosas o variantes patológicas no vistas durante el entrenamiento.

2. Metodología: MASS (MAsk-guided Self-Supervised learning)

MASS introduce un nuevo marco de aprendizaje auto-supervisado que utiliza la segmentación en contexto (in-context segmentation) como tarea pretext, eliminando la necesidad de anotaciones humanas.

A. Generación de Máscaras sin Anotación (Annotation-Free Mask Generation)

En lugar de usar etiquetas de expertos, MASS utiliza herramientas de segmentación automática (específicamente SAM2, entrenado en imágenes naturales) para generar máscaras agnósticas a la clase:

1. Preprocesamiento: Se convierten los volúmenes médicos 3D (CT, MRI, PET) en entradas de 3 canales utilizando ventanas de intensidad complementarias para preservar el contraste de tejidos.
2. Generación 2D: Se muestrean rebanadas 2D y SAM2 genera cientos de propuestas de máscaras basadas en límites de intensidad y coherencia morfológica.
3. Propagación 3D: Estas máscaras 2D se propagan a través del volumen completo utilizando la capacidad de predicción de video de SAM2, creando miles de máscaras 3D por volumen.
   Nota: Estas máscaras son ruidosas y no tienen etiquetas semánticas (no dicen "hígado" o "tumor"), pero capturan regiones anatómicamente y patológicamente significativas.

B. Aprendizaje Guiado por Máscaras (Mask-Guided Learning)

El marco adopta un enfoque de segmentación en contexto (ICS):

• Entrada: Se toma una imagen 3D no etiquetada y sus máscaras generadas automáticamente.
• Vistas: Se crean dos vistas aumentadas: una referencia $(x_s, y_s)$ y una consulta $(x_q, y_q)$.
• Mecanismo:
  1. El modelo codifica la imagen de referencia y extrae una incrustación de tarea ($T$) que captura "qué" estructura se debe segmentar basándose en la máscara de referencia.
  2. Esta incrustación guía al decodificador para predecir la máscara correspondiente en la imagen de consulta ($x_q$).
• Objetivo de Aprendizaje: Al entrenar en miles de estas tareas de segmentación en contexto, el modelo aprende a identificar qué define semánticamente una estructura médica (combinación de apariencia, forma, contexto espacial y relaciones anatómicas) sin necesidad de saber el nombre de la estructura. El modelo aprende la invariancia semántica frente a variaciones de apariencia y posición.

3. Contribuciones Clave

1. Nuevo Marco Auto-supervisado para 3D: MASS es el primer enfoque que utiliza la segmentación en contexto con máscaras generadas automáticamente para aprender representaciones ricas semánticamente en imágenes médicas 3D sin anotaciones humanas.
2. Escalabilidad sin Anotaciones: Demuestra eficacia desde preentrenamiento a pequeña escala (20-200 escaneos) hasta grandes conjuntos de datos multimodales (5,000 volúmenes de CT, MRI y PET).
3. Capacidades de Few-Shot (Pocos Ejemplos):
   • Inferencia en contexto: El modelo preentrenado puede segmentar nuevas estructuras directamente sin ajuste (finetuning), superando a métodos base en configuraciones de 1-shot.
   • Ajuste eficiente: Con solo el 20-40% de datos etiquetados para ajuste fino, MASS iguala el rendimiento de la supervisión completa y supera a otros métodos auto-supervisados en más de 20 puntos de Dice en escenarios de datos extremadamente limitados.
4. Generalización a Tareas No Vistas: En tareas de clasificación de patologías no vistas (usando el codificador congelado), MASS iguala o supera el entrenamiento supervisado completo con miles de muestras, demostrando que aprende conocimiento médico transferible más allá de la segmentación.

4. Resultados Experimentales

Los experimentos se realizaron en múltiples conjuntos de datos (BCV, AMOS, KiTS, BraTS, etc.) y modalidades (CT, MRI, PET).

• Segmentación Few-Shot:
  • En el conjunto de datos BCV (órganos abdominales), MASS con ajuste fino (10 shots) alcanzó un 84.2% de Dice, igualando la supervisión completa (85.5%) y superando a la mejor línea base auto-supervisada (OM-MG) en más de 20 puntos.
  • MASS-IC (sin ajuste) superó a todos los métodos base en configuraciones de 1-shot en anatomías bien definidas.
• Preentrenamiento a Gran Escala:
  • El preentrenamiento en 5,000 volúmenes multimodales mejoró el rendimiento en comparación con el preentrenamiento en un solo conjunto de datos, validando la escalabilidad.
  • MASS superó a métodos supervisados masivos (como SuPreM) en conjuntos de datos fuera de distribución (ej. Pelvis, BraTS), demostrando una mejor generalización cruzada.
• Clasificación (Codificador Congelado):
  • En la detección de hemorragias intracraneales (RSNA ICH) con solo el 5% de datos de entrenamiento, MASS (75.4% AUC) superó al entrenamiento desde cero (72.8%).
  • En detección de trauma, MASS superó significativamente al entrenamiento completo con solo el 30% de los datos.
• Análisis de Fallos: Se observó que MASS tiene dificultades iniciales con patologías muy variables (como tumores) en modo "sin ajuste" debido a la falta de correspondencia inter-paciente, pero el ajuste fino con pocos ejemplos corrige esto rápidamente, confirmando que el modelo sí ha aprendido características discriminativas de las patologías.

5. Significado e Impacto

El trabajo MASS representa un cambio de paradigma en el desarrollo de modelos fundacionales para imágenes médicas:

• Democratización: Permite el preentrenamiento de modelos a gran escala en instituciones que no tienen recursos para generar miles de anotaciones de expertos.
• Generalización Abierta: Al no estar limitado por taxonomías predefinidas, MASS puede aprender sobre cualquier estructura anatómica o patológica presente en los datos, abriendo la puerta a modelos fundacionales verdaderamente generales.
• Eficiencia de Datos: Demuestra que la supervisión débil (máscaras automáticas ruidosas) es suficiente para aprender representaciones semánticas robustas, reduciendo drásticamente la dependencia de datos etiquetados para tareas de alto rendimiento.

En resumen, MASS establece una ruta viable hacia modelos fundacionales de imágenes médicas 3D que son escalables, generalizables y no dependen de la costosa anotación experta, utilizando la segmentación en contexto guiada por máscaras automáticas como motor de aprendizaje.