Federated Learning with Energy-Based Structured Probabilistic Inference

Resumen Técnico: Aprendizaje Federado con Inferencia Probabilística Estructurada Basada en Energía

Declaración del Problema
El Aprendizaje Federado (FL) depende típicamente de la Promediación Federada (FedAvg), donde el servidor agrega las actualizaciones de los modelos de los clientes utilizando pesos proporcionales a las cantidades de muestras locales. Aunque es simple, este heurístico se vuelve subóptimo cuando los clientes poseen distribuciones de datos heterogéneas (non-IID). En tales entornos, las reglas de ponderación fijas tratan todas las actualizaciones escaladas como igualmente confiables, lo que provoca deriva del cliente (client drift), convergencia lenta o entrenamiento inestable. Las soluciones existentes, como los términos proximales (FedProx), las variables de control (SCAFFOLD) o las reglas de agregación robustas (por ejemplo, media recortada, mediana geométrica), abordan modos de falla específicos pero generalmente dependen de heurísticos fijos. No logran modelar explícitamente tanto la confiabilidad individual de los clientes como las interacciones entre los clientes dentro de un marco de inferencia unificado.

Metodología
Los autores proponen reformular el problema de agregación del lado del servidor como una tarea de inferencia probabilística estructurada basada en energía utilizando Campos Aleatorios Condicionales (CRFs). La idea central es inferir "etiquetas de confiabilidad" discretas para cada cliente basándose en sus actualizaciones de modelo locales, en lugar de depender únicamente en las cantidades de muestras.

1. Formulación de Grafo:

   • Nodos: Clientes participantes en una ronda de comunicación.
   • Observaciones ($x$): Los vectores de actualización de modelos locales (deltas) y las cantidades de muestras.
   • Variables Latentes ($y$): Etiquetas de confiabilidad discretas (por ejemplo, $\{0, 0.25, 1\}$), donde valores más altos indican una mayor confianza en la actualización.

2. Función de Energía:
   El CRF define una función de energía $E(y|x)$ compuesta por dos potenciales:

   • Potenciales Unarios ($\psi_u$): Miden la confiabilidad individual del cliente. El servidor computa una actualización de referencia robusta (mediana coordenada a coordenada de todos los deltas de los clientes). Posteriormente, deriva tres estadísticas normalizadas para cada cliente: acuerdo de coseno con la referencia, distancia relativa de la referencia y una puntuación de anomalía de la norma de actualización robusta. Estas se combinan mediante una función sigmoide para generar una puntuación de confiabilidad, la cual inicializa el potencial unario.
   • Potenciales Pareados ($\psi_p$): Fomentan la consistencia entre clientes con actualizaciones similares. La afinidad entre el cliente $i$ y el $j$ se define por la similitud de coseno de sus actualizaciones, regulada por el mínimo de sus puntuaciones de confiabilidad individuales. Los clientes con alta afinidad son penalizados si se les asignan etiquetas de confiabilidad divergentes.

3. Inferencia y Agregación:

   • Dado que la inferencia de Máxima Verosimilitud a Posteriori (MAP) exacta es intratable para un grafo totalmente conectado, los autores emplean una aproximación de campo medio (mean-field). Esta actualiza iterativamente la distribución posterior de las etiquetas de confiabilidad para cada cliente basándose en las puntuaciones unarias y los costos esperados de los vecinos.
   • El peso de agregación final para un cliente es el producto de su peso base de FedAvg (cantidad de muestras) y su confiabilidad esperada inferida. Este peso refinado se utiliza para computar el nuevo modelo global.

Contribuciones Clave

• Formulación: El artículo formula la agregación de FL como un problema de inferencia estructurada basada en energía sobre variables de confiabilidad de los clientes.
• Instanciación: Implementa un CRF a nivel de cliente totalmente conectado donde los potenciales unarios se derivan de estadísticas de actualización robustas (basadas en la mediana) y los potenciales pareados se derivan de la afinidad de las actualizaciones.
• Algoritmo: Propone un algoritmo de agregación de campo medio compatible con los bucles de entrenamiento estándar de FedAvg, que no requiere cambios en la optimización del lado del cliente.
• Evaluación: El método se evalúa frente a líneas base estándar de FL (FedAvg, FedProx) y métodos de agregación robustos (Krum, media recortada, mediana geométrica, RFA).

Resultados Experimentales
Los autores evaluaron el enfoque en MNIST, CIFAR-10 y CIFAR-100 bajo condiciones severas de non-IID (sesgo de etiqueta Dirichlet con $\alpha=0.003$ y $\alpha=0.1$).

• Desempeño: El enfoque guiado por CRF (CRF-FedAvg y CRF-FedProx) superó consistentemente a las líneas base de agregación fija y robusta en todos los conjuntos de datos.
• Significancia de las Ganancias: Mientras que las mejoras en MNIST fueron modestas (ya que la mayoría de los métodos alcanzaron una alta precisión), las ganancias fueron más pronunciadas en CIFAR-10 y CIFAR-100. CRF-FedProx logró la mayor precisión en CIFAR-10, y CRF-FedAvg logró la mayor en CIFAR-100.
• Convergencia: Las curvas de pérdida de validación indicaron que la agregación guiada por CRF convergió de manera más rápida y estable que las líneas base robustas, particularmente en las rondas de comunicación tardías. Esto sugiere que el método preserva eficazmente las direcciones de actualización útiles mientras suprime las no confiables, evitando el estancamiento conservador visto en métodos como la mediana geométrica.

Significancia y Reivindicaciones
El artículo afirma que tratar la agregación como un problema de inferencia estructurada ofrece una ventaja distinta sobre los heurísticos fijos. Al modelar explícitamente tanto la confiabilidad de la actualización individual como la compatibilidad cliente-cliente, el servidor puede asignar pesos que son adaptativos a la geometría específica de la actualización de la ronda actual.

• Los autores enfatizan que este enfoque es complementario a las técnicas de optimización del lado del cliente (como FedProx); mientras que FedProx restringe la deriva local, la agregación guiada por CRF refina cómo se combinan las actualizaciones post-entrenamiento.
• El método se presenta como un complemento (plug-in) del lado del servidor que no altera el entrenamiento ni los protocolos de comunicación del cliente.
• Los autores señalan limitaciones, incluyendo el uso de potenciales diseñados a mano (en lugar de aprendidos), hiperparámetros fijos en todos los experimentos, y el costo computacional del grafo totalmente conectado ($O(m^2)$), que podría requerir esparcimiento para federaciones de gran escala. También declaran que la evaluación actual se centra en la heterogeneidad estadística en lugar de la robustez adversaria (Byzantine).