CoRe: Joint Optimization with Contrastive Learning for Medical Image Registration

Resumen Técnico: CoRe - Optimización Conjunta con Aprendizaje Contrastivo para el Registro de Imágenes Médicas

1. Planteamiento del Problema

El registro de imágenes médicas es fundamental para establecer correspondencias anatómicas y semánticas densas entre imágenes (por ejemplo, en diferentes momentos, sujetos o modalidades). Sin embargo, el proceso enfrenta dos desafíos principales:

• Variabilidad en la apariencia de la imagen: Diferencias en protocolos de adquisición, hardware o modalidades que hacen que las comparaciones píxel a píxel sean poco fiables.
• Deformaciones anatómicas no lineales: Cambios estructurales complejos debidos a la progresión de enfermedades, intervenciones quirúrgicas o diferencias inter-sujeto.

Los métodos existentes a menudo dependen de métricas de información mutua (sensibles al ruido), descriptores de características manuales (limitados en generalización) o aprendizaje supervisado (que requiere datos alineados costosos). Además, enfoques híbridos recientes suelen preentrenar extractores de características de forma independiente, sin optimizarlos conjuntamente con la tarea de registro, lo que puede resultar en representaciones no óptimas para la estimación de campos de deformación.

2. Metodología: El Marco CoRe

El artículo propone CoRe (Contrastive learning for medical image Registration), un marco unificado que integra el aprendizaje contrastivo auto-supervisado directamente en la tubería de registro deformable mediante una optimización conjunta.

• Arquitectura: El sistema consta de un extractor de características convolucional ($G$), una cabeza de proyección ($P$) y un módulo de optimización diferenciable ($H$).
• Paradigma de Extracción Pre-Deformación: Las características se extraen de las imágenes fija y móvil antes de estimar el campo de deformación.
• Optimización Conjunta: A diferencia de métodos que preentrenan el extractor y luego lo congelan, CoRe optimiza simultáneamente $G$ y el módulo de registro minimizando una función de pérdida combinada:
  $$L = L_{reg} + \alpha \cdot L_c$$
  Donde:
  • $L_{reg}$ (Pérdida de Registro): Utiliza un esquema de auto-entrenamiento con pseudo-etiquetas. Minimiza el error cuadrático medio (MSE) entre el campo de desplazamiento predicho y un campo de desplazamiento refinado (pseudo-ground truth) generado iterativamente.
  • $L_c$ (Pérdida Contrastiva de Equivariancia): Se basa en una restricción de equivariancia geométrica. Aplica transformaciones afines aleatorias a las imágenes y exige que las características extraídas de la imagen transformada sean consistentes con las características transformadas de la imagen original. Se implementa mediante la pérdida InfoNCE, asegurando que las representaciones de las mismas ubicaciones anatómicas permanezcan robustas ante deformaciones geométricas.
• Ventaja Clave: No requiere datos de entrenamiento con correspondencias conocidas ni pares multimodales alineados, utilizando únicamente aumentaciones geométricas de la misma imagen para la señal auto-supervisada.

3. Contribuciones Clave

1. Marco de Optimización Conjunta: Propone integrar el aprendizaje contrastivo auto-supervisado directamente en el flujo de registro, alineando el aprendizaje de características con los requisitos de la tarea de registro descendente.
2. Sinergia de Objetivos: Demuestra que la optimización simultánea de los objetivos de contraste y registro produce resultados superiores a entrenar con cualquiera de los objetivos por separado, destacando la sinergia entre el aprendizaje de representaciones y el registro.
3. Robustez y Rendimiento: Valida el método en conjuntos de datos abdominales y torácicos (intra e inter-paciente), logrando un rendimiento superior frente a enfoques competitivos sin necesidad de preentrenamiento separado.

4. Resultados Experimentales

El método fue evaluado en dos conjuntos de datos desafiantes:

• AbdomenCT (Inter-paciente): 30 escaneos 3D con 13 estructuras anatómicas etiquetadas.
• RadChestCT (Intra-paciente): 371 pares de escaneos longitudinales de tórax.

Métricas de Evaluación: Coeficiente de DICE (precisión de alineación), Desviación Estándar del Logaritmo del Jacobiano (SDLogJ, para suavidad del campo) y tiempo de inferencia.

Hallazgos Cuantitativos (Tabla 1):

• CoRe logró el mayor puntaje DICE en ambos conjuntos de datos:
  • AbdomenCT: 52.1 ± 7.5 (superando a SAMConvex con 49.2 y RegCyc con 51.1).
  • RadChestCT: 89.4 ± 6.6 (superando a DEEDs con 88.5 y SAMConvex con 88.1).
• La suavidad de los campos de deformación (SDLogJ) se mantuvo comparable a los métodos más avanzados.
• Estudio de Ablación (Tabla 2): Confirmó que la optimización conjunta ($L_{reg} + L_c$) supera significativamente a los enfoques donde se entrena el extractor solo con pérdida de registro o solo con pérdida contrastiva.
• Ablación de Aumentaciones (Tabla 3): Se descubrió que, en datos de CT monomodales, las aumentaciones de intensidad no lineales no aportan beneficios adicionales e incluso pueden ser perjudiciales; la equivariancia geométrica es el factor crítico para la robustez.

5. Significado e Impacto

El trabajo de CoRe representa un avance significativo en el registro de imágenes médicas al demostrar que no es necesario separar las fases de aprendizaje de características y registro. Al integrar la restricción de equivariancia geométrica directamente en el entrenamiento del registro, el modelo aprende representaciones que son intrínsecamente robustas a las deformaciones de los tejidos y semánticamente coherentes para la optimización.

Esto permite lograr un rendimiento de vanguardia (State-of-the-Art) en tareas complejas de registro inter-paciente sin depender de anotaciones costosas o preentrenamientos externos, ofreciendo un marco eficiente y preciso para la alineación de imágenes médicas en entornos clínicos y de investigación.