Resumen Técnico: FedProIn – Mitigación del Desplazamiento de Cliente para Prototipos Aprendibles en Imágenes Médicas Federadas

Declaración del Problema

El Aprendizaje Federado (FL, por sus siglas en inglés) en imágenes médicas se ve significativamente obstaculizado por la heterogeneidad estadística derivada de las variaciones en los protocolos de escaneo, los métodos de adquisición y las diversas poblaciones de pacientes en diferentes centros médicos. Esta distribución de datos no IID (no independientes e idénticamente distribuidos) induce un desplazamiento de cliente (client drift), donde los modelos locales divergen del modelo global durante la optimización. Específicamente, instancias idénticas son mapeadas a regiones distintas en el espacio de características por diferentes clientes, y las representaciones locales se sesgan hacia las características específicas de cada cliente.

Los enfoques existentes enfrentan limitaciones:

• La regularización basada en restricciones (por ejemplo, FedProx, MOON) tiene dificultades para alinear las representaciones semánticas en espacios de alta dimensión.
• Los métodos basados en prototipos estándar a menudo asumen un único prototipo global por clase (mediante vectores de características medios), lo cual es inadecuado para datos médicos heterogéneos.
• Los métodos de múltiples prototipos a menudo introducen una alta complejidad computacional o dependen de representaciones unimodales que no logran optimizar eficazmente los límites de decisión.

Metodología: FedProIn

Los autores proponen FedProIn (Aprendizaje Federado con agregación de prototipos con influencia), un marco que utiliza Prototipos Aprendibles (LPs) como parámetros optimizables del modelo en lugar de vectores medios estáticos. El modelo global consiste en un codificador de características $f_\theta$ y un conjunto de prototipos $P \in \mathbb{R}^{C \times M \times d}$, donde $M$ es el número de prototipos por clase.

1. Descomposición del Desplazamiento de Cliente

FedProIn aborda explícitamente el desplazamiento de cliente descomponiéndolo en dos componentes y aplicando pérdidas de regularización específicas:

• Desplazamiento de Características (Feature Drift): Abordado mediante la Pérdida de Divergencia de Características (FDL). Los clientes mantienen una copia congelada del extractor de características global ($\theta^{t-1}$). El codificador local es regularizado para minimizar el error cuadrático medio entre las incrustaciones de características locales y las características globales de referencia, asegurando la alineación semántica.
• Desplazamiento de Prototipos (Prototype Drift): Abordado mediante la Pérdida Contrastiva de Prototipos (PCL). Esta pérdida acerca los prototipos locales actualizados a los prototipos globales actuales, mientras los aleja de sus estados locales previos. Esto evita que los prototipos se sesguen excesivamente hacia los cambios de datos locales.

El objetivo local total combina la pérdida de clasificación primaria ($\ell_{Task}$) con estos términos de regularización:
$$\ell_{Total} = \ell_{Task} + \lambda_{FDL}\ell_{FDL} + \lambda_{PCL}\ell_{PCL}$$

Para manejar el desequilibrio de clases en conjuntos de datos médicos de cola larga (long-tailed), el método emplea muestreo balanceado por clase y aumentación adaptativa para asegurar que los subtipos de enfermedades raras estén adecuadamente representados durante las actualizaciones de los prototipos.

2. Agregación de Influencia Normalizada (NIA)

La agregación estándar (por ejemplo, FedAvg) es insuficiente para múltiples prototipos porque pondera las actualizaciones únicamente según el tamaño del conjunto de datos, ignorando la utilización de los prototipos. Esto puede amplificar el ruido de los prototipos subutilizados o suprimir las representaciones especializadas.

FedProIn introduce la Agregación de Influencia Normalizada (NIA):

• Matriz de Influencia ($I_i$): Durante el entrenamiento local, cada cliente rastrea una puntuación de influencia para cada prototipo, contando cuántas veces un prototipo específico sirvió como el vecino más cercano para las muestras locales.
• Agregación: Los prototipos globales se agregan ponderándolos proporcionalmente a sus puntuaciones de uso ($I_{k,c}^i$) en lugar de solo al tamaño del conjunto de datos del cliente. Esto reduce el peso de las actualizaciones ruidosas de prototipos no utilizados y preserva los prototipos especializados relevantes para subconjuntos específicos de clientes.

Principales Contribuciones

1. Formulación de Prototipos Aprendibles: Adoptar los prototipos como parámetros optimizables (conjuntamente con el extractor de características) en lugar de vectores medios estáticos para maximizar los márgenes de separación de clases.
2. Estrategia de Doble Regularización: Mitigar explícitamente el desplazamiento de cliente descomponiéndolo en desplazamiento de características y de prototipos, abordados mediante FDL y PCL respectivamente.
3. Agregación de Influencia Normalizada (NIA): Un novedoso mecanismo de agregación que pondera adaptativamente los prototipos de los clientes basándose en su utilidad local, abordando problemas de utilización heterogénea y activación dispersa en entornos de múltiples prototipos.
4. Validación Empírica: Evaluación exhaustiva en conjuntos de datos de imágenes médicas públicos demostrando robustez contra la heterogeneidad estadística.

Resultados Experimentales

El marco fue evaluado en dos conjuntos de datos de imágenes médicas de cola larga: HAM10000 (imágenes dermatoscópicas) y Matek-19 (imágenes de células sanguíneas periféricas). Los experimentos se realizaron bajo configuraciones IID y Non-IID (Dirichlet $\alpha=0.5$) con $K=10$ clientes.

Aspectos Destacados del Rendimiento:

• HAM10000: FedProIn alcanzó un 83.5% de precisión (IID) y un 81.1% de precisión (Non-IID), superando a todos los modelos de referencia, incluyendo FedAvg, FedProx, MOON y otros métodos basados en prototipos (FedProto, FedTGP, FedPLVM).
• Matek-19: FedProIn alcanzó un 96.2% de precisión (IID) y un 95.8% de precisión (Non-IID).
• Robustez: Bajo condiciones Non-IID, los métodos de prototipos existentes sufrieron caídas significativas de rendimiento (por ejemplo, una pérdida de precisión del 6–7% en HAM10000), mientras que FedProIn restringió las pérdidas relativas a apenas un 2.4% y 0.4% en HAM10000 y Matek-19, respectivamente.
• Estudio de Ablación: Los resultados confirmaron que tanto FDL como PCL son necesarios para un rendimiento óptimo. Eliminar FDL causó la caída más pronunciada en el rendimiento, indicando que el desplazamiento de características es el modo de falla dominante. El modelo completo superó significativamente a las variantes de componente único y al caso degenerado (equivalente a FedAvg).

Significado y Reivindicaciones

El artículo afirma que FedProIn mitiga eficazmente el desplazamiento de cliente en imágenes médicas federadas mediante el uso de prototipos aprendibles y agregación consciente de la influencia. Los autores sostienen que su enfoque:

• Logra una generalización superior en comparación con los modelos base de FL existentes y los métodos de personalización, particularmente en entornos altamente heterogéneos (Non-IID).
• Aborda con éxito las limitaciones de las suposiciones de prototipo único y la agregación sin ponderar en entornos de múltiples prototipos.
• Proporciona una solución robusta para las imágenes médicas donde la heterogeneidad de los datos y el desequilibrio de clases son prevalentes.

Los autores señalan que el trabajo futuro explorará la adaptación de FedProIn para FL personalizado (eliminando la dependencia del modelo global) y la extensión del marco a configuraciones de FL heterogéneo con diversas arquitecturas de modelos. El código se pone a disposición para apoyar la reproducibilidad.