HypCBC: Domain-Invariant Hyperbolic Cross-Branch Consistency for Generalizable Medical Image Analysis

Resumen Técnico: HypCBC

Planteamiento del Problema
Las redes neuronales profundas enfrentan desafíos significativos para generalizar más allá de las distribuciones de entrenamiento, un problema crítico en el análisis de imágenes médicas donde la escasez de datos y los cambios de covariable (derivados de diferentes hardware, protocolos de imagen y poblaciones de pacientes) dificultan un despliegue clínico fiable. Los remedios existentes suelen depender del aprendizaje de representaciones euclidianas, lo cual asume una geometría plana que no logra capturar las estructuras compleas y jerárquicas inherentes a los datos clínicos. Además, las redes hiperbólicas de extremo a extremo (end-to-end), aunque teóricamente adecuadas para datos jerárquicos, históricamente han sufrido de inestabilidad numérica y una reducción en la eficiencia de entrenamiento en grandes conjuntos de datos.

Metodología
Los autores proponen HypCBC (Consistencia de Doble Rama Hiperbólica), un marco que aprovecha las variedades hiperbólicas para una generalización de dominio robusta sin requerir etiquetas de dominio. El método opera sobre un backbone euclidiano preentrenado y congelado (por ejemplo, ViT) e introduce una estrategia de proyección hiperbólica ligera:

1. Proyección Hiperbólica: En lugar de entrenar una red hiperbólica de extremo a extremo, el método proyecta las incrustaciones (embeddings) euclidianas de un backbone congelado en un modelo de bola de Poincaré ($D^d_c$) utilizando un mapa exponencial. Esto combina la estabilidad de las características euclidianas preentrenadas con las ventajas geométricas del espacio hiperbólico (curvatura negativa constante).
2. Arquitectura de Dos Ramas: La arquitectura emplea dos ramas paralelas:
   • Rama de Alta Dimensión ($d_1=128$): Captura características finas y específicas de la clase para la inferencia.
   • Rama de Baja Dimensión ($d_2=2$): Actúa como un cuello de botella de información. El análisis empírico sugiere que este espacio de baja dimensión impone naturalmente la invariancia de dominio al descartar el ruido específico del dominio mientras retiene la información discriminativa de la etiqueta.
3. Consistencia de Doble Rama (CBC): Se introduce una novedosa restricción de regularización no supervisada. Los logits de la rama de baja dimensión (invariante al dominio) supervisan la rama de alta dimensión mediante una pérdida de divergencia de Kullback–Leibler (KL). Esto transfiere conocimiento agnóstico al dominio a la rama principal de inferencia sin depender del entrenamiento adversarial o de etiquetas de dominio explícitas.
4. Objetivo de Entrenamiento: La pérdida total combina las pérdidas de entropía cruzada para ambas ramas y la pérdida de consistencia, ponderada por un hiperparámetro $\lambda$.

Contribuciones Clave

• Validación Exhaustiva del Aprendizaje Hiperbólico: El artículo proporciona la primera validación extensa del aprendizaje de representaciones hiperbólicas para imágenes médicas, demostrando mejoras estadísticamente significativas sobre las líneas base euclidianas en once conjuntos de datos de distribución interna (in-distribution) que abarcan nueve modalidades de imagen, tamaños de muestra desde $10^2$ hasta $10^5$ y diversas granularidades de etiquetas.
• Generalización de Dominio No Supervisada: La introducción de HypCBC, una estrategia de doble rama con consistencia invariante al dominio, que supera a los métodos euclidianos de vanguardia (incluyendo enfoques basados en aumentación y minimización de riesgo invariante) en tres benchmarks de generalización de dominio.
• Perspectivas de Ablación: El estudio revela que una rama de baja dimensión de 2D equilibra eficazmente la invariancia de dominio con la discriminación de etiquetas. Además, demuestra que las ganancias de robustez se deben a las propiedades geométricas de la variedad hiperbólica y a la restricción de consistencia, y no simplemente a un aumento en el recuento de parámetros.

Resultados Experimentales

• Rendimiento de Distribución Interna (ID): A través de tres modelos ViT (ViT-S, DeiT3-S, DINOv2-S) y once conjuntos de datos médicos, los clasificadores hiperbólicos (HypERM) superaron consistentemente a los clasificadores euclidianos (ERM). Las ganancias fueron estadísticamente significativas ($p < 0.05$), con mejoras notables en los conjuntos de datos de OCT (+2.70%), Path (+1.48%) y Tissue (+1.36%) utilizando DINOv2.
• Generalización de Dominio (DG): En tres benchmarks—Fitzpatrick17k (dermatología), Camelyon17-WILDS (histopatología) y una configuración de imagen de retina entre conjuntos de datos—HypCBC logró el mejor rendimiento.
  • Superó al mejor método euclidiano (AugMix) por un +0.68% en el conjunto de prueba de Retina.
  • Logró una mejora promedio de +2.1% AUC en los tres benchmarks de DG en comparación con las líneas base euclidianas.
  • A diferencia de métodos como VREx, HypCBC logra estas ganancias de manera totalmente no supervisada, sin acceso a etiquetas de dominio durante el entrenamiento.
• Hallazgos de Ablación:
  • Dimensionalidad: Un cuello de botella de 2D proporcionó la mayor invariancia de dominio (el AUC de clasificación de dominio más bajo) manteniendo una alta discriminación de etiquetas. Las incrustaciones hiperbólicas mantuvieron las ganancias de precisión de etiquetas hasta $d=128$, mientras que las incrustaciones euclidianas se estancaron.
  • Consistencia: La regularización de consistencia de doble rama produjo ganancias de AUC positivas en todos los conjuntos de datos y tamaños de cuello de botella en el espacio hiperbólico, mientras que la consistencia euclidiana produjo resultados desiguales o negativos.

Significancia y Reivindicaciones
Los autores afirman que HypCBC ofrece una "alternativa sencilla pero poderosa" para construir sistemas de IA médica robustos. El trabajo destaca que la curvatura negativa constante de las variedades hiperbólicas refleja naturalmente las estructuras jerárquicas de los datos clínicos, una característica que la geometría euclidiana plana no puede capturar. Al proyectar las características euclidianas congeladas en un espacio hiperbólico y aplicar la consistencia de doble rama, el método logra un entrenamiento estable y una generalización superior a dominios no vistos. El artículo enfatiza que estas mejoras no son artefactos de un aumento en la capacidad del modelo, sino que son intrínsecas a la geometría hiperbólica y a la restricción de consistencia propuesta.

Limitaciones
Los autores señalan que la implementación actual utiliza un parámetro de curvatura estático ($c=1.0$) y un backbone congelado. Sugieren que una curvatura aprendible o el ajuste fino (fine-tuning) del backbone podrían mejorar aún más el rendimiento, pero a costa de la complejidad. Además, el método se evalúa actualmente solo en tareas de clasificación multiclase de etiqueta única, aunque el mecanismo de consistencia a nivel de logits es conceptualmente extensible a entornos de etiquetas múltiples.