Resumen Técnico: Cada Cliente es un Entorno: De-confounding Federado para el Pronóstico Espacio-Temporal

1. Planteamiento del Problema

El pronóstico espacio-temporal (STF, por sus siglas en inglés) es crítico para aplicaciones de computación urbana como la predicción de tráfico y el monitoreo de la calidad del aire. Aunque los modelos de aprendizaje profundo han logrado un alto rendimiento, típicamente requieren la recolección centralizada de datos, lo cual suele ser inviable debido a regulaciones de privacidad, restricciones de comunicación y problemas de propiedad de los datos. El Aprendizaje Federado (FL) ofrece una solución al entrenar modelos a través de clientes distribuidos sin compartir los datos brutos.

Sin embargo, los métodos existentes de STF federado tratan la heterogeneidad entre clientes principalmente como un desafío de optimización que debe mitigarse mediante modelos personalizados. Los autores argumentan que esta perspectiva es errónea porque la heterogeneidad surge fundamentalmente de diversas condiciones ambientales (por ejemplo, clima, eventos especiales, patrones de hora pico) específicas de la ubicación de cada cliente. En consecuencia, los modelos locales capturan patrones de pronóstico específicos del entorno que están entrelazados con efectos de confusión específicos del cliente. Estos modelos fallan al generalizar cuando las condiciones ambientales cambian entre ubicaciones o tiempos. Además, los métodos de inferencia causal para STF tienen dificultades en entornos federados porque carecen del acceso a los datos brutos centralizados necesarios para desenredar los confundidores ambientales.

2. Metodología: El Marco Atlas

El artículo propone Atlas, un novedoso marco de de-confounding federado. La idea central es tratar a cada cliente federado no como un estorbo, sino como un entorno causal distinto que proporciona observaciones complementarias del mismo sistema espacio-temporal subyacente.

2.1 Formulación Causal

Los autores formulan un Modelo Causal Estructural (SCM) para el STF federado:

• Variables: Observaciones históricas ($X_k$), horizontes futuros ($Y_k$) y condiciones ambientales específicas del cliente ($E_k$).
• Confusión (Confounding): $E_k$ introduce correlaciones espurias a través de la ruta de puerta trasera (back-door path) $X_k \leftarrow E_k \rightarrow Y_k$.
• La Limitación del Ajuste Local: Para eliminar la confusión mediante el ajuste de puerta trasera, se debe estimar $P(Y|do(X)) = \sum P(Y|X, E=e)P(E=e)$. En un entorno federado, los clientes individuales solo observan un subconjunto sesgado de estratos ambientales ($P_k(E)$) y carecen de la distribución marginal global $P(E)$. Esto conduce a dos sesgos específicos:
  1. Estratos Faltantes: Los clientes pierden regímenes ambientales que no están presentes en sus datos locales.
  2. Pesos Sesgados: Las estimaciones locales están ponderadas por distribuciones locales en lugar de la distribución global.

2.2 El Mecanismo Atlas

Atlas aborda estos problemas aprovechando la agregación federada como un ajuste de puerta trasera implícito.

1. Libro de Códigos de Prototipos (Prototype Codebook): El marco aprende un libro de códigos de prototipos global compartido $Q = [q_1, \dots, q_I]$, donde cada prototipo $q_i$ corresponde a un estrato ambiental latente.
2. Codificación Local: Cada cliente $k$ codifica su entrada $X_k$ en una representación latente $z_k$ utilizando un backbone personalizado. Luego, mapea $z_k$ al libro de códigos mediante un mecanismo de asignación suave (Eq. 5), produciendo una representación condicionada por el entorno $\hat{z}_k$. Esto actúa como una aproximación a nivel de parámetros del ajuste de puerta trasera local.
3. Alineación y Agregación: Dado que los espacios de prototipos son intercambiables (por ejemplo, el espacio $i$ en el Cliente A podría corresponder a un régimen diferente que el espacio $i$ en el Cliente B), el servidor realiza una alineación de prototipos antes de la agregación.
   • El servidor calcula una matriz de costos basada en distancias de coseno entre los prototipos de los clientes y los prototipos de anclaje globales.
   • Resuelve un problema de asignación lineal (usando el algoritmo húngaro) para encontrar la permutación óptima $\Pi_k$ que alinea los espacios del cliente con los espacios de anclaje globales.
   • El servidor agrega los libros de códigos alineados: $Q \leftarrow \sum w_k \Pi_k Q_k$.
4. Garantía Teórica: Los autores derivan un límite de de-confounding federado, demostrando que el error está controlado linealmente por la confusión promedio entre los clientes, en lugar de la fuerza de cualquier cliente individual. Esto implica que los diversos entornos de los clientes se cancelan mutuamente los sesgos específicos al agregarse.

2.3 Implementación

• Arquitectura: Los clientes utilizan un codificador local (convolución a nivel de ventana + ResMLP) y un decodificador. Solo se intercambia el libro de códigos de prototipos ligero ($I \times D$ parámetros), no los pesos completos del modelo.
• Entrenamiento: En cada ronda, los clientes actualizan sus codificadores locales, decodificadores y libros de códigos locales. Suben el libro de códigos actualizado al servidor, el cual lo alinea y agrega para formar el nuevo libro de códigos global.

3. Contribuciones Clave

1. Un Marco de De-confounding Federado: El artículo introduce una perspectiva causal que trata a los clientes como entornos causales distribuidos. Atlas aprende efectos de confusión ambiental compartidos a través de clientes heterogéneos para mejorar la generalización.
2. La Federación como Ajuste de Puerta Trasera Implícito: Teóricamente, el artículo revela que la agregación de libros de códigos de prototipos establece las condiciones para mitigar la confusión ambiental. Deriva un límite que muestra que el error de de-confounding depende de la fuerza de confusión promedio, no del sesgo de un cliente individual.
3. Validación Empírica: Experimentos extensos en cinco conjuntos de datos del mundo real (METR-LA, PEMS-BAY, PEMS03, PEMS04, KnowAir) demuestran que Atlas supera a las líneas base federadas y centralizadas existentes. Se demuestra que el marco es transferible, interpretable y eficiente en comunicación.

4. Resultados Experimentales

4.1 Rendimiento de Pronóstico

Atlas supera consistentemente a todas las líneas base federadas (incluyendo FedAvg, FedProx, FedPer, FedGRU, FedGTP y métodos especializados como Fuels y FedDis) en los cinco conjuntos de datos.

• En METR-LA, Atlas reduce el MAE a 3.30 en comparación con la mejor línea base federada (FedGTP con 3.54) y se acerca a los métodos centralizados (DCRNN con 3.05) dentro de una brecha del 4.9%–6.7%.
• En PEMS-BAY, Atlas logra un MAE de 1.68, superando a la mejor línea base federada (FedGTP con 1.73).

4.2 Estudios de Ablación

• Agregación Federada: Eliminar la agregación (entrenamiento solo local) provoca una caída masiva en el rendimiento (el MAE aumenta un 77% en METR-LA), confirmando que un solo cliente no puede capturar regímenes ambientales estables.
• Arquitectura Compartida: Forzar un codificador/decodificador unificado entre clientes (A.2, A.3) degrada el rendimiento, indicando que el modelado temporal específico del cliente debe preservarse.
• Alineación: Eliminar el paso de alineación de prototipos (A.4) conduce a caídas consistentes en el rendimiento, probando la necesidad de alinear la semántica ambiental antes de la agregación.
• Asignación Suave: Reemplazar la asignación suave con asignación dura (A.5) degrada el rendimiento, resaltando la necesidad de modelar transiciones suaves entre regímenes.

4.3 Transferibilidad y Eficiencia

• Transferibilidad: El libro de códigos aprendido se transfiere eficazmente a clientes no vistos. Al pre-entrenarse con solo el 5% de los clientes, el libro de códigos congelado rinde de manera comparable a un modelo entrenado desde cero con ese mismo subconjunto, y significativamente mejor que una inicialización aleatoria.
• Eficiencia: Atlas logra un equilibrio favorable entre computación y comunicación. Los clientes suben solo el libro de códigos (por ejemplo, 640 parámetros), lo cual es dos órdenes de magnitud más pequeño que métodos como FedAvg o FedGRU que suben los pesos completos del modelo.
• Robustez: El libro de códigos es robusto al ruido Gaussiano y Rademacher, con una degradación mínima del rendimiento, lo que sugiere que captura semánticas ambientales estables.

4.4 Interpretabilidad

La visualización en el conjunto de datos KnowAir muestra que las asignaciones de prototipos se correlacionan con regímenes ambientales distintos (por ejemplo, periodos limpios vs. contaminados) y condiciones meteorológicas (temperatura, presión, humedad), no solo con la variable objetivo (PM2.5). Esto confirma que el libro de códigos captura contextos ambientales latentes.

5. Significado y Reivindicaciones

El artículo afirma que Atlas representa un cambio de paradigma en el STF federado:

• De Estorbo a Recurso: En lugar de suprimir la heterogeneidad entre clientes como un obstáculo de optimización, Atlas la explota como evidencia distribuida de regímenes ambientales.
• Generalización Causal: Al enmarcar teóricamente la federación como un mecanismo de ajuste de puerta trasera, el método permite que los modelos generalicen bajo cambios ambientales donde los métodos tradicionales de FL personalizado fallan.
• Viabilidad Práctica: El marco ofrece una solución práctica para el pronóstico espacio-temporal sensible a la privacidad, proporcionando un alto rendimiento con un mínimo de sobrecarga de comunicación y representaciones ambientales interpretables.

Los autores reconocen limitaciones, señalando que los prototipos discretos pueden tener dificultades con cambios continuos o extremadamente de alta dimensión, y que la complejidad de alineación $O(I^3)$ podría ser un cuello de botella para escalas de prototipos muy grandes. El trabajo futuro pretende explorar estrategias de alineación escalables y libros de códigos dinámicos.