Resumen Técnico: Evidencia Mecanística de Representaciones Preservadas pero Desalineadas en FedAvg No-IID

1. Planteamiento del Problema

Federated Averaging (FedAvg) es el procedimiento de optimización estándar para el Aprendizaje Federado (FL), pero sufre una degradación significativa del rendimiento bajo datos de clientes no-IID (heterogeneidad estadística), particularmente ante el sesgo de etiquetas (label skew). Las explicaciones existentes atribuyen este fallo al desvío del cliente (client drift), la inconsistencia de objetivos y la divergencia de pesos, caracterizando el proceso de optimización pero dejando el efecto interno de la agregación mayormente sin observar.

Una pregunta crítica sigue abierta: ¿Refleja la degradación en la precisión global la erradicación de las representaciones aprendidas por el cliente, o refleja un desalineamiento donde las estructuras internas se preservan pero no logran conectarse eficazmente con la vía de predicción final? Este artículo investiga esta cuestión utilizando herramientas de Interpretabilidad Mecanística (MI) para determinar si el FedAvg no-IID destruye el conocimiento aprendido o simplemente lo desconecta del clasificador.

2. Metodología

Los autores analizan modelos de visión entrenados en clientes dispersos (SimpleCNN y ResNet10) sobre CIFAR-10 y Fashion-MNIST. Para asegurar la interpretabilidad, los clientes entrenan bajo un esquema de dispersión basado en magnitud (atenuación hacia un 99.9% de dispersión), con controles de modelos densos que confirman que los hallazgos no son artefactos de la dispersión. El estudio emplea cuatro diagnósticos mecanísticos:

2.1. Descubrimiento de Circuitos

Utilizando un gating binario diferenciable, los autores identifican subredes mínimas (circuitos) suficientes para predecir clases específicas. Miden:

• Estabilidad intra-cliente: Solapamiento de circuitos a través de las rondas de comunicación.
• Consistencia inter-cliente: Similitud de circuitos entre diferentes clientes.
• Preservación local-a-global: La Intersección sobre la Unión (IoU) entre el circuito de un cliente y el circuito correspondiente en el modelo global agregado.
• Intervenciones: Suficiencia (precisión cuando solo el circuito está activo) y Necesidad (caída de precisión cuando el circuito es ablacionado).

2.2. Probing Lineal (Linear Probing)

Se entrena un clasificador lineal sobre las activaciones de la penúltima capa del modelo global. Esto pone a prueba si la información discriminativa de clase permanece linealmente accesible en el backbone, incluso si el cabezal (head) nativo del modelo no logra utilizarla.

2.3. Ajuste Fino Solo del Cabezal (Head-Only Finetuning)

El backbone se congela y solo el cabezal de clasificación final se ajusta finamente (finetuning) sobre un subconjunto IID pequeño y balanceado (200 muestras). Esto aísla si la pérdida de precisión se localiza en la capa de lectura o en el backbone de características.

2.4. Autoencoders Dispersos Universales (USAE)

Un USAE se entrena conjuntamente sobre las activaciones de modelos globales tanto IID como no-IID para aprender un diccionario de características compartido. Esto prueba si los dos modelos organizan sus características en una base compartida y permite la reconstrucción cruzada para evaluar la compatibilidad de las características.

3. Resultados Clave

3.1. Estabilidad y Divergencia de Circuitos

• Estabilidad Local: La dinámica de entrenamiento local es estable independientemente de la heterogeneidad de los datos. Los clientes convergen a circuitos altamente estables y específicos de clase (IoU $\approx$ 0.90) incluso bajo un sesgo extremo ($\alpha=0.05$).
• Divergencia Inter-Cliente: A medida que aumenta la heterogeneidad, los clientes aprenden diferentes circuitos para la misma clase. Bajo un sesgo extremo, el solapamiento de circuitos inter-cliente cae a un IoU $\approx$ 0.45.
• Preservación Parcial: El modelo global no promedia los circuitos hacia el ruido. La preservación local-a-global (IoU $\approx$ 0.55) sigue siendo mayor que el solapamiento inter-cliente, lo que indica que el modelo global hereda una estructura parcial de los clientes en lugar de converger a una solución única.

3.2. Los Circuitos Funcionales Persisten a pesar de la Precisión Cero

En entornos no-IID extremos, el modelo global puede alcanzar una precisión cercana a cero para clases específicas (por ejemplo, 0.0% para "Gato" o "Ciervo"). Sin embargo:

• Suficiencia: Los circuitos extraídos para estas clases logran una suficiencia casi perfecta (100% de precisión) cuando se aíslan.
• Necesidad: Ablar estos circuitos colapsa la precisión a casi cero.
• Conclusión: Las vías computacionales para estas clases existen y son causalmente responsables de la predicción, pero el modelo falla al enrutar las entradas a través de ellas durante la inferencia estándar.

3.3. Separabilidad Lineal y Desalineación del Cabezal

• Probes Lineales: Mientras que la precisión global cae significativamente bajo el sesgo (por ejemplo, del 68.3% al 43.1%), la precisión del probe lineal permanece alta y estable (cayendo solo al 64.7%). La brecha entre el probe y la precisión global se ensancha monótonamente con la heterogeneidad (+21.7 puntos en $\alpha=0.05$).
• Ajuste Fino del Cabezal: El ajuste fino únicamente del cabezal de clasificación sobre un conjunto pequeño y balanceado recupera una precisión sustancial (+10.2 puntos en $\alpha=0.05$), aunque no restaura completamente el rendimiento IID. Esto confirma que el backbone retiene características útiles, pero el cabezal por defecto está desalineado.

3.4. Base de Características Compartida (USAE)

• Solapamiento de Conceptos: El 98.7% de los conceptos del diccionario de 2,048 del USAE se activan de forma no trivial en ambos modelos (IID y no-IID), indicando una base de características mayoritariamente compartida.
• Transferencia Cruzada de Modelos: Cuando las activaciones no-IID se decodifican en el espacio de características IID, el cabezal IID funciona bien. Por el contrario, el cabezal no-IID falla incluso cuando se le alimentan activaciones IID. Esto confirma que las características están presentes y son compatibles, pero el mecanismo de lectura no-IID es el punto de fallo.

4. Significado y Reivindicaciones

El artículo argumenta que la degradación de FedAvg no-IID no se explica totalmente por la erradicación de la representación. En su lugar, el mecanismo principal es el desalineamiento:

1. Preservación: Las estructuras internas aprendidas por los clientes (circuitos) y las bases de características se preservan en gran medida en el modelo global agregado.
2. Desalineamiento: La vía de predicción final (el cabezal de clasificación) se desalinea con estas representaciones preservadas. El modelo "conoce" las clases pero no puede enrutar las entradas a los circuitos correctos ni utilizar las características de manera efectiva.

Implicaciones:

• Reencuadre del Problema: El problema no es necesariamente que los clientes "olviden" o que los pesos diverjan hacia cuencas incompatibles que destruyan el conocimiento, sino que el proceso de agregación interrumpe la relación estructural entre el backbone y el cabezal.
• Direcciones Algorítmicas:
  • Realineación del Cabezal Post-Agregación: El ajuste fino simple del cabezal sobre un conjunto pequeño y balanceado puede recuperar un rendimiento significativo sin cambiar el protocolo federado.
  • Agregación Consciente de Circuitos: Los métodos futuros podrían alinear las neuronas que implementan circuitos funcionales similares entre clientes antes de promediar (por ejemplo, mediante métodos de simetría de permutación) para reducir el cambio estructural desde la fuente.

Limitaciones:
Los autores señalan que, si bien la realineación del cabezal recupera un rendimiento significativo, persiste una brecha residual respecto al rendimiento IID, lo que sugiere que la geometría del backbone también se degrada ligeramente. Los hallazgos están actualmente limitados a arquitecturas CNN/ResNet en conjuntos de datos de imágenes, y el análisis opera a nivel de canal en lugar de granularidad de subcanal.