Federated Learning with Energy-Based Structured Probabilistic Inference

Résumé Technique : Apprentissage Fédéré avec Inférence Probabiliste Structurée basée sur l'Énergie

Énoncé du Problème
L'apprentissage fédéré (FL) repose généralement sur la moyenne fédérée (FedAvg), où le serveur agrège les mises à jour des modèles des clients en utilisant des poids proportionnels au nombre d'échantillons locaux. Bien que simple, cette heuristique devient sous-optimale lorsque les clients possèdent des distributions de données hétérogènes (non-IID). Dans ces contextes, les règles de pondération fixes traitent toutes les mises à jour mises à l'échelle comme étant également fiables, ce qui entraîne une dérive des clients, une convergence lente ou une instabilité de l'entraînement. Les solutions existantes, telles que les termes proximaux (FedProx), les variables de contrôle (SCAFFOLD) ou les règles d'agrégation robustes (par exemple, la moyenne tronquée, la médiane géométrique), traitent des modes de défaillance spécifiques mais reposent généralement sur des heuristiques fixes. Elles ne parviennent pas à modéliser explicitement à la fois la fiabilité individuelle des clients et les interactions entre les clients au sein d'un cadre d'inférence unifié.

Méthodologie
Les auteurs proposent de reformuler le problème d'agrégation côté serveur comme une tâche d'inférence probabiliste structurée basée sur l'énergie utilisant des champs aléatoires conditionnels (CRF). L'idée centrale est d'inférer des « étiquettes de fiabilité » discrètes pour chaque client sur la base de leurs mises à jour de modèles locales, plutôt que de se fier uniquement aux nombres d'échantillons.

1. Formulation du Graphe :

   • Nœuds : Clients participants dans un tour de communication.
   • Observations ($x$) : Les vecteurs de mise à jour des modèles locaux (deltas) et les nombres d'échantillons.
   • Variables Latentes ($y$) : Étiquettes de fiabilité discrètes (par exemple, $\{0, 0.25, 1\}$), où les valeurs plus élevées indiquent une plus grande confiance dans la mise à jour.

2. Fonction d'Énergie :
   Le CRF définit une fonction d'énergie $E(y|x)$ composée de deux potentiels :

   • Potentiels Unaires ($\psi_u$) : Mesurent la fiabilité individuelle de chaque client. Le serveur calcule une mise à jour de référence robuste (médiane coordonnée de tous les deltas des clients). Il dérive ensuite trois statistiques normalisées pour chaque client : l'accord par cosinus avec la référence, la distance relative par rapport à la référence, et un score d'anomalie de norme de mise à jour robuste. Ces éléments sont combinés via une fonction sigmoïde pour générer un score de fiabilité, qui initialise le potentiel unaire.
   • Potentiels Pare-à-Pare ($\psi_p$) : Encouragent la cohérence entre les clients ayant des mises à jour similaires. L'affinité entre les clients $i$ et $j$ est définie par la similitude de cosinus de leurs mises à jour, modulée par le minimum de leurs scores de fiabilité individuels. Les clients à haute affinité sont pénalisés s'ils se voient assigner des étiquettes de fiabilité divergentes.

3. Inférence et Agrégation :

   • Puisque l'inférence MAP (Maximum a Posteriori) exacte est intraitable pour un graphe totalement connecté, les auteurs emploient une approximation par champ moyen (mean-field). Celle-ci met à jour de manière itérative la distribution a posteriori des étiquettes de fiabilité pour chaque client en fonction des scores unaires et des coûts attendus de ses voisins.
   • Le poids d'agrégation final pour un client est le produit de son poids de base FedAvg (nombre d'échantillons) et de sa fiabilité attendue inférée. Ce poids raffiné est utilisé pour calculer le nouveau modèle global.

Principales Contributions

• Formulation : Le papier formule l'agrégation du serveur FL comme un problème d'inférence structurée basée sur l'énergie sur les variables de fiabilité des clients.
• Instanciation : Il implémente un CRF au niveau des clients totalement connecté où les potentiels unaires sont dérivés de statistiques de mise à jour robustes (basées sur la médiane) et les potentiels pare-à-pare sont dérivés de l'affinité des mises à jour.
• Algorithme : Il propose un algorithme d'agrégation par champ moyen compatible avec les boucles d'entraînement standard de FedAvg, ne nécessant aucun changement de l'optimisation côté client.
• Évaluation : La méthode est évaluée par rapport aux références standards de FL (FedAvg, FedProx) et aux méthodes d'agrégation robustes (Krum, moyenne tronquée, médiane géométrique, RFA).

Résultats Expérimentaux
Les auteurs ont évalué l'approche sur MNIST, CIFAR-10 et CIFAR-100 sous des conditions non-IID sévères (biais de labels de Dirichlet avec $\alpha=0.003$ et $\alpha=0.1$).

• Performance : L'approche guidée par CRF (CRF-FedAvg et CRF-FedProx) surpasse systématiquement les bases d'agrégation fixes et robustes sur tous les jeux de données.
• Signification des gains : Alors que les améliorations sur MNIST sont modestes (car la plupart des méthodes atteignent une précision élevée), les gains sont plus prononcés sur CIFAR-10 et CIFAR-100. CRF-FedProx a obtenu la précision la plus élevée sur CIFAR-10, et CRF-FedAvg a obtenu la plus élevée sur CIFAR-100.
• Convergence : Les courbes de perte de validation indiquent que l'agrégation guidée par CRF converge plus rapidement et plus stablement que les bases robustes, particulièrement dans les derniers tours de communication. Cela suggère que la méthode préserve efficacement les directions de mise à jour utiles tout en supprimant les directions peu fiables, évitant ainsi le plateau conservateur observé avec des méthodes comme la médiane géométrique.

Signification et Revendications
Le papier affirme que traiter l'agrégation comme un problème d'inférence structurée offre un avantage distinct par rapport aux heuristiques fixes. En modélisant explicitement tant la fiabilité individuelle des mises à jour que la compatibilité entre clients, le serveur peut assigner des poids qui sont adaptatifs à la géométrie spécifique de la mise à jour du tour actuel.

• Les auteurs soulignent que cette approche est complémentaire aux techniques d'optimisation côté client (comme FedProx) ; tandis que FedProx contraint la dérive locale, l'agrégation guidée par CRF affine la manière dont les mises à jour sont combinées après l'entraînement.
• La méthode est présentée comme un module complémentaire (plug-in) côté serveur qui ne modifie pas l'entraînement ou les protocoles de communication des clients.
• Les auteurs notent des limitations, notamment l'utilisation de potentiels conçus à la main (plutôt que appris), des hyperparamètres fixes à travers les expériences, et le coût computationnel du graphe totalement connecté ($O(m^2)$), qui pourrait nécessiter une sparsification pour des fédérations à très grande échelle. Ils précisent également que l'évaluation actuelle se concentre sur l'hétérogénéité statistique plutôt que sur la robustesse face aux attaques adverses (Byzantine).