PRISM-VFL: A Differentially Private Vertical Federated Framework for Heterogeneous Multi-Task Clinical Prediction

Resumen Técnico: PRISM-VFL

Planteamiento del Problema
El Aprendizaje Federado Vertical (VFL, por sus siglas en inglés) permite el entrenamiento de modelos colaborativos entre instituciones que poseen características complementarias de pacientes sin compartir datos brutos. Sin embargo, los marcos de trabajo de VFL existentes enfrentan dos limitaciones críticas en entornos clínicos: (1) carecen de soporte para la optimización concurrente de tareas de predicción heterogéneas (Aprendizaje Multitarea, o MTL) bajo garantías formales de privacidad, y (2) a menudo asumen una única parte activa o estructuras de tareas uniformes. En la atención médica del mundo real, las instituciones pueden poseer conjuntos de características disjuntos (por ejemplo, signos vitales frente a resultados de laboratorio) para los mismos pacientes mientras persiguen objetivos clínicos distintos (por ejemplo, predicción de mortalidad frente a fenotipado). Además, la integración de la Privacidad Diferencial (DP) en el MTL federado introduce compromisos complejos; la aplicación uniforme de ruido puede degradar la utilidad para tareas de alta prioridad, mientras que la interacción entre la sensibilidad específica de la tarea y las representaciones compartidas permanece poco explorada.

Metodología
El estudio introduce PRISM-VFL (Privacy-Resilient Institutional Split Multi-task VFL), un marco diseñado para la predicción clínica multitarea heterogénea.

• Arquitectura: El sistema sigue un paradigma de aprendizaje dividido (split learning). Cada sitio participante (simulado como tres hospitales) procesa sus características de series temporales locales utilizando un codificador de Memoria de Corto Plazo (LSTM) de dos capas, produciendo un embedding de 64 dimensiones. Estos embeddings se transmiten a un servidor central.
• Procesamiento en el Servidor: El servidor concatena los embeddings de todos los sitios en un vector de 192 dimensiones. Este vector es procesado por un módulo de Mezcla de Expertos con Puertas Múltiples (MMoE, Multi-gate Mixture-of-Experts), que consta de cuatro subredes de expertos compartidos y una puerta softmax aprendida para cada tarea. Esta arquitectura permite que las tareas se dirijan a través de diferentes especializaciones de expertos, mitigando la transferencia negativa.
• Mecanismo de Privacidad: El marco aplica el Gradiente Descendente Estocástico con Privacidad Diferencial (DP-SGD) específicamente en la capa de corte. Los gradientes de embedding por muestra ($\partial L/\partial \langle k$) devueltos desde el servidor a los clientes son recortados y perturbados con ruido Gaussiano antes de actualizar los codificadores locales. El estudio evalúa tanto la asignación de ruido uniforme como la asignación de ruido estratificada por tareas, donde los multiplicadores de ruido se ajustan según la urgencia clínica (por ejemplo, menor ruido para la Mortalidad Intrahospitalaria para preservar la utilidad).
• Conjunto de Datos y Simulación: El marco fue evaluado en una versión verticalmente particionada del conjunto de datos MIMIC-III. Las 14 variables clínicas se distribuyeron en tres sitios simulados:
  • Sitio A: Signos vitales (7 variables) con la tarea de Mortalidad Intrahospitalaria (IHM).
  • Sitio B: Resultados de laboratorio (4 variables) con la tarea de estratificación de riesgo de descompensación.
  • Sitio C: Mediciones estáticas al momento del ingreso (3 variables) con la tarea de Fenotipado (25 grupos de enfermedades).
  • La alineación de pacientes se logró mediante una simulación de Intersección de Conjuntos Privados (PSI), reteniendo solo pacientes con estancias en la UCI $\ge$ 48 horas.
• Evaluación: El estudio utilizó un conjunto integral de métricas, incluyendo AUC-ROC, puntuación de Brier y una novedosa puntuación de Utilidad Multitarea Ajustada por Privacidad-Comunicación (PCMU). También introdujo un límite de inferencia de etiquetas multitarea derivado teóricamente, $g(\sigma, \rho) = \Phi(C\sqrt{1+\rho}/\sigma)$, y lo validó contra ataques empíricos en el espacio de embeddings (inferencia de etiquetas e inferencia de membresía).

Contribuciones Clave

1. Desarrollo del Marco: PRISM-VFL es la primera arquitectura de VFL que soporta simultáneamente conjuntos de tareas heterogéneas a través de sitios particionados verticalmente con garantías explícitas de DP, utilizando MMoE para el enrutamiento específico de tareas.
2. Análisis de Privacidad por Tarea: El estudio proporciona un análisis granular de los compromisos entre privacidad y utilidad, demostrando que un multiplicador de ruido uniforme único no logra capturar la inestabilidad específica de cada tarea. Introduce un método para evaluar los umbrales de utilidad clínica ($\epsilon^*$) para tareas individuales dentro de un entorno multitarea.
3. Límites Teóricos y Empíricos: Se deriva un límite de referencia teórico para la inferencia de etiquetas multitarea y se valida contra una suite de ataques empíricos, mostrando que el enrutamiento de MMoE mantiene los riesgos de inferencia multitarea dentro de los límites de la tarea única cuando el acoplamiento de gradientes es bajo.
4. Referencia Reproducible: Los autores establecen un benchmark de división vertical de MIMIC-III reproducible que particiona tanto las características como los objetivos de predicción, facilitando futuras investigaciones en VFL multitarea de múltiples sitios.

Resultos

• Desempeño: PRISM alcanzó un AUC-ROC medio de 0.781 para IHM, 0.632 (con alta varianza entre semillas) para Descompensación, y 0.613 (macro-AUC) para Fenotipado.
• Ganancia Multitarea: Se observaron ganancias significativas para la Descompensación (+0.093 AUC) en comparación con las líneas base de tarea única, atribuidas al contexto fisiológico auxiliar proporcionado por otros sitios. IHM y Fenotipado mostraron ganancias marginales (+0.003 y +0.002, respectivamente). No se detectó transferencia negativa sistemática en las configuraciones.
• Compromiso Privacidad-Utilidad:
  • IHM: Se mantuvo clínicamente viable bajo restricciones de privacidad moderadas ($\epsilon \in (1, 2)$), manteniendo un AUC por encima del suelo de utilidad de 0.75.
  • Descompensación y Fenotipado: No lograron alcanzar sus respectivos suelos de utilidad clínica (0.70 y 0.65) incluso en ausencia de restricciones de DP ($\epsilon = \infty$), lo que indica que la arquitectura actual y la distribución de datos son insuficientes para estas tareas específicas en esta configuración.
  • Estabilidad: La estocasticidad de DP causó una inestabilidad asimétrica. IHM mostró un pico de inestabilidad en $\epsilon=1$ antes de recuperarse, mientras que la Descompensación y el Fenotipado se degradaron progresivamente.
• Eficacia de la Privacidad: Bajo una privacidad fuerte ($\epsilon = 0.5$), los ataques de inferencia de etiquetas se suprimieron hasta un rendimiento cercano al azar (AUC 0.498). Los ataques de inferencia de membresía permanecieron ineficaces en todos los presupuestos, un resultado atribuido a las propiedades arquitectónicas de la configuración de VFL más que a la DP por sí sola.
• Ablaciones: El estudio encontró que el enrutamiento de puertas aprendidas en MMoE fue el principal motor de los beneficios multitarea, no el aumento de capacidad. Curiosamente, las dimensiones de embedding más bajas (32 vs. 128) preservaron mejor la utilidad de IHM bajo fuertes restricciones de privacidad.

Significancia y Reivindicaciones
El artículo afirma que PRISM-VFL demuestra la viabilidad del aprendizaje federado vertical multitarea y heterogéneo con preservación de la privacidad. Sin embargo, los autores enfatizan que la preparación clínica es condicional a la tarea. Bajo las condiciones probadas, IHM es la única tarea que cumple con los umbrales de utilidad clínica bajo una privacidad moderada, mientras que la Descompensación y el Fenotipado requieren mayor desarrollo arquitectónico o enriquecimiento de datos antes de su despliegue.

El estudio posiciona a PRISM no solo como un benchmark de desempeño, sino como un modelo metodológico para evaluar qué problemas de predicción clínica son adecuados para el despliegue de aprendizaje multitarea vertical federado. Destaca que las estrategias de privacidad uniformes son insuficientes para tareas heterogéneas y que los sistemas futuros requieren mecanismos de calibración adaptativos a la tarea. El trabajo establece una línea base para la investigación de VFL de múltiples sitios, reconociendo que los resultados de un único conjunto de datos (MIMIC-III) probablemente representan una estimación optimista del desempeño en el mundo real debido a la falta de cambios en la distribución que se encuentran en despliegues reales de múltiples sitios.