Resumen Técnico: Aprendizaje Federado con Restricción de Subespacio mediante Adaptación de Bajo Rango

1. Planteamiento del Problema

El Aprendizaje Federado (FL, por sus siglas en inglés) permite el entrenamiento colaborativo a través de clientes descentralizados sin compartir datos brutos, pero enfrenta desafíos significativos cuando las distribuciones de datos de los clientes son heterogóneas. En el promedio federado estándar (FedAvg), la heterogeneidad estadística conduce al "desvío del cliente" (client drift), donde las actualizaciones de los modelos locales divergen, causando inestabilidad y una convergencia lenta.

Con el auge de los modelos fundacionales de gran escala, la Adaptación de Bajo Rango (LoRA) se ha convertido en una técnica popular de ajuste fino eficiente en parámetros para el FL. LoRA restringe las actualizaciones de los pesos ($\Delta W$) a una descomposición de bajo rango ($BA$), reduciendo los costos de comunicación y computación. Sin embargo, este artículo plantea la hipótesis de que, bajo distribuciones de datos heterogéneas, diferentes clientes aprenden actualizaciones de bajo rango que abarcan subespacios distintos. Cuando estas actualizaciones geométricamente desalineadas se agregan de forma ingenua en el servidor, pueden cancelarse parcialmente entre sí, un fenómeno denominado "cancelación silenciosa". Esta cancelación estructural reduce la magnitud efectiva de la actualización global, obstaculizando la convergencia incluso cuando los clientes individuales realizan un progreso sustancial local.

2. Metodología

Los autores proponen LoRA con Regularización de Subespacio (Subspace-Reg), un marco diseñado para alinear las actualizaciones locales de los clientes con un subespacio de referencia global compartido.

2.1 Función Objetivo

Cada cliente $k$ optimiza una función objetivo regularizada por subespacio que incluye una penalización de alineación proximal. La función es:
$$L_k^{SR}(w_k) = \min_{A_k, B_k} \left\{ L_k(w_k) + \frac{\mu}{2} \|B_k A_k - \Delta W_g^{(t)}\|_F^2 + \frac{\lambda}{2} (\|A_k\|_F^2 + \|B_k\|_F^2) \right\}$$
Donde:

• $L_k(w_k)$ es la pérdida local.
• $\Delta W_g^{(t)}$ es la actualización de bajo rango global actual (referencia del servidor).
• $\mu$ controla la fuerza de la alineación del subespacio.
• $\lambda$ regulariza los factores de bajo rango.
• El término de la norma de Frobenius restringe las actualizaciones locales para que permanezcan cerca de la actualización de referencia del servidor, fomentando una agregación constructiva.

2.2 Algoritmo

El entrenamiento sigue una estructura estándar de FL con modificaciones:

1. Inicialización: El servidor inicializa una base de bajo rango global $B_g$ con columnas ortonormales.
2. Paso del Servidor: El servidor selecciona clientes, calcula una actualización de referencia global mediante la SVD de la suma ponderada de las actualizaciones de clientes anteriores, y transmite el modelo y la base de referencia.
3. Paso del Cliente: Cada cliente inicializa sus factores locales con la referencia global y realiza $E$ épocas locales. El descenso de gradiente local minimiza la pérdida regularizada por subespacio, la cual incluye:
   • La pérdida de la tarea principal.
   • Un término que penaliza la desviación del producto $B_k A_k$ respecto a la referencia global $B_g A_g$.
   • Un término opcional que penaliza la desviación de la base $B_k$ respecto a $B_g$.
4. Agregación: Los clientes envían sus actualizaciones ($\Delta W_k = B_k A_k$) al servidor, el cual las agrega y actualiza el modelo global.

3. Análisis Teórico

El artículo proporciona un análisis teórico para comprender cómo la alineación de subespacios influye en la convergencia bajo heterogeneidad.

• Límite de Cancelación Geométrica: Los autores derivan que el cuadrado de la norma de la actualización agregada depende de los términos cruzados entre las actualizaciones de los clientes. Bajo heterogeneidad, si las actualizaciones están desalineadas, estos términos cruzados se vuelven negativos, reduciendo la magnitud efectiva de la actualización (cancelación silenciosa).
• Garantías de Convergencia: Bajo supuestos estándar (suavidad, heterogeneidad acotada, varianza acotada y actualizaciones de bajo rango acotadas), el artículo demuestra un límite de convergencia para la norma cuadrada del gradiente promedio.
  • El límite incluye un término $C_1 \sigma^2 / \mu^2$ que representa la penalización por desalineación.
  • A diferencia del LoRA-FedAvg estándar, que sufre de un término de cancelación fijo, Subspace-Reg permite que esta penalización se haga arbitrariamente pequeña aumentando el parámetro de alineación $\mu$.
  • Los autores derivan condiciones para $\mu$ y el número de rondas de comunicación $T$ para alcanzar una precisión de gradiente objetivo $\epsilon$.

4. Resultados Experimentales

El método fue evaluado en dos modelos preentrenados (RoBERTa-large y SmolLM-360M) utilizando el conjunto de datos HellaSwag en un entorno no-IID (10 clientes, Dirichlet $\beta=0.5$). Los experimentos se realizaron en 3 semillas aleatorias (42, 43, 44), para un total de 24 ejecuciones.

4.1 Rendimiento en RoBERTa-large

Subspace-Reg superó significativamente a todos los modelos de referencia (FedAvg, redistribución SVD y FedSVD):

• Precisión (Accuracy): Logró la mayor precisión media máxima (0.454 ± 0.023) y la precisión media final (0.429 ± 0.011), superando al mejor modelo de referencia (FedSVD) por un margen de ~0.118 en precisión máxima.
• Pérdida (Loss): Logró la menor pérdida final (1.363).
• Alineación: Logró un Solapamiento de Base (Basis Overlap) casi perfecto (≈0.9999), comparado con el 0.958–0.991 de los modelos de referencia.

4.2 Rendimiento en SmolLM-360M

Los resultados en SmolLM-360M revelaron que las ganancias de precisión dependen del modelo:

• Precisión (Accuracy): FedAvg logró la mayor precisión (0.581), mientras que Subspace-Reg quedó en segundo lugar (0.564).
• Alineación: A pesar de la falta de mejora en la precisión, Subspace-Reg mantuvo un Solapamiento de Base casi perfecto (0.99989), consistente con los resultados de RoBERTa.
• Conclusión: La alineación geométrica fue exitosa, pero no se tradujo en ganancias de precisión para la tarea específica de este modelo.

4.3 Hallazgos Diagnósticos Clave

• Solapamiento de Base vs. Similitud de Coseno: Mientras que los métodos basados en SVD lograron métricas basadas en el coseno más altas (Similitud de Coseno Cliente-Global, Similitud de Coseno Par), no lograron las mismas mejoras de precisión que Subspace-Reg en RoBERTa. Los autores identifican el Solapamiento de Base (alineación directa de subespacios) como el diagnóstico crítico, en lugar de la similitud de coseno.
• Robustez: El solapamiento de base casi perfecto fue consistente en todas las semillas y ambos modelos, validando la hipótesis de la alineación geométrica incluso cuando el rendimiento de la tarea variaba.

5. Significación y Reivindicaciones

El artículo reivindica las siguientes contribuciones y significados:

1. Identificación de la Cancelación Silenciosa: Destaca que la desalineación de subespacios en el FL basado en LoRA conduce a una agregación destructiva, un fenómeno previamente poco explorado.
2. Objetivo Regularizado por Subespacio: Propone una penalización de alineación proximal que fomenta explícitamente que las actualizaciones locales permanezcan cerca de un subespacio de referencia global, extendiendo el concepto de FedProx de la magnitud de los parámetros a la estructura geométrica.
3. Mejora Teórica: Proporciona un límite de convergencia que muestra que la penalización por desalineación es ajustable, ofreciendo una garantía estrictamente mejor que el LoRA-FedAvg estándar bajo heterogeneidad.
4. Validación Empírica: Demuestra que el método propuesto logra un solapamiento de base casi perfecto (≈0.9999) de manera consistente. En RoBERTa-large, esto se traduce en ganancias sustanciales de precisión. En SmolLM-360M, confirma que la alineación geométrica puede lograrse independientemente de la precisión de la tarea, sugiriendo que los beneficios de la alineación dependen del modelo.
5. Perspectiva Diagnóstica: El trabajo establece que el solapamiento de base es un indicador más relevante de la calidad de la agregación que la similitud de coseno en entornos de LoRA federado.

Los autores concluyen que, si bien la alineación de subespacios es un mecanismo geométrico robusto, su impacto en el rendimiento final de la tarea depende de la arquitectura y capacidad del modelo específico, lo que justifica una mayor investigación sobre por qué estos beneficios varían entre modelos.