Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

Resumen Técnico: FedCGNM y FedHOO para el Desequilibrio de Clases en Aprendizaje Federado

Planteamiento del Problema
El desequilibrio de clases sigue siendo un desafío crítico en el Aprendizaje Federado (FL), particularmente cuando la distribución global de etiquetas es de cola larga (long-tailed). En tales escenarios, las clases minoritarias sufren un desempeño predictivo deficiente debido a su subrepresentación en los datos de entrenamiento. Las técnicas centralizadas estándar, como el reponderado de pérdida (loss reweighting), el muestreo avanzado o el aumento generativo, son difíciles de implementar en FL. Las restricciones de privacidad impiden el intercambio de datos, y los generadores sintéticos a menudo no logran producir muestras realistas para regímenes dispersos o tareas de dominio específico (por ejemplo, la detección de defectos en semiconductores). Además, los enfoques existentes a nivel de optimización, como la Normalización por Clase (PCN), enfrentan limitaciones en entornos de múltiples clases: son susceptibles al ruido direccional de las clases minoritarias e introducen desajustes de escala al sumar muchos vectores unitarios, lo que desestabiliza la convergencia a medida que aumenta el número de clases.

Metodología
El artículo propone un enfoque de dos vertientes para abordar estos problemas: un optimizador del lado del cliente llamado FedCGNM (Federated Class-Grouped Normalized Momentum) y un sintonizador de hiperparámetros auxiliar llamado FedHOO.

1. FedCGNM (Optimizador):

   • Agrupación de Clases (Class Grouping): En lugar de normalizar los gradientes para cada clase individual ($C$), FedCGNM particiona las clases en un pequeño número de grupos ($H \ll C$, típicamente mayoría frente a minoría). La regla de agrupación se deriva de la minimización de la varianza intra-grupo de las proporciones de las clases, una estrategia que ha demostrado ser un proxy confiable para la búsqueda exhaustiva de umbrales.
   • Momento Normalizado por Grupo (Group-Wise Normalized Momentum): Para cada grupo, el algoritmo mantiene un vector de momento. Durante el entrenamiento local, se computan gradientes estocásticos por grupo y se actualiza el momento. Crucialmente, el vector de momento de cada grupo se normaliza a longitud unitaria (con una constante de estabilización $\delta$) antes de la sumatoria.
   • Dirección de Actualización: La dirección de actualización final es la suma de estos momentos de grupo normalizados. Este diseño iguala la magnitud del gradiente entre los grupos mayoritarios y minoritarios, al tiempo que mitiga el ruido direccional inherente a los gradientes de clases raras. Al operar sobre un puñado de grupos en lugar de clases individuales, evita el desajuste de escala y la inestabilidad asociados con PCN.
   • Alineación del Cliente: Al forzar que cada actualización de grupo tenga una norma unitaria, las direcciones de actualización agregadas entre clientes heterogéneos se vuelven más alineadas, abordando una causa principal de degradación del desempeño en FL.

2. FedHOO (Sintonizador de Hiperparámetros):

   • Contexto: Seleccionar tasas de remuestreo apropiadas no es trivial en FL debido al espacio de búsqueda combinatorio entre clientes y al riesgo de desestabilizar la optimización con cronogramas fluctuantes.
   • Mecanismo: FedHOO es un algoritmo basado en bandidos de brazo X (X-armed bandit) diseñado para federaciones de pequeña escala (donde $K$ es pequeño). Explota la naturaleza lineal de la agregación en FL. En cada ronda de comunicación, el servidor instruye a cada cliente para que entrene con dos tasas de remuestreo candidatas.
   • Eficiencia: Mediante la combinación lineal de las actualizaciones de los clientes devueltas, el servidor puede inferir el desempeño de validación de todas las $2^K$ combinaciones posibles de tasas locales sin requerir $2^K$ rondas de entrenamiento separadas. Esto permite la identificación rápida de tasas de muestreo efectivas en las etapas tempranas del entrenamiento, manteniendo la privacidad (los clientes solo devuelven métricas de validación escalares).

Análisis Teórico
El artículo proporciona un análisis de convergencia para FedCGNM que considera explícitamente las tasas de remuestreo variables en el tiempo. Bajo supuestos estándar de suavidad, gradientes acotados e imparcialidad, los autores demuestran que FedCGNM alcanza una tasa de punto estacionario estándar de $O(T^{-1/2})$. El análisis destaca que, para que la convergencia se cumpla, la variación acumulada de las tasas de muestreo a través de las rondas ($V_T$) debe satisfacer condiciones específicas, lo que implica que un remuestreo agresivo pero inestable puede dificultar la convergencia.

Contribuciones Clave

1. Optimizador FedCGNM: El primer optimizador del lado del cliente que agrupa clases y aplica impulso de norma unitaria por grupo. Equilibra la influencia de la mayoría/minoría, reduce el ruido y mitiga el desajuste de escala.
2. Agrupación Consciente de la Varianza (Variance-Aware Grouping): Una estrategia de principios ligeros que enmarca la partición de clases como la minimización de la varianza intra-grupo, produciendo divisiones óptimas comparables a la búsqueda exhaustiva.
3. Análisis de Convergencia: Un marco teórico novedoso que combina FL, normalización de grupo y tasas de muestreo dinámicas, demostrando el impacto de los cronogramas de tasas en la convergencia.
4. FedHOO: El primer algoritmo basado en la optimización para determinar tasas de muestreo local combinatorias en FL utilizando un marco de bandidos de brazo X que aprovecha el paralelismo federado para una exploración eficiente.

Resultados Experimentales
Los métodos fueron evaluados en cuatro benchmarks de cola larga públicos (CIFAR-10-LT, CIFAR-100-LT, Adult Income, UNSW-NB15) y un conjunto de datos propietario de defectos de chips de semiconductores (CDD).

• Desempeño: FedCGNM superó consistentemente a las líneas base (incluyendo FedAvg, FedProx, CLIMB, FedGraB y variantes de PCN) en todos los entornos. En regímenes de pocos clientes ($K=5$), la combinación de FedCGNM con FedHOO produjo los mejores resultados, con mejoras de hasta un 29% sobre FedAvg con remuestreo.
• Aplicación Industrial: En el conjunto de datos CDD (tasa de defecto del 1.7%), el método propuesto logró una mejora del 16% sobre la mejor línea base, demostrando efectividad en entornos industriales reales y sensibles a la privacidad.
• Robustez: El enfoque mostró resiliencia en regímenes de muchos clientes con participación parcial y bajo diferentes grados de heterogeneidad de datos (Non-IID).
• Ablación: Los experimentos confirmaron que dos grupos ($H=2$) ofrecen el mejor equilibrio entre sesgo y varianza, y que factores de impulso moderados ($\beta \approx 0.5$) producen un desempeño óptimo.

Significancia y Reivindicaciones
El artículo afirma que FedCGNM y FedHOO forman un marco práctico para mitigar el desequilibrio de clases global en entornos de aprendizaje federado que preservan la privacidad, sin depender del intercambio de datos o la generación sintética. La significancia radica en abordar el desequilibrio de clases a nivel de optimización, lo cual es más generalmente aplicable que los trucos a nivel de datos restringidos por la privacidad o las características del conjunto de datos. Los autores enfatizan que su enfoque es particularmente efectivo en federaciones de pequeña escala (comunes en despliegues industriales) donde el ajuste conjunto de hiperparámetros es factible y crítico, mientras que la estrategia de agrupación asegura la estabilidad y alineación en escenarios de múltiples clases donde las técnicas de normalización previas fallan. El trabajo avanza en el campo al proporcionar tanto un optimizador robusto como un método eficiente y teóricamente fundamentado para ajustar estrategias de remuestreo en presencia de desequilibrio de clases.