Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

Resumen Técnico: Imputación Probabilística de Características y Agregación Federada Multimodal Consciente de la Incertidumbre

1. El Problema

El aprendizaje federado (FL) permite entrenar modelos colaborativos en aplicaciones de salud preservando la privacidad de los datos. Sin embargo, un desafío fundamental en entornos clínicos reales es la heterogeneidad de modalidades:

• Disparidad de Recursos: Mientras que grandes centros médicos suelen tener datos completos (imágenes de rayos X y reportes textuales), clínicas más pequeñas o hospitales rurales a menudo solo poseen un subconjunto de modalidades (generalmente solo imágenes).
• Limitaciones de los Métodos Actuales: Las soluciones existentes utilizan redes de imputación de características deterministas para sintetizar la modalidad faltante (ej. generar un reporte textual basado en una imagen).
• Riesgo Crítico: Estos métodos deterministas producen estimaciones puntuales (un solo valor) sin medidas de fiabilidad. En aplicaciones médicas de seguridad crítica, esto es peligroso porque el modelo puede "alucinar" características con alta confianza incluso cuando la predicción es incorrecta, propagando errores silenciosos al clasificador final. Además, las estrategias de agregación federada estándar (como FedAvg) tratan todas las actualizaciones de los clientes por igual, sin considerar la calidad de la imputación.

2. Metodología Propuesta

Los autores proponen P-FIN (Probabilistic Feature Imputation Network) y una estrategia de agregación llamada Fed-UQ-Avg. El enfoque se divide en dos niveles:

A. Red de Imputación de Características Probabilística (P-FIN)
En lugar de predecir un vector fijo para la modalidad faltante, P-FIN modela la distribución condicional $p(z_T | z_I)$ como una Gaussiana Heterocedástica.

• Arquitectura: Utiliza un codificador de imágenes (ResNet-50) y un codificador de texto (BERT). La imputación se realiza mediante un transformador que toma las características de la imagen y un token de consulta aprendible para predecir dos salidas:
  1. Media ($\mu$): La estimación de la característica faltante.
  2. Varianza ($\sigma^2$): La incertidumbre por dimensión asociada a esa estimación.
• Entrenamiento: Se utiliza una función de pérdida especial llamada $\beta$-NLL (Negative Log-Likelihood). Esta función evita que el modelo minimice el error simplemente prediciendo una varianza infinita (un "atajo" común en modelos probabilísticos), forzando al modelo a aprender representaciones significativas mientras calibra su incertidumbre.

B. Fusión Local Consciente de la Incertidumbre
En los clientes unimodales (solo imágenes), las características imputadas se fusionan con las características de la imagen observada antes de la clasificación:

• Puerta de Incertidumbre (Gating): Se calcula una puerta $g = \text{sigmoid}(-\log \sigma^2)$. Si la incertidumbre ($\sigma^2$) es alta, la puerta se cierra (se acerca a 0), atenuando o suprimiendo las dimensiones de la característica imputada que no son confiables.
• Atención Cruzada: Las características de imagen y texto (filtradas por la puerta) se fusionan mediante mecanismos de atención, permitiendo que el modelo pondere dinámicamente la contribución de cada modalidad.

C. Agregación Federada Global (Fed-UQ-Avg)
Para mejorar la agregación global, se introduce un nuevo esquema de ponderación que no solo considera el tamaño del conjunto de datos, sino también la calidad de la imputación:

• Peso de Confianza: Se calcula una métrica de confianza basada en la incertidumbre media del cliente ($\bar{\sigma}^2_k$). Los clientes con baja incertidumbre (imputación de alta calidad) reciben un peso mayor.
• Fórmula de Ponderación: El peso final $W_k$ es una combinación lineal del peso basado en datos ($w_{data}$) y el peso basado en confianza ($w_{conf}$):
  $$W_k = (1 - \alpha)w_{data} + \alpha w_{conf}$$
  Esto asegura que los clientes con datos de baja calidad o imputaciones poco fiables no degraden el modelo global.

3. Contribuciones Clave

1. P-FIN: Primera red de imputación de características probabilística diseñada específicamente para entornos federados multimodales, que proporciona estimaciones de incertidumbre calibradas junto con las características sintetizadas.
2. Mecanismo de Doble Nivel: Integración de la incertidumbre tanto localmente (supresión de características no fiables antes de la clasificación) como globalmente (agregación ponderada por calidad).
3. Uso de $\beta$-NLL: Aplicación exitosa de esta función de pérdida para entrenar redes de imputación que evitan el colapso de la varianza y producen incertidumbre significativa.
4. Fed-UQ-Avg: Una nueva estrategia de agregación que prioriza las actualizaciones de clientes con imputaciones confiables, mitigando el impacto de los clientes con datos incompletos y de baja calidad.

4. Resultados Experimentales

Los experimentos se realizaron en tareas de clasificación de rayos X torácicos utilizando tres conjuntos de datos públicos: CheXpert, NIH Open-I y PadChest. Se simularon entornos federados con diferentes ratios de clientes unimodales a multimodales (8:2, 6:4, 4:6).

• Rendimiento Superior: El enfoque propuesto (P-FIN + Fed-UQ-Avg) superó consistentemente a todas las líneas base deterministas.
  • En la configuración más desafiante (80% de clientes sin texto, ratio 8:2), se logró una mejora absoluta del 5.36% en el AUC (Área bajo la curva) en comparación con el método determinista más avanzado (FIN + FedAvg).
  • La mejora fue de aproximadamente 2.4% a 2.8% adicional al comparar P-FIN con FedAvg estándar, demostrando el valor crítico de la agregación consciente de la incertidumbre.
• Calibración: El análisis de la incertidumbre mostró un Error de Calibración Esperado (ECE) bajo (0.0422), indicando que las predicciones de incertidumbre son fiables.
• Correlación Error-Incertidumbre: Se observó una fuerte correlación monótona: a medida que aumenta la incertidumbre predicha, aumenta el error real de imputación, validando que el modelo sabe cuándo "no sabe".

5. Significado e Impacto

Este trabajo es fundamental para el despliegue seguro de la IA en la salud:

• Seguridad Clínica: Aborda el requisito regulatorio (ej. FDA) de cuantificar la incertidumbre. Permite a los radiólogos confiar en las predicciones cuando la incertidumbre es baja y revisarlas manualmente cuando es alta, evitando diagnósticos erróneos por "alucinaciones" del modelo.
• Robustez en Heterogeneidad: Proporciona una solución robusta para la realidad clínica donde los datos son incompletos y no están distribuidos uniformemente (non-IID), permitiendo que hospitales con recursos limitados participen en el aprendizaje federado sin degradar el modelo global.
• Paradigma de "Aprender a Desconfiar": Introduce un cambio de paradigma donde el modelo no solo aprende a predecir, sino también a cuantificar su propia fiabilidad, creando sistemas multimodales federados más resilientes y dignos de confianza.