Unsupervised Source-Free Ranking of Biomedical Segmentation Models Under Distribution Shift

Resumen Técnico

1. El Problema

En el campo de la imagen biomédica, la anotación de datos es costosa y laboriosa, lo que limita el entrenamiento de modelos desde cero. Como solución, se recurre al uso de modelos preentrenados (reutilización de modelos) disponibles en repositorios como el BioImage Model Zoo o desafíos públicos. Sin embargo, surge un desafío crítico: cómo seleccionar el modelo más adecuado para un nuevo conjunto de datos objetivo sin tener acceso a etiquetas de ground truth en ese dominio objetivo.

Los métodos existentes para estimar la transferibilidad suelen:

• Requerir datos etiquetados en el dominio objetivo.
• Depender del acceso al espacio de características (requiriendo acceso al modelo, lo que no es posible con modelos "caja negra").
• Estar diseñados principalmente para clasificación, fallando al aplicarse a segmentación (especialmente segmentación de instancias).
• No ser agnósticos a la arquitectura o al entrenamiento.

El objetivo es desarrollar un marco no supervisado, libre de fuentes (source-free) y agnóstico al modelo que pueda clasificar modelos de segmentación semántica y de instancias en repositorios heterogéneos.

2. Metodología: CMR (Consistency Model Ranking)

Los autores proponen CMR, un método que evalúa la calidad de un modelo basándose en la consistencia de sus predicciones bajo perturbaciones, sin necesidad de etiquetas.

• Principio Fundamental: La idea central es que un modelo que se adapta bien a un dominio objetivo tendrá fronteras de decisión robustas. Si se aplican pequeñas perturbaciones a la entrada o a las características del modelo, las predicciones de un buen modelo deberían permanecer estables (consistentes). Un modelo con un rendimiento pobre o con fronteras de decisión inestables sufrirá cambios drásticos en sus predicciones (flips de clase) ante las mismas perturbaciones.
• Tipos de Perturbaciones:
  • Espacio de Entrada: Ruido gaussiano, cambios de brillo, contraste y corrección gamma.
  • Espacio de Características: Aplicación de Dropout en tiempo de prueba (Test-Time Dropout - TTD) en capas específicas (cuello de botella, conexiones de salto o todas las capas).
• Métricas de Consistencia:
  • Para Segmentación Semántica:
    • CMR-EI (Soft): Basado en la Invarianza Efectiva, utiliza la probabilidad softmax y la consistencia de etiquetas.
    • CMR-NHD (Hard): Basado en la distancia de Hamming normalizada (equivalente al IoU) contando los cambios de píxeles en el fondo restringido.
  • Para Segmentación de Instancias:
    • Dado que las etiquetas de instancias son permutables y no tienen correspondencia fija en el espacio de características, no se pueden comparar píxel a píxel directamente.
    • Se utiliza el Adapted Rand Score (ARS) restringido al primer plano. Este mide la concordancia entre pares de píxeles asignados a la misma instancia en la predicción original y la perturbada.
• Proceso de Clasificación:
  1. Se aplica el modelo al conjunto de datos objetivo no etiquetado.
  2. Se generan predicciones perturbadas (N perturbaciones).
  3. Se calcula la puntuación de consistencia promedio (usando la mediana sobre las imágenes para robustez).
  4. Los modelos se ordenan de mayor a menor puntuación de consistencia.

3. Contribuciones Clave

1. Primer método unificado: Es la primera metodología que clasifica tanto modelos de segmentación semántica como de instancias de manera no supervisada y libre de fuentes.
2. Compatibilidad con Caja Negra: No requiere acceso a los pesos del modelo, ni a sus características internas, ni a datos de entrenamiento. Funciona solo con las predicciones de salida, lo que lo hace compatible con APIs de inferencia y contenedores.
3. Evaluación en Escenarios Realistas: Se valida en una amplia gama de tareas biomédicas (núcleos, células, mitocondrias) en 2D y 3D, incluyendo microscopía de luz y electrónica, y volúmenes CBCT.
4. Validación en Desafíos Públicos: El método logró replicar con éxito la clasificación de los participantes del desafío ToothFairy (segmentación 3D de mandíbula) sin acceso a las etiquetas de prueba, demostrando su utilidad práctica.

4. Resultados Experimentales

Los autores evaluaron CMR frente a métricas de transferibilidad supervisadas (como LEEP, LogME) y basadas en UDA (Transfer Score), así como métodos de calidad de control (SEG, NuNo).

• Correlación con el Rendimiento Real:
  • En segmentación semántica (mitocondrias y núcleos), CMR mostró una fuerte correlación (Pearson > 0.9, Kendall > 0.7) con el rendimiento real (F1-score) en el dominio objetivo.
  • Las métricas supervisadas (que requieren etiquetas) y las basadas en características fallaron consistentemente en la segmentación biomédica, confirmando que las técnicas derivadas de la clasificación no son directamente transferibles.
• Segmentación de Instancias:
  • CMR-ARS superó significativamente al método SEG (basado en ensembles y detección de centroides), que es sensible a parámetros y no captura la calidad a nivel de píxel.
  • Logró clasificar correctamente modelos heterogéneos (desde U-Nets especializados hasta modelos generalistas como Cellpose-SAM y µSAM).
• Adaptación de Dominio No Supervisada (UDA):
  • El método también fue efectivo para clasificar modelos después de aplicar técnicas de UDA (como Mean Teacher o AdaBN), superando a la métrica Transfer Score (SOTA actual para validación de UDA).
• Robustez:
  • El método es estable ante un amplio rango de fuerzas de perturbación (ruido, dropout), aunque se recomienda evitar perturbaciones extremas que rompan completamente la señal.

5. Significado e Impacto

Este trabajo aborda una brecha crítica en la adopción de IA en biomedicina: la selección de modelos en repositorios.

• Eficiencia: Permite a los investigadores identificar el mejor modelo para sus datos sin gastar semanas en anotación o fine-tuning.
• Generalización: Al ser agnóstico a la arquitectura y al entrenamiento, facilita la integración de modelos de diferentes orígenes (modelos generalistas vs. especializados).
• Aplicabilidad: Es especialmente valioso en escenarios donde las etiquetas son escasas o costosas (como en microscopía electrónica 3D o tomografía computarizada), permitiendo una reutilización escalable de modelos preentrenados.

En conclusión, CMR proporciona una herramienta práctica, eficiente y teóricamente fundamentada para la gestión de modelos en el ecosistema emergente de repositorios de modelos biomédicos, garantizando que los modelos seleccionados sean robustos y adecuados para el dominio objetivo específico.