Robust Multi-Source Covid-19 Detection in CT Images

1. El Problema

La detección de COVID-19 mediante modelos de aprendizaje profundo en tomografías computarizadas (TC) de tórax suele funcionar bien cuando los datos de entrenamiento y prueba provienen de la misma institución. Sin embargo, estos modelos sufren un desplazamiento de dominio (domain shift) significativo cuando se aplican a datos de múltiples centros médicos.

Las causas principales de este problema son:

• Heterogeneidad de los datos: Diferencias en los escáneres, protocolos de imagen, espesor de las capas y poblaciones de pacientes entre hospitales.
• Desbalance de fuentes: En el conjunto de datos utilizado, tres centros aportan aproximadamente 330 escaneos de entrenamiento cada uno, mientras que el cuarto aporta solo 234.
• Sesgo en el aprendizaje: Los métodos existentes optimizan únicamente la clasificación de COVID-19 (positivo/negativo). Esto hace que el modelo aprenda representaciones sesgadas hacia los centros que aportan más datos, utilizando "artefactos de adquisición" (como el brillo o el kernel de reconstrucción) como señales engañosas en lugar de características patológicas reales. Al evaluar el modelo por fuente, el rendimiento cae drásticamente en los centros minoritarios o no vistos.

2. Metodología

Los autores proponen un enfoque de aprendizaje multi-tarea diseñado para forzar al modelo a aprender características invariantes al origen de los datos. La arquitectura consta de tres componentes principales:

A. Preprocesamiento (SSFL + KDS)

Siguiendo el marco SSFL (Slice Sampling for Federated Learning) con muestreo basado en densidad de kernel (KDS):

1. Extracción de pulmones: Se aplica un filtro mínimo, umbralización y relleno morfológico para aislar la región pulmonar y eliminar ruido.
2. Muestreo de cortes (KDS): Dado que la longitud de los escaneos varía enormemente (de 50 a 700 cortes), el método KDS selecciona 8 cortes representativos por escaneo. En lugar de un submuestreo uniforme, utiliza una estimación de densidad de kernel gaussiana para seleccionar cortes que capturen la anatomía completa (especialmente las zonas de cambio rápido como la base y el hilio), evitando perder regiones diagnósticamente informativas.

B. Arquitectura del Modelo (Multi-Head)

• Backbone Compartido: Se utiliza EfficientNet-B7 preentrenado en ImageNet. Cada uno de los 8 cortes seleccionados pasa independientemente por esta red, generando 8 vectores de características de 2560 dimensiones.
• Agregación: Los vectores se promedian (mean pooling) elemento a elemento para obtener una representación única a nivel de escaneo.
• Dos Cabezas de Clasificación:
  1. Cabeza de Detección de COVID-19: Predice la probabilidad de infección (positivo/negativo) usando una capa lineal con activación sigmoide.
  2. Cabeza de Identificación de Fuente (Auxiliar): Predice de cuál de los 4 centros proviene el escaneo. Esta cabeza se descarta durante la inferencia; su único propósito es regularizar el aprendizaje de características durante el entrenamiento.

C. Función de Pérdida Ajustada por Logits (Logit-Adjusted Loss)

Para abordar el desbalance en la cantidad de datos por centro, no se utiliza una pérdida de entropía cruzada estándar para la identificación de la fuente. En su lugar, se aplica una pérdida de entropía cruzada ajustada por logits:

• Se añade un término de corrección basado en la frecuencia logarítmica de cada fuente ($\log(p(d))$) a los logits antes de la función softmax.
• Objetivo: Esto reduce la confianza requerida para predecir fuentes minoritarias y aumenta la penalización por ignorarlas, evitando que el gradiente de la tarea auxiliar sesgue el backbone hacia los centros mayoritarios.
• Función de pérdida total: $\mathcal{L} = \mathcal{L}_{CE} + \gamma \cdot \mathcal{L}_{LA}$, donde $\gamma$ controla el peso de la tarea auxiliar.

3. Contribuciones Clave

1. Formulación Multi-Tarea: Plantear la detección de COVID-19 multi-fuente como un problema de aprendizaje multi-tarea, emparejando el objetivo diagnóstico con la identificación de la fuente. Esto obliga al extractor de características a desentrelazar la señal de la enfermedad de los artefactos específicos de cada institución.
2. Uso de Logit-Adjustment: Demostrar que aplicar un ajuste de logits en la tarea de identificación de fuente es crucial. Sin este ajuste, la tarea multi-tarea con pérdida estándar (Cross-Entropy) empeora el rendimiento en comparación con un modelo de tarea única, ya que refuerza el sesgo hacia los centros mayoritarios.
3. Mejora del Rendimiento: Lograr un aumento de 1.86 puntos porcentuales en la puntuación final sobre la línea base de tarea única y 1.98 puntos sobre la mejor variante multi-tarea estándar.

4. Resultados

El modelo se evaluó en un conjunto de validación de 308 escaneos de 4 centros médicos distintos.

• Puntuación Final (Final Score): El método propuesto (Multi-task + Logit-Adjusted) alcanzó una puntuación de 0.8194, superando a la línea base de tarea única (0.8008) y a la variante multi-tarea con pérdida estándar (máximo 0.7996).
• Métricas Globales:
  • F1 Score: 0.9098.
  • AUC-ROC: 0.9647.
  • Sensibilidad: 0.9062.
  • Especificidad: 0.9389.
• Análisis por Fuente:
  • Las mayores mejoras se observaron en las fuentes 0 y 2 (donde la línea base tenía más margen de mejora), con aumentos de 3.34 y 4.87 puntos porcentuales respectivamente en el F1 promedio por fuente.
  • La fuente 3 tuvo un rendimiento bajo en la puntuación final debido a la ausencia de casos positivos en su partición de validación (un problema estructural del dataset, no del modelo), aunque el modelo logró una alta especificidad (0.9888) en los casos negativos de este centro.
• Hiperparámetro $\gamma$: La sensibilidad mostró que un valor de $\gamma = 0.5$ es óptimo. Valores más bajos no regularizan lo suficiente, y valores más altos (ej. 1.0) hacen que el modelo se centre demasiado en identificar la fuente, degradando la capacidad de detección de la enfermedad.

5. Significado e Impacto

Este trabajo es significativo porque aborda directamente el problema de la generalización en entornos clínicos reales, donde los datos provienen de múltiples instituciones con características dispares.

• Justicia y Equidad: Al corregir el desbalance de datos mediante el ajuste de logits, el modelo evita que las instituciones con menos datos queden subrepresentadas o que el modelo falle sistemáticamente en ellas.
• Robustez: Demuestra que la simple adición de una tarea auxiliar no es suficiente; la diseño de la función de pérdida es crítica. Si la tarea auxiliar no se equilibra correctamente, puede ser contraproducente.
• Aplicabilidad Clínica: El enfoque permite crear sistemas de cribado automatizado más fiables que no dependen de que el hospital de destino sea idéntico al hospital de entrenamiento, un requisito esencial para la implementación a gran escala de la IA en medicina.

El código fuente del proyecto está disponible públicamente, facilitando la reproducción y el avance en la investigación de detección de enfermedades basada en imágenes médicas multi-fuente.