Resumen Técnico: Alineación de Distribuciones para Aprendizaje Federado de Un Solo Paso mediante Transporte Óptimo

1. Planteamiento del Problema

El artículo aborda los desafíos inherentes al Aprendizaje Federado de Un Solo Paso (OSFL, por sus siglas en inglés), un régimen de comunicación donde los clientes interactúan con un servidor central exactamente una vez. El Aprendizaje Federado (FL) estándar mitiga la heterogeneidad de los datos mediante la optimización iterativa cliente-servidor, pero el OSFL carece de esta capacidad de refinamiento.

El problema central identificado es el efecto acumulativo de dos tipos específicos de heterogeneidad estadística:

1. Desplazamiento de Dominio (Domain Shift): Los clientes observan diferentes distribuciones de entrada $P_k(x)$ debido a variaciones en las condiciones de adquisición o entornos.
2. Desplazamiento de Etiquetas (Label Shift): Los clientes exhiben diferentes marginales de etiquetas $P_k(y)$, lo que conduce a proporciones de clases variables.

En ausencia de corrección iterativa, estos desplazamientos inducen representaciones de características desalineadas y funciones de decisión de posterior de clase inconsistentes entre los clientes. Los métodos de OSFL existentes (que dependen de la destilación, la síntesis de datos o el ensamble) a menudo asumen representaciones alineadas o abordan estos desplazamientos de forma separada. En consecuencia, las discrepancias en la geometría de las características y los sesgos condicionales de clase permanecen sin corregir, lo que degrada el rendimiento del modelo global agregado.

2. Metodología: SLOT-Align

Los autores proponen SLOT-Align (Alineación de Transporte Óptimo de un solo round, libre de aprendizaje), un marco de trabajo sensible a la geometría diseñado para armonizar las distribuciones de características de los clientes antes del entrenamiento de OSFL descendente. El método opera como un paso de preprocesamiento y se basa en un codificador congelado compartido $f_\theta$ (preentrenado en un conjunto de datos externo) para mapear todas las entradas locales en un espacio latente común.

El flujo de trabajo consta de cuatro etapas clave:

A. Extracción de Estadísticas de Características Locales

Cada cliente $k$ calcula estadísticas compactas de primer y segundo orden de las características $z = f_\theta(x)$ extraídas de sus datos locales:

• Media: $\mu_k = \mathbb{E}[f_\theta(x)]$
• Covarianza: $\Sigma_k$ (regularizada mediante contracción de Ledoit–Wolf para asegurar la definición positiva).
  Estas estadísticas sirven como un proxy Gaussiano $N(\mu_k, \Sigma_k)$ para la distribución de características inducida, capturando las discrepancias dominantes en ubicación y dispersión causadas por los desplazamientos de dominio y de etiqueta.

B. Construcción de la Referencia Global

El servidor agrega estas estadísticas locales para construir una distribución de referencia global:

• Media Global: Un promedio ponderado de las medias de los clientes.
• Covarianza Global: Calculada como el baricentro de Bures–Wasserstein de las covarianzas de los clientes. Esto se resuelve mediante una iteración de punto fijo para encontrar el minimizador único de la suma de las distancias de Bures al cuadrado, equilibrando efectivamente la variabilidad específica de cada cliente con la consistencia global.

C. Mapeo de Transporte Óptimo Geodésico

Cada cliente construye un mapa de transporte para alinear su aproximación Gaussiana local con la referencia global.

• Mapa Óptimo: Bajo un costo cuadrático, el mapa de transporte óptimo entre dos Gaussianas es afín y tiene una solución de forma cerrada que involucra las raíces cuadradas de las matrices de covarianza.
• Interpolación: Para evitar la sobrecorrección (dado que la verdadera distribución de características no es estrictamente Gaussiana), el método introduce un parámetro de interpolación $\tau \in [0, 1]$. La transformación final $T^{(\tau)}_k$ es una combinación convexa del mapa de identidad y el mapa de transporte óptimo completo. Esto crea una interpolación de desplazamiento de Wasserstein, moviendo la distribución del cliente a lo largo de una geodésica de velocidad constante hacia la referencia global.

D. Integración

Los clientes aplican la transformación $z' = T^{(\tau)}_k(z)$ a sus características locales. Estas características alineadas se introducen luego en cualquier algoritmo de OSFL existente (por ejemplo, FedAvg, FedCGS, FedPFT) sin modificar los procedimientos de optimización descendentes.

3. Contribuciones Clave

1. Formalización: El artículo formaliza el OSFL como un problema de alineación de distribuciones bajo un desplazamiento conjunto de dominio y etiqueta, destacando cómo estos desplazamientos limitan la eficacia de los enfoques de un solo paso.
2. Propuesta de Algoritmo: Introduce SLOT-Align, un marco de trabajo basado en transporte óptimo y libre de aprendizaje, que armoniza las distribuciones de características de los clientes mediante un único intercambio de estadísticas compactas. No requiere datos auxiliares, generadores entrenables ni optimización iterativa.
3. Validación Empírica: Experimentos extensos demuestran que SLOT-Align mejora consistentemente la precisión y la robustez a través de múltiples benchmarks (Office-Home, Digits, DomainNet), diversos backbones preentrenados (ViT, ResNet) y baselines de OSFL de última generación.

4. Resultados Experimentales

Los autores evaluaron SLOT-Align bajo distintos grados de desplazamiento de etiqueta (controlado por el parámetro de concentración de Dirichlet $\alpha$) y diversidad de dominio.

• Ganancias de Rendimiento: Bajo desplazamiento conjunto de dominio y etiqueta ($\alpha=0.1$), SLOT-Align mejoró consistentemente la precisión promedio macro para todos los baselines probados (O-FedAvg, FedCGS, FedPFT). Por ejemplo, en el conjunto de datos Office-Home, mejoró O-FedAvg en un +7.86% y FedPFT en un +6.55%.
• Robustez: El método se mantuvo efectivo bajo un sesgo de etiqueta más fuerte ($\alpha=0.05$ y $\alpha=0.01$), mostrando ganancias incluso en regímenes de heterogeneidad severa.
• Ablación sobre la Fuerza de Alineación: Se encontró que el parámetro de interpolación $\tau$ es crítico. Los valores intermedios (por ejemplo, $\tau=0.4$) produjeron el mejor equilibrio, mientras que $\tau=1.0$ (transporte completo) a veces degradó el rendimiento, sugiriendo que una alineación parcial es preferible a una corrección total.
• Desplazamiento de Solo Dominio: En configuraciones con etiquetas balanceadas (solo desplazamiento de dominio), SLOT-Align proporcionó ganancias marginales o un rendimiento neutro, confirmando que el método apunta específicamente a la interacción compleja de los desplazamientos de dominio y de etiqueta.
• Eficiencia: El costo de comunicación es mínimo ($O(m + m^2)$ por cliente para media y covarianza), y la sobrecarga computacional es insignificante en comparación con la extracción de características y el entrenamiento descendente.

5. Significado y Reivindicaciones

El artículo afirma que SLOT-Align ofrece una solución modular y computacionalmente eficiente para un cuello de botella crítico en OSFL: la incapacidad de corregir el desalineamiento de la distribución en un solo round de comunicación.

• Modularidad: No reemplaza los flujos de trabajo de OSFL existentes, sino que actúa como una capa de preprocesamiento compatible, lo que lo hace aplicable a una amplia gama de métodos actuales y futuros que dependen de codificadores congelados.
• Sensibilidad a la Geometría: Al aprovechar la geometría del espacio 2-Wasserstein y los baricentros de Bures–Wasserstein, corrige explícitamente las discrepancias tanto en la ubicación (media) como en la forma (covarianza) de las distribuciones de características.
• Practicidad: El método evita los riesgos de privacidad y la sobrecarga computacional asociados con la síntesis de datos o la destilación iterativa, basándose únicamente en resúmenes estadísticos compactos.

Los autores concluyen que, si bien el OSFL está limitado por presupuestos estrictos de comunicación, SLOT-Align mitiga eficazmente la degradación del rendimiento causada por las distribuciones de datos heterogéneas, particularmente cuando los desplazamientos de dominio y de etiqueta interactúan para distorsionar las estadísticas de las características.