Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

Aquí se presenta un resumen técnico detallado del artículo "Optimización de subespacio para aprendizaje federado eficiente bajo datos heterogéneos" de Zhu et al.

1. Planteamiento del Problema

El artículo aborda el "triple desafío" que enfrenta el Aprendizaje Federado (FL) en la era de los modelos a gran escala (cientos de millones a miles de millones de parámetros):

1. Cómputo: Los pasos de entrenamiento local son computacionalmente pesados.
2. Memoria: Almacenar estados completos del modelo y variables auxiliares en dispositivos periféricos con recursos limitados es prohibitivo.
3. Comunicación: Transmitir actualizaciones de dimensión completa a través de un ancho de banda limitado es ineficiente.

Estos desafíos se ven exacerbados por la heterogeneidad de datos (No-IID). En configuraciones No-IID, los objetivos locales se desvían del objetivo global, provocando una "deriva del cliente" (client drift).

• Soluciones existentes y limitaciones:
  • Corrección de heterogeneidad (p. ej., SCAFFOLD): Utiliza variables de control para corregir la deriva, pero requiere almacenar y comunicar estados auxiliares de dimensión completa, anulando las ganancias de eficiencia para modelos grandes.
  • Compresión (p. ej., cuantización): Reduce los bits transmitidos, pero a menudo requiere búferes de error de dimensión completa, fallando en reducir la huella de memoria.
  • Métodos de subespacio/bajo rango: Reducen la dimensionalidad para la eficiencia, pero carecen típicamente de mecanismos fundamentados para corregir la deriva del cliente. Las combinaciones ingenuas a menudo descartan información histórica de corrección cuando el subespacio cambia, lo que lleva a inestabilidad.

La brecha: No existe ningún algoritmo que realice una corrección de heterogeneidad nativa dentro de un subespacio de baja dimensión mientras preserva la robustez de los métodos de dimensión completa como SCAFFOLD.

2. Metodología: Subespacio-SCAFFOLD (SSF)

Los autores proponen SSF, un algoritmo que integra la corrección de heterogeneidad directamente en un subespacio de optimización de baja dimensión.

Mecanismos principales

1. Proyección de subespacio:

   • En cada ronda $t$, se genera una matriz de proyección ortogonal aleatoria compartida $P_t \in \mathbb{R}^{r \times d}$ (donde $r \ll d$).
   • La optimización (actualizaciones locales y agregación) ocurre estrictamente en el subespacio de $r$ dimensiones.
   • Los gradientes se proyectan: $g_{proj} = P_t g$.

2. Variables de control proyectadas (Corrección de heterogeneidad):

   • SSF imita el mecanismo de SCAFFOLD pero opera sobre cantidades proyectadas.
   • Los clientes mantienen variables de control locales $c_i$ y el servidor mantiene una variable de control global $c$.
   • La dirección de actualización local se corrige utilizando la diferencia entre las variables de control locales y globales proyectadas: $P_t g - c_{i, proj} + c_{proj}$.

3. Relleno de preservación de residuos (Innovación clave):

   • A diferencia de los métodos de subespacio estándar que descartan información fuera del subespacio activo, SSF mantiene un estado de control de dimensión completa ($c_i$ y $c$).
   • Regla de actualización: Cuando el subespacio cambia, las variables de control se actualizan como:
     $$c^{t+1} = (I - P_t^\top P_t)c^t + P_t^\top P_t (\text{nueva información de gradiente})$$
   • Esto significa que el complemento ortogonal (residuo) de la variable de control se preserva sin cambios, mientras que solo se actualiza el componente del subespacio activo.
   • Relleno (Backfill): Después de la actualización del subespacio, el modelo de espacio completo se reconstruye combinando las nuevas coordenadas proyectadas con el residuo preservado: $x^{t+1} = P_t^\top x^{t+1}_{proj} + x^t_{res}$.

Flujo del algoritmo

1. Descomposición: El servidor divide el modelo y las variables de control en partes proyectadas y residuales.
2. Pasos locales: Los clientes reconstruyen el modelo completo utilizando el modelo de subespacio transmitido y el residuo del servidor. Calculan gradientes en el espacio completo, los proyectan y realizan actualizaciones corregidas por heterogeneidad en el subespacio.
3. Agregación: El servidor agrega los puntos finales proyectados en el subespacio.
4. Relleno y actualización: El servidor eleva el modelo de nuevo al espacio completo y actualiza las variables de control, preservando los componentes residuales.

3. Contribuciones clave

1. Diseño de algoritmo (SSF):

   • El primer algoritmo de FL que realiza corrección de heterogeneidad nativa en un subespacio de baja dimensión.
   • Introduce una actualización estilo relleno que retiene los componentes residuales de las variables de control, asegurando que la información histórica de corrección de deriva no se pierda cuando el subespacio rota o cambia.

2. Análisis teórico:

   • Demuestra que, bajo supuestos estándar de suavidad y varianza acotada, SSF logra una tasa de convergencia no asintótica de $\tilde{O}(1/T + 1/\sqrt{NKT})$.
   • Esta tasa coincide con la aceleración lineal de SCAFFOLD de dimensión completa, demostrando que las restricciones de subespacio no degradan el comportamiento de convergencia asintótica en relación con la deriva del cliente.
   • El análisis se mantiene sin asumir independencia entre el proyector aleatorio y los gradientes estocásticos.

3. Validación empírica:

   • Demostrado tanto en un problema de juguete controlado de regresión matricial como en un punto de referencia de aprendizaje profundo (CIFAR-100 con ResNet-110).
   • Mostró que SSF reduce significativamente los costos de memoria, comunicación y cómputo, manteniendo una precisión comparable a la de SCAFFOLD de dimensión completa.

4. Resultados experimentales

Problema de juguete (Regresión matricial)

• Robustez: Bajo niveles variables de heterogeneidad (bajo, medio, alto), SSF superó consistentemente a FedSub (un método de subespacio de referencia sin preservación de residuos) y a FedAvg.
• Estabilidad: FedSub exhibió inestabilidad numérica y divergencia (errores NaN) en dimensiones de subespacio más altas ($r=50$) debido a la pérdida de información ortogonal. SSF permaneció estable en todas las dimensiones.
• Precisión: Con una relación de subespacio de $r/d = 0.2$, SSF logró errores relativos muy cercanos a Full-SCAFFOLD, significativamente mejores que Full-FedAvg.

Aprendizaje profundo (CIFAR-100 / ResNet-110)

• Clasificación: Full-SCAFFOLD > SSF > FedAvg > FedSub.
• Rendimiento: SSF logró una precisión de prueba final de 45.43%, superando a FedAvg (34.35%) y FedSub (23.17%) por un amplio margen.
• Eficiencia: SSF redujo con éxito la sobrecarga de comunicación y memoria mientras preservaba la robustez de SCAFFOLD frente a la heterogeneidad de datos.

5. Significado e impacto

• Cerrando la brecha: SSF resuelve la tensión fundamental entre robustez (manejo de datos No-IID) y eficiencia (manejo de modelos grandes). Demuestra que no se necesitan estados auxiliares de dimensión completa para corregir eficazmente la deriva del cliente.
• Escalabilidad: Al reducir la huella de comunicación y memoria de $O(d)$ a $O(r)$ (donde $r \ll d$) mientras se mantienen garantías de convergencia teóricas, SSF hace viable el Aprendizaje Federado de modelos grandes en dispositivos periféricos con recursos limitados.
• Perspectiva teórica: El mecanismo de "preservación de residuos" proporciona un nuevo paradigma para la optimización de subespacio, mostrando que mantener memoria de dimensión completa para las variables de control (incluso si solo se accede parcialmente) es crucial para la estabilidad en entornos de subespacio dinámicos.

En conclusión, SSF ofrece una solución práctica y teóricamente sólida para desplegar Aprendizaje Federado a gran escala y heterogéneo, logrando "lo mejor de ambos mundos": la eficiencia de los métodos de subespacio y la robustez de la corrección mediante variables de control.