Multi-Dataset Cross-Domain Knowledge Distillation for Unified Medical Image Segmentation, Classification, and Detection

Resumen Técnico: Destilación de Conocimiento Transdominio Multi-Conjunto de Datos para Segmentación, Clasificación y Detección Unificada de Imágenes Médicas

Enunciado del Problema
El aprendizaje profundo ha revolucionado el análisis de imágenes médicas, sin embargo, los modelos entrenados en conjuntos de datos individuales a menudo fallan al generalizar a nuevos entornos clínicos debido a desplazamientos de dominio causados por variaciones en los protocolos de imagen, el hardware de los escáneres y las poblaciones de pacientes. Aunque el aprendizaje por transferencia y la augmentación de datos ofrecen soluciones parciales, a menudo dependen de un pre-entrenamiento de fuente única o de datos sintéticos que no logran capturar detalles anatómicos de gran detalle y la variabilidad entre conjuntos de datos. Además, los enfoques existentes de destilación de conocimiento (KD) suelen utilizar un único profesor sobre datos homogéneos para una sola tarea (por ejemplo, segmentación), descuidando la naturaleza diversa y multi-tarea de los flujos de trabajo clínicos que abarcan segmentación, clasificación y detección de objetos a través de modalidades heterogéneas (RM, TC, Rayos X).

Metodología
Los autores proponen un marco unificado, agnóstico a la tarea, basado en profesor-alumno diseñado para agregar conocimiento de múltiples conjuntos de datos de imágenes médicas heterogéneos. El marco opera en tres etapas secuenciales:

1. Entrenamiento del Profesor Adversarial de Dominio:

   • Se entrena un profesor objetivo exclusivamente en el conjunto de datos objetivo utilizando una pérdida específica de la tarea (por ejemplo, pérdida Dice para segmentación).
   • Los profesores fuente se entrenan en sus respectivos conjuntos de datos fuente, pero aumentados con un discriminador de dominio. Utilizando un objetivo adversarial de dominio, estos codificadores se ven obligados a producir características invariantes al dominio que confunden al discriminador, alineando las distribuciones de características de la fuente y el objetivo. Esta etapa soporta tanto la adaptación supervisada (usando datos objetivos etiquetados) como la no supervisada.

2. Construcción Conjunta del Profesor mediante Fusión por Atención Cruzada:

   • Se construye un profesor conjunto ($\mathcal{T}^*$) fusionando las características de codificador y cuello de botella de múltiples niveles del profesor objetivo y de todos los profesores fuente.
   • Esta fusión utiliza mecanismos de atención cruzada para agregar representaciones complementarias de diferentes dominios. Las capas de codificador y cuello de botella de todos los profesores constituyentes se congelan durante este proceso para preservar el conocimiento aprendido.
   • Se adjunta una nueva cabeza específica de la tarea (decodificador, clasificador o cabeza de detección) al esqueleto fusionado y se entrena desde cero en el conjunto de datos objetivo. Esto permite que el profesor conjunto se especialice en la tarea objetivo mientras retiene características invariantes al dominio.

3. Destilación Multi-Nivel Guiada por Currículo:

   • Se entrena un modelo alumno compacto en el conjunto de datos objetivo para heredar conocimiento del profesor conjunto.
   • El objetivo de entrenamiento combina una pérdida de tarea principal con múltiples términos de destilación: alineación de características (MSE), similitud coseno, pérdida contrastiva (divergencia KL en distribuciones de similitud) y extensiones específicas de la tarea (por ejemplo, transferencia de atención, destilación de logits, alineación de características de ROI para detección).
   • Un factor de currículo ($c_f$) programa dinámicamente el peso de las pérdidas de destilación. El entrenamiento comienza con un enfoque en la pérdida de tarea supervisada, aumentando gradualmente la influencia de los términos de destilación para refinar las representaciones internas del alumno sin desestabilizar el aprendizaje inicial.

El marco está diseñado para ser flexible, soportando arquitecturas UNet, TResUNet, MedViT, EfficientNet, Faster R-CNN y RF-DETR en tareas de segmentación, clasificación y detección.

Contribuciones Clave

1. Marco Unificado: Una arquitectura única y escalable capaz de manejar segmentación, clasificación a nivel de imagen y detección de objetos a través de conjuntos de datos médicos heterogéneos.
2. Arquitectura de Profesor Conjunto: Un método novedoso para fusionar características de múltiples niveles de múltiples profesores fuente mediante atención cruzada, permitiendo una transferencia robusta de representaciones transdominio sin requerir un único modelo pre-entrenado monolítico.
3. Destilación Basada en Currículo: Una estrategia de destilación multi-nivel que combina pérdidas de segmentación, contrastivas, de alineación y de similitud, regida por un programa de currículo para garantizar una optimización estable.
4. Evaluación Exhaustiva: Experimentos extensos a través de tres tareas, tres modalidades (RM, TC, Rayos X) y numerosos conjuntos de datos (incluyendo BrainMetShare, ISLES, BraTS, LiTS, KiTS, COVIDx-CXR, OASIS, ADNI y DeepLesion).

Resultados Experimentales
El método propuesto se evaluó contra dos líneas base: modelos específicos de conjunto de datos (entrenados independientemente en cada conjunto de datos) y modelos multi-cabeza multi-conjunto de datos (esqueleto compartido con cabezas específicas de conjunto de datos).

• Segmentación: En conjuntos de datos de RM y TC, los modelos alumno superaron consistentemente a ambas líneas base. Para tareas basadas en TC, el alumno logró mejoras promedio en IoU de ~2.5% sobre las líneas base específicas de conjunto de datos y ~4.8% sobre los modelos multi-cabeza, con reducciones significativas en la Distancia de Hausdorff (HD95), lo que indica predicciones de límites más nítidas.
• Clasificación: En tareas de clasificación de enfermedades pulmonares y demencia, los modelos alumnos destilados lograron una mayor precisión que las líneas base específicas de conjunto de datos y multi-cabeza. La destilación tardía-$k$ (transferencia de conocimiento desde capas más profundas) generalmente produjo resultados superiores.
• Detección de Objetos: En conjuntos de datos de TC pulmonar, el marco mejoró la Precisión Promedio Media (mAP). El profesor conjunto a menudo logró el rendimiento más alto, mientras que el alumno destilado retuvo la mayor parte de estas ganancias, ofreciendo un modelo más compacto adecuado para implementación.
• Estudios de Ablación: Los resultados confirmaron que (i) la adaptación supervisada del objetivo mejora la precisión de los límites; (ii) el rendimiento escala con el número de conjuntos de datos fuente; y (iii) la extracción de características multi-nivel (codificador + cuello de botella) produce mejores resultados que la extracción de un solo nivel.
• Visualización: Las visualizaciones t-SNE demostraron que los modelos alumnos destilados aprenden agrupaciones de características más compactas y mejor separadas en comparación con las líneas base específicas de conjunto de datos en las tres tareas.

Significado y Afirmaciones
El artículo afirma que la destilación de conocimiento multi-conjunto de datos sirve como un enfoque escalable y agnóstico a la tarea para mejorar el rendimiento en dominios heterogéneos de imágenes médicas. Al agregar representaciones invariantes al dominio de diversas fuentes y destilarlas en un alumno compacto, el marco aborda los desafíos de la escasez de datos, los costos de anotación y los desplazamientos de dominio. Los autores enfatizan que su enfoque produce mejoras consistentes sobre líneas base sólidas, demostrando una mayor robustez ante desplazamientos distribucionales y una generalización superior. El trabajo se posiciona como una extensión sustancial de investigaciones anteriores centradas en la segmentación, ahora validada en todo el espectro de tareas de análisis de imágenes médicas.