Resumen Técnico: COSMOS

Enunciado del Problema

El Aprendizaje Federado (FL) enfrenta desafíos significativos en entornos heterogéneos donde los clientes difieren no solo en la distribución de datos (heterogeneidad estadística), sino también en la arquitectura del modelo (heterogeneidad arquitectónica). Si bien se han desarrollado métodos de Aprendizaje Federado Personalizado (PFL) para adaptar modelos a clientes individuales, la mayoría de los enfoques existentes asumen compatibilidad estructural o requieren que los clientes revelen arquitecturas de modelos propietarias. Esto crea una barrera práctica, ya que los clientes a menudo desean utilizar modelos arbitrarios o propietarios sin revelar parámetros o gradientes.

Los enfoques actuales de FL agnósticos al modelo, como FedMD y COMET, dependen de la comunicación a nivel de salida (por ejemplo, etiquetas suaves o duras) utilizando un conjunto de datos no etiquetado compartido. Sin embargo, estos métodos sufren limitaciones específicas:

1. Coordinación Pasiva del Servidor: A menudo tratan al servidor como un coordinador pasivo, sin aprovechar los recursos computacionales del servidor para la personalización activa.
2. Agrupación Heurística: Métodos como COMET dependen de la agrupación K-means, lo que requiere conocimiento previo de la cantidad de clusters ($K$) y carece de mecanismos formales para garantías de personalización cuando la diversidad de clientes es alta.
3. Falta de Garantías Teóricas: La literatura existente de PFL agnóstica al modelo carece generalmente de garantías teóricas rigurosas sobre la contracción del riesgo, dependiendo a menudo de supuestos simplificados (por ejemplo, modelos lineales, distribuciones gaussianas) o proporcionando únicamente resultados estándar de convergencia a la estacionariedad.

Metodología: Marco COSMOS

Los autores proponen COSMOS (Modelos de Servidor Basados en Salida Agrupada), un marco de PFL agnóstico al modelo que permite la personalización del lado del servidor utilizando únicamente comunicación de pseudo-etiquetas. El marco opera bajo el supuesto de que los clientes poseen datos etiquetados privados y que está disponible un conjunto de datos público no etiquetado compartido $U$ (por ejemplo, datos extraídos de la web o sintéticos).

COSMOS procede en dos fases distintas:

1. Fase de Pre-entrenamiento y Agrupación (PTC)

Esta fase establece el mapeo inicial cliente-servidor y los modelos específicos de cluster:

• Pre-entrenamiento Local: Los clientes entrenan sus modelos locales $f_i$ en sus datos privados $D_i$ y generan pseudo-etiquetas en el conjunto de datos público $U$.
• Agrupación Codiciosa: El servidor agrupa a los clientes basándose en la similitud de sus pseudo-etiquetas en $U$. A diferencia de K-means, COSMOS emplea un algoritmo de agrupación codiciosa que no requiere que $K$ se especifique de antemano. Selecciona iterativamente un cliente con los vecinos más "cercanos" (dentro de un umbral de distancia $B_0$) para formar un cluster hasta que todos los clientes estén asignados. Este proceso es endógeno al procedimiento de entrenamiento.
• Destilación del Lado del Servidor: Para cada cluster $G_k$, el servidor agrega las pseudo-etiquetas de sus clientes miembros para crear una distribución media de pseudo-etiquetas. Luego entrena un modelo "maestro" dedicado $h_k$ para ese cluster utilizando estas etiquetas agregadas y un término de regularización específico para garantizar la suavidad local y la robustez.
• Destilación del Lado del Cliente: Los clientes reciben las pseudo-etiquetas del modelo maestro de su cluster asignado y ajustan finamente sus modelos locales utilizando estas etiquetas.

2. Fase de Ajuste Fino Federado Iterativo (IFFT)

Tras la fase PTC, el sistema entra en un bucle iterativo:

• Ajuste Fino Local: Los clientes refinan sus modelos en sus datos etiquetados privados.
• Destilación del Servidor: Los clientes envían pseudo-etiquetas actualizadas al servidor. El servidor actualiza los modelos maestros específicos del cluster basándose en las nuevas pseudo-etiquetas agregadas.
• Destilación del Cliente: Los clientes refinan aún más sus modelos locales utilizando las pseudo-etiquetas actualizadas del maestro.

Crucialmente, los modelos del servidor se entrenan utilizando un término de regularización derivado de la teoría del aprendizaje semi-supervisado (SSL), que penaliza las inconsistencias de predicción bajo transformaciones de datos (por ejemplo, rotaciones, traslaciones) para garantizar la robustez.

Contribuciones Teóricas

El artículo proporciona el primer análisis teórico de extremo a extremo de la contracción del riesgo de personalización para PFL agnóstico al modelo. Al aprovechar herramientas de la teoría de SSL (específicamente conectividad basada en expansión y error de pseudo-etiqueta acotado), los autores derivan condiciones suficientes para la contracción exponencial del riesgo.

Los componentes teóricos clave incluyen:

• Supuestos: El análisis asume pesos de importancia acotados entre las distribuciones pública y privada, objetivos localmente bien condicionados y un "margen de confianza" donde los clientes son confiables en sus predicciones sobre los datos públicos. También depende de una propiedad de expansión de la distribución de datos, asegurando que las pseudo-etiquetas porten información no trivial sobre la verdad fundamental.
• Contracción del Riesgo: Los autores demuestran que bajo estas condiciones, el riesgo de personalización tanto para los modelos del servidor como para los modelos de los clientes se contrae exponencialmente con el número de iteraciones. Específicamente, el límite de riesgo disminuye a una tasa determinada por la calidad del cluster y la cobertura del conjunto de datos público, cayendo eventualmente por debajo del riesgo del entrenamiento local por sí solo.

Resultados Experimentales

Los autores evaluaron COSMOS en cuatro benchmarks de clasificación de imágenes: EMNIST-balanced, CIFAR-10, CIFAR-100 y Tiny ImageNet. Los experimentos utilizaron una mezcla de arquitecturas de clientes (AlexNet, ResNet, MobileNet, SqueezeNet) y un modelo de servidor VGG16.

• Rendimiento: COSMOS superó consistentemente a todos los baselines agnósticos al modelo (FedMD, FedCT, COMET) en todos los benchmarks. La brecha de rendimiento se amplió a medida que aumentaba la complejidad de la tarea (por ejemplo, de EMNIST a Tiny ImageNet).
• Heterogeneidad: COSMOS permaneció robusto ante diversos grados de heterogeneidad de datos (controlados por el parámetro Dirichlet $\alpha$) y diferentes mezclas de arquitecturas de clientes. Mantuvo una ventaja de rendimiento significativa incluso cuando los clientes utilizaban modelos más débiles.
• Eficiencia de Comunicación: Al comunicar únicamente pseudo-etiquetas (duras o suaves) en lugar de parámetros del modelo o gradientes, COSMOS redujo el ancho de banda de comunicación en 1–2 órdenes de magnitud en comparación con los métodos de intercambio de parámetros.
• Entornos Homogéneos: Incluso en entornos homogéneos donde todos los clientes comparten la misma arquitectura, COSMOS permaneció competitivo con los métodos PFL de última generación, demostrando que su diseño agnóstico al modelo no incurre en una penalización de rendimiento cuando existe homogeneidad.

Significado y Afirmaciones

El artículo afirma que COSMOS aborda un cuello de botella crítico en FL al permitir una personalización de alto rendimiento sin requerir conocimiento estructural de los modelos de los clientes. Su importancia radica en tres áreas principales:

1. Despliegue Práctico: Elimina las barreras estructurales que limitan el PFL en ecosistemas diversos y del mundo real donde los clientes utilizan arquitecturas propietarias o arbitrarias.
2. Avance Teórico: Proporciona las primeras garantías generales de contracción del riesgo para PFL agnóstico al modelo, yendo más allá de los resultados estándar de convergencia para probar la reducción exponencial del riesgo bajo supuestos realistas.
3. Escalabilidad: Demuestra que el aprendizaje personalizado del lado del servidor utilizando pseudo-etiquetas es un paradigma viable y escalable para entornos altamente heterogéneos, ofreciendo ganancias sustanciales en eficiencia de comunicación.

Los autores concluyen que, aunque el trabajo actual se centra en una estructura cliente-servidor, el marco sienta las bases para futuros sistemas de aprendizaje descentralizado, mejorando aún más la relevancia práctica del aprendizaje federado.