Resumen Técnico: MuCALD-SplitFed

Enunciado del Problema

La segmentación de imágenes médicas es crítica para los flujos de trabajo clínicos, pero los sistemas prácticos de inteligencia artificial enfrentan obstáculos significativos en entornos descentralizados. Los hospitales y clínicas operan con datos heterogéneos (modos, protocolos y demografías variables) y rara vez comparten datos crudos debido a restricciones de privacidad. Aunque el Aprendizaje Federado (FL) y el Aprendizaje Federado Dividido (SplitFed) permiten el entrenamiento colaborativo sin compartir datos, las implementaciones estándar asumen que todos los clientes resuelven la misma tarea. Esta suposición de "no IID" es poco realista en entornos clínicos reales, donde las instituciones suelen manejar diferentes tareas específicas de dominio.

Los enfoques existentes de FL multi-tarea y SplitFed intentan abordar esto permitiendo que los clientes manejen diferentes objetivos. Sin embargo, estos métodos sufren frecuentemente de:

1. Inestabilidad de la Convergencia: Los gradientes de tareas heterogéneas entran en conflicto durante la agregación, lo que conduce a actualizaciones oscilatorias.
2. Entrelazamiento de Características: Las representaciones latentes compartidas en los puntos de división mezclan características específicas de la tarea, causando interferencia entre tareas.
3. Vulnerabilidades de Privacidad: La transmisión de activaciones latentes en los puntos de división es susceptible a ataques de reconstrucción, inferencia de membresía e inferencia de propiedades, riesgos que se amplifican por la diversidad de semánticas de tareas.

Las soluciones actuales suelen depender de representaciones basadas en correlación, que no logran distinguir entre estructuras relevantes para la tarea y correlaciones espurias, limitando la robustez ante cambios en la distribución.

Metodología: MuCALD-SplitFed

Los autores proponen MuCALD-SplitFed, un marco SplitFed multi-tarea diseñado para la segmentación de imágenes médicas que preserva la privacidad. La arquitectura extiende el modelo SplitFed estándar (particionado en componentes Front-End [FE], Lado del Servidor [SS] y Back-End [BE]) integrando tres etapas novedosas:

1. Construcción de Etiquetas Proxy y Descubrimiento de Grafos Causales

Para facilitar el aprendizaje causal débilmente supervisado sin requerir etiquetas causales explícitas, el sistema extrae características morfológicas e intensivas específicas del dominio de cada conjunto de datos para servir como "etiquetas proxy".

• Conjuntos de Datos: Blastocyst (embrión), FHPsAOPMSB (cabeza fetal), HAM10K (lesiones cutáneas), MosMed (TC de pulmón) y Kvasir-SEG (pólipos).
• Proceso: Utilizando Notears-MLP, el sistema aprende grafos causales específicos del conjunto de datos a partir de estas características proxy. Se seleccionan las tres relaciones causales más fuertes por conjunto de datos para guiar el modelo (por ejemplo, "Tamaño de la lesión $\to$ Asimetría").

2. Módulo de Representación Causal y Difusión (CRDM)

Este módulo opera en el lado del cliente (FE) para generar representaciones latentes robustas y que preserven la privacidad antes de la transmisión al servidor.

• Codificador Exógeno: Mapea los mapas de características del FE ($z_{fe}$) a variables exógenas de baja dimensión ($u$).
• SCM Neuronal: Aprende una factorización causal para producir latentes estructuralmente causales ($z_{causal}$).
• Difusión Forward: Aplica un paso de difusión forward a las características del codificador ($z_{fe}$) para ofuscar las activaciones transmitidas, creando latentes ruidosos ($z_{fe\_noisy}$).
• Decodificador de Difusión Causal (LDDM-1): Elimina el ruido de los latentes ruidosos condicionados a los latentes causales ($z_{causal}$) para generar $z_{denoised}$.
• Beneficio: Este proceso expone menos información reconstructiva que las activaciones estándar de SplitFed, mitigando los ataques de reconstrucción e inferencia de membresía.

3. Alineación Causal Adversarial de Dominio (DACA)

Ubicado en el lado del servidor (SS), este bloque impone invarianza de dominio para suprimir sesgos específicos de la tarea.

• Mecanismo: Emplea una Capa de Reversión de Gradiente (GRL) y un discriminador de dominio ligero.
• Objetivo de Entrenamiento: La red de segmentación intenta "engañar" al discriminador, produciendo representaciones agnósticas al dominio, mientras que el discriminador penaliza la filtración específica del dominio. Esto mejora la generalización entre tareas y la privacidad.

Optimización

El modelo se entrena minimizando conjuntamente una función de pérdida compuesta (Ecuación 1) que incluye:

• Pérdida de segmentación ($L_{seg}$)
• Pérdidas de alineación de proxy ($L_{proxy}$)
• Pérdidas de regularización de difusión ($L_{diff}$)
• Términos de divergencia KL para latentes exógenos y causales ($L_{KL}$)
• Pérdidas de dominio adversarias ($L_{adv}$)

Después de cada ronda de comunicación, los modelos SS, LDDM, codificadores y SCM Neuronales se agregan, mientras que los modelos FE y BE permanecen locales para preservar la personalización.

Contribuciones Clave

1. Integración Causal en SplitFed: La primera integración de modelado causal dentro del espacio latente de SplitFed específicamente para segmentación multi-tarea en entornos clínicos heterogéneos.
2. Arquitectura Mejorada en Privacidad: Una combinación novedosa de representaciones causales, difusión latente y desruido adversario para reducir la filtración de información en los puntos de división.
3. Desentrelazamiento de Características: Un mecanismo para separar factores causales relevantes para la tarea de la información específica del dominio, mejorando tanto el rendimiento de segmentación como la privacidad.
4. Evaluación Exhaustiva: Pruebas extensas en cinco conjuntos de datos médicos heterogéneos, comparando con SplitFed de referencia, FL personalizado de última generación (SoTA) y métodos de FL multi-tarea.

Resultados Experimentales

El marco se evaluó en cinco clientes utilizando conjuntos de datos diversos (Blastocyst, HAM10K, FHPsAOPMSB, MosMed, Kvasir-SEG) durante 24 rondas de comunicación.

• Estabilidad de la Convergencia: SplitFed de referencia no logró converger consistentemente, exhibiendo puntuaciones inestables de Intersección sobre Unión (IoU) (por ejemplo, IoU N/B promedio final de 0.001 para UNet). En contraste, MuCALD-SplitFed convergió de manera estable, logrando un IoU N/B promedio final de 0.487 para la arquitectura UNet.
• Rendimiento de Segmentación: MuCALD-SplitFed superó significativamente a SplitFed de referencia y a métodos SoTA (FedPer, FedRep, MOCHA, etc.).
  • Puntuación Dice: MuCALD-SplitFed logró 0.816 (promedio entre clientes), en comparación con 0.394 para SplitFed de referencia.
  • IoU N/B: MuCALD-SplitFed logró 0.579, en comparación con 0.045 para SplitFed de referencia.
• Métricas de Privacidad: Las métricas de calidad de reconstrucción (MSE, PSNR, SSIM) indicaron que los mapas de características en MuCALD-SplitFed estaban mucho más ofuscados. Por ejemplo, el PSNR promedio para ataques de reconstrucción disminuyó significativamente (por ejemplo, ~35.7 para MS frente a ~7.0 para BS), haciendo que los ataques de reconstrucción e inferencia de membresía fueran efectivamente inviables.
• Estudios de Ablación: Eliminar el componente de difusión causó la mayor degradación del rendimiento, destacando su papel crítico en la estabilidad y robustez. Tanto CRDM como DACA proporcionaron beneficios complementarios.

Significado y Afirmaciones

El artículo posiciona a MuCALD-SplitFed como una evolución necesaria para la IA médica descentralizada, abordando la desalineación entre las suposiciones estándar de FL y los flujos de trabajo clínicos reales. Los autores afirman que, al ir más allá de las representaciones basadas en correlación hacia la difusión latente causal, el marco logra con éxito:

• Resolver la inestabilidad de la convergencia causada por conflictos de gradientes en configuraciones multi-tarea.
• Mitigar los riesgos de privacidad inherentes a las comunicaciones en los puntos de división.
• Mejorar la generalización a través de dominios heterogéneos sin requerir el intercambio de datos crudos.

El trabajo se presenta como un "primer paso" hacia un aprendizaje SplitFed multi-tarea integral que preserve la privacidad. Los autores notan modestamente que el trabajo futuro implicará evaluar modelos adversarios más fuertes, evaluar la escalabilidad a poblaciones de clientes más grandes y extender el marco a configuraciones causales más amplias y mecanismos de privacidad adaptativos.