Resumen Técnico: Aprendizaje Federado mediante Inferencia Bayesiana Variacional

1. Declaración del Problema

El Aprendizaje Federado (FL, por sus siglas en inglés) permite el aprendizaje automático distribuido preservando la privacidad del cliente. Sin embargo, los enfoques estándar de FL, como el Promedio Federado (FedAvg), suelen sufrir una degradación del rendimiento al enfrentarse a dos desafíos críticos del mundo real:

1. Heterogeneidad de Datos (No-i.i.d.): Las distribuciones de datos de los clientes varían significamente debido a diferencias en los hábitos, ubicaciones y preferencias de los usuarios.
2. Escasez de Datos: Los clientes suelen poseer datos limitados debido al alto costo de recolección, lo que conduce al sobreajuste (overfitting) del modelo.

Las soluciones existentes intentan abordar estos problemas mediante FL Personalizado (PFL) y FL Agrupado (CFL). Sin embargo, los modelos PFL suelen sobreajustarse ante datos insuficientes, y los modelos de FL Bayesiano (BFL) estándar tienen dificultades con la heterogeneidad extrema. Además, el costo de inferencia de las Redes Neuronales Bayesianas (BNN) es alto debido a la necesidad de múltiples pasadas hacia adelante para estimar la incertidumbre. Este artículo aborda la necesidad de un marco que gestione simultáneamente la heterogeneidad, la escasez de datos y la eficiencia de la inferencia.

2. Metodología

Los autores proponen un marco de optimización de dos niveles que integra BNN en PFL y CFL mediante Inferencia Bayesiana Variacional (VI). A diferencia de la VI tradicional donde se asume un prior fijo para los parámetros, este marco utiliza la distribución global entrenada del servidor como la distribución prior para cada cliente. Esto evita el desajuste entre los priors asumidos y las distribiones reales.

El marco consta de tres algoritmos específicos:

2.1 pFedBayes (Aprendizaje Federado Bayesiano Personalizado)

• Mecanismo: Diseñado para escenarios de datos limitados. Asume que todos los parámetros de la red siguen distribuciones Gaussianas independientes e idénticamente distribuidas (i.i.d.).
• Optimización: Los clientes minimizan la suma del error de reconstrucción sobre sus datos locales y la divergencia de Kullback-Leibler (KL) entre su distribución local y la distribución global. El servidor agrega las distribuciones locales para actualizar la distribución global.
• Objetivo: Lograr la personalización mientras se aprovecha el prior global para prevenir el sobreajuste.

2.2 sFedBayes (Aprendizaje Federado Bayesiano Personalizado Disperso)

• Mecanismo: Diseñado para mejorar la eficiencia de la inferencia. Introduce dispersión (sparsity) asumiendo que los parámetros siguen una distribución Bernoulli-Gaussian (un prior de tipo "spike-and-slab").
• Optimización: Similar a pFedBayes pero incluye una estructura de dispersión aprendible. Los parámetros se podan si la probabilidad de inclusión es baja.
• Objetivo: Reducir los costos de almacenamiento y computación durante la inferencia mediante la poda de parámetros redundantes, manteniendo el rendimiento del modelo.

2.3 cFedBayes (Aprendizaje Federado Bayesiano Agrupado)

• Mecanismo: Diseñado para la diversidad estadística extrema. En lugar de aprender un único prior global, el servidor aprende múltiples distribuciones globales (una por cada grupo o cluster).
• Optimización: Los clientes seleccionan dinámicamente la distribución global (el grupo) que mejor se ajusta a sus datos privados (minimizando la divergencia KL) y actualizan su distribución local basándose en ese prior específico.
• Objetivo: Proporcionar garantías adaptativas a nivel de grupo en lugar de un único límite uniforme, acomodando mejor a los clientes con distribuciones de datos muy diferentes.

3. Contribuciones Clave

1. Algoritmos Novedosos: El artículo introduce tres algoritmos de BFL personalizado de dos niveles: pFedBayes, sFedBayes y cFedBayes.
2. Garantías Teóricas:
   • Los autores proporcionan límites de error de generalización para los tres enfoques.
   • Demuestran que las tasas de error de generalización promedio alcanzan la optimalidad minimax salvo por un factor logarítmico.
   • Específicamente, se muestra que cFedBayes logra un límite de error de generalización a nivel de grupo, ofreciendo garantías adaptativas superiores al límite uniforme único de pFedBayes.
3. Eficiencia y Dispersión: Se propone sFedBayes para abordar el alto costo de inferencia de las BNN. El análisis teórico muestra que alcanza tasas de convergencia minimax óptimas en el régimen disperso.
4. Algoritmos de Gradiente Estocástico: Se desarrollan algoritmos de gradiente estocástico correspondientes para los tres métodos, utilizando trucos de reparametrización (por ejemplo, Gumbel-softmax para variables discretas en sFedBayes) para permitir el entrenamiento de extremo a extremo.

4. Resultados Experimentales

Los autores evaluaron los métodos propuestos en cuatro conjuntos de datos de referencia (MNIST, FMNIST, CIFAR-10, SVHN) bajo configuraciones no-i.i.d. y de agrupamiento, comparándolos con líneas base de vanguardia (ej. FedAvg, FedProx, pFedMe, IFCA, FedCFL).

• Rendimiento con Datos Limitados: pFedBayes superó significativamente a otros métodos personalizados, particularmente en conjuntos de datos pequeños. Por ejemplo, en el conjunto de datos pequeño de CIFAR-10, pFedBayes mejoró la precisión personalizada en un 12.43% en comparación con el mejor algoritmo SOTA.
• Beneficios de la Dispersión: sFedBayes demostró que una dispersión moderada (por ejemplo, reteniendo ~35-45% de los parámetros) puede mantener o incluso mejorar ligeramente la precisión personalizada en comparación con el modelo denso pFedBayes, actuando como un regularizador implícito contra el sobreajuste.
• Efectividad del Agrupamiento: cFedBayes alcanzó las precisiones más altas en conjuntos de datos agrupados (ej. MNIST-MNIST⊥, CIFAR10-CIFAR10⊥), superando a las líneas base basadas en agrupamiento como FedCFL e IFCA por aproximadamente un 1–2% en precisión global.
• Estimación de la Incertidumbre: La naturaleza Bayesiana de los modelos permite la cuantificación de la incertidumbre, la cual los autores señalan que podría usarse para filtrar modelos de baja confianza durante la agregación.

5. Significado y Reivindicaciones

El artículo sostiene que los enfoques propuestos abordan eficazmente los desafíos duales de la heterogeneidad y los datos limitados en el Aprendizaje Federado. Al utilizar la distribución global como un prior, los modelos evitan los problemas de asumir priors incompatibles mientras aprovechan la robustez de la inferencia Bayesiana.

Los autores enfatizan que su análisis teórico establece la optimalidad minimax, proporcionando una base rigurosa para las mejoras empíricas observadas. Afirman que cFedBayes ofrece un avance significativo al proporcionar garantías adaptativas a nivel de grupo en lugar de un límite de "talla única". Finalmente, posicionan a sFedBayes como una solución práctica para entornos con recursos limitados al reducir los costos de inferencia sin sacrificar la precisión.

El trabajo se presenta como una extensión de una versión previa de conferencia, con mejoras mayores que incluyen la introducción de la dispersión (sFedBayes), el desarrollo de un marco agrupado con límites a nivel de grupo (cFedBayes) y una extensa validación experimental contra modelos SOTA actuales.