Résumé Technique : Preuves Mécanistes de Représentations Préservées-mais-Désalignées dans le FedAvg Non-IID

1. Énoncé du Problème

Federated Averaging (FedAvg) est la procédure d'optimisation standard pour l'apprentissage fédéré (Federated Learning - FL), pourtant elle souffre d'une dégradation significative des performances sous des données clients non-IID (hétérogénéité statistique), particulièrement en cas de biais de labels (label skew). Les explications existantes attribuent cet échec au dérive des clients (client drift), à l'incohérence des objectifs et à la divergence des poids, caractérisant ainsi le processus d'optimisation mais laissant l'effet interne de l'agrégation largement inobservé.

Une question cruciale reste ouverte : la dégradation de la précision globale reflète-t-elle l'effacement des représentations apprises par les clients, ou reflète-t-elle un désalignement où les structures internes sont préservées mais ne parviennent pas à se connecter efficacement avec la voie de prédiction finale ? Cet article étudie cette question en utilisant des outils d'interprétabilité mécaniste (Mechanistic Interpretability - MI) pour déterminer si le non-IID FedAvg détruit la connaissance apprise ou s'il se contente de la déconnecter du classifieur.

2. Méthodologie

Les auteurs analysent des modèles de vision entraînés sur des clients avec une parcimonie (sparsity) de type SimpleCNN et ResNet10 sur CIFAR-10 et Fashion-MNIST. Pour garantir l'interprétabilité, les clients s'entraînent sous un programme de parcimonie basé sur la magnitude (recuit vers 99,9 % de parcimonie), les contrôles sur modèles denses confirmant que les résultats ne sont pas des artefacts de la parcimonie. L'étude emploie quatre diagnostics mécanistes :

2.1. Découverte de Circuits

En utilisant un masquage binaire différentiable, les auteurs identifient des sous-réseaux minimaux (circuits) suffisants pour prédire des classes spécifiques. Ils mesurent :

• Stabilité intra-client : Le chevauchement des circuits à travers les cycles de communication.
• Consistance inter-client : La similitude des circuits entre différents clients.
• Préservation local-global : L'Intersection-over-Union (IoU) entre le circuit local d'un client et le circuit correspondant dans le modèle global agrégé.
• Interventions : Suffisance (précision lorsque seul le circuit est actif) et Nécessité (chute de précision lorsque le circuit est ablaté).

2.2. Sondage Linéaire (Linear Probing)

Un classifieur linéaire est entraîné sur les activations de la pénultième couche du modèle global. Cela teste si l'information discriminante de classe reste linéairement accessible dans le backbone, même si la tête de prédiction native du modèle ne parvient pas à l'utiliser.

2.3. Fine-tuning de la Tête Uniquement

Le backbone est gelé, et seule la tête de classification finale est soumise à un fine-tuning sur un petit sous-ensemble IID équilibré (200 échantillons). Cela isole si la perte de précision est localisée dans la couche de lecture (readout layer) ou dans le backbone de caractéristiques.

2.4. Autoencodeurs de Caractéristiques Universels (USAE)

Un USAE est entraîné conjointement sur les activations de modèles globaux IID et non-IID pour apprendre un dictionnaire de caractéristiques partagé. Cela teste si les deux modèles organisent leurs caractéristiques selon une base commune et permet la reconstruction croisée pour évaluer la compatibilité des caractéristiques.

3. Résultats Clés

3.1. Stabilité et Divergence des Circuits

• Stabilité Locale : La dynamique d'entraînement locale est stable, indépendamment de l'hétérogénéité des données. Les clients convergent vers des circuits spécifiques aux classes hautement stables (IoU $\approx$ 0,90), même sous un biais extrême ($\alpha=0,05$).
• Divergence Inter-Clients : À mesure que l'hétérogénéité augmente, les clients apprennent des circuits différents pour une même classe. Sous un biais extrême, le chevauchement des circuits inter-clients chute à un IoU $\approx$ 0,45.
• Préservation Partielle : Le modèle global n'agrège pas les circuits en un bruit. La préservation local-global (IoU $\approx$ 0,55) reste plus élevée que le chevauchement inter-client, indiquant que le modèle global hérite d'une structure partielle des clients plutôt que de converger vers une solution unique.

3.2. Les Circuits Fonctionnels Persistent Malgré une Précision Nulle

Dans des contextes non-IID extrêmes, le modèle global peut atteindre une précision proche de zéro pour certaines classes (ex: 0,0 % pour "Chat" ou "Cerf"). Cependant :

• Suffisance : Les circuits extraits pour ces classes atteignent une suffisance quasi parfaite (100 % de précision) lorsqu'ils sont isolés.
• Nécessité : L'ablation de ces circuits fait chuter la précision à presque zéro.
• Conclusion : Les voies de calcul pour ces classes existent et sont causalement responsables de la prédiction, mais le modèle échoue à acheminer les entrées à travers elles lors de l'inférence standard.

3.3. Séparabilité Linéaire et Désalignement de la Tête

• Sondes Linéaires : Bien que la précision globale chute significativement sous le biais (ex: de 68,3 % à 43,1 %), la précision de la sonde linéaire reste élevée et stable (ne chutant qu'à 64,7 %). L'écart entre la sonde et la précision globale s'élargit de manière monotone avec l'hétérogénéité (+21,7 points à $\alpha=0,05$).
• Fine-tuning de la Tête : Le fine-tuning de la tête de classification uniquement sur un petit ensemble équilibré récupère une part substantielle de la précision (+10,2 points à $\alpha=0,05$), bien qu'il ne restaure pas totalement la performance IID. Cela confirme que le backbone conserve des caractéristiques utiles, mais que la tête par défaut est désalignée.

3.4. Base de Caractéristiques Partagée (USAE)

• Chevauchement de Concepts : 98,7 % des 2 048 concepts du dictionnaire USAE s'activent de manière non triviale dans les modèles IID et non-IID, indiquant une base de caractéristiques largement partagée.
• Transfert Croisé : Lorsque les activations non-IID sont décodées dans l'espace de caractéristiques IID, la tête IID performe bien. Inversement, la tête non-IID échoue même lorsqu'elle est alimentée par des activations IID. Cela confirme que les caractéristiques sont présentes et compatibles, mais que le mécanisme de lecture (readout) non-IID est le point de défaillance.

4. Signification et Revendications

L'article soutient que la dégradation du FedAvg non-IID n'est pas entièrement expliquée par l'effacement des représentations. Au lieu de cela, le mécanisme principal est le désalignement :

1. Préservation : Les structures internes apprises par les clients (circuits) et les bases de caractéristiques sont largement préservées dans le modèle global agrégé.
2. Désalignement : La voie de prédiction finale (la tête de classification) devient désalignée avec ces représentations préservées. Le modèle "connaît" les classes mais ne parvient pas à acheminer les entrées vers les bons circuits ou à utiliser les caractéristiques efficacement.

Implications :

• Recadrage du Problème : Le problème n'est pas nécessairement que les clients "oublient" ou que les poids divergent vers des bassins incompatibles qui détruisent la connaissance, mais que le processus d'agrégation perturbe la relation structurelle entre le backbone et la tête.
• Directions Algorithmiques :
  • Réalignement de la Tête Post-Agrégation : Un simple fine-tuning de la tête sur un petit ensemble équilibré peut récupérer une performance significative sans changer le protocole fédéré.
  • Agrégation Sensible aux Circuits : Les méthodes futures pourraient aligner les neurones implémentant des circuits fonctionnels similaires entre les clients avant la moyenne (par exemple, via des méthodes de symétrie de permutation) afin de réduire le décalage structurel à la source.

Limites :
Les auteurs notent que bien que le réalignement de la tête récupère une performance significative, un écart résiduel par rapport à la performance IID subsiste, suggérant que la géométrie du backbone se dégrade également légèrement. Les conclusions sont actuellement limitées aux architectures CNN/ResNet sur des jeux de données d'images, et l'analyse opère au niveau des canaux plutôt qu'à la granularité des sous-canaux.