Résumé Technique : COSMOS

Énoncé du Problème

L'apprentissage fédéré (FL) fait face à des défis majeurs dans des environnements hétérogènes où les clients diffèrent non seulement par la distribution de leurs données (hétérogénéité statistique), mais aussi par l'architecture de leurs modèles (hétérogénéité architecturale). Bien que des méthodes d'apprentissage fédéré personnalisé (PFL) aient été développées pour adapter les modèles aux clients individuels, la plupart des approches existantes supposent une compatibilité structurelle ou exigent que les clients divulguent leurs architectures de modèles propriétaires. Cela crée un obstacle pratique, car les clients souhaitent souvent utiliser des modèles arbitraires ou propriétaires sans révéler leurs paramètres ou leurs gradients.

Les approches actuelles de FL agnostiques au modèle, telles que FedMD et COMET, reposent sur une communication au niveau des sorties (par exemple, étiquettes douces ou dures) en utilisant un ensemble de données public non étiqueté. Cependant, ces méthodes souffrent de limitations spécifiques :

1. Coordination passive du serveur : Elles traitent souvent le serveur comme un coordinateur passif, ne parvenant pas à exploiter les ressources informatiques du serveur pour une personnalisation active.
2. Regroupement heuristique : Des méthodes comme COMET reposent sur le clustering K-means, qui nécessite une connaissance préalable du nombre de clusters ($K$) et manque de mécanismes formels pour garantir la personnalisation lorsque la diversité des clients est élevée.
3. Absence de garanties théoriques : La littérature existante sur le PFL agnostique au modèle manque généralement de garanties théoriques rigoureuses concernant la contraction du risque, s'appuyant souvent sur des hypothèses simplifiées (par exemple, modèles linéaires, distributions gaussiennes) ou ne fournissant que des résultats standards de convergence vers un point stationnaire.

Méthodologie : Cadre COSMOS

Les auteurs proposent COSMOS (Clustered Output-based Server Models), un cadre PFL agnostique au modèle qui permet une personnalisation côté serveur en utilisant uniquement la communication d'étiquettes pseudo. Le cadre fonctionne sous l'hypothèse que les clients possèdent des données étiquetées privées et qu'un ensemble de données public non étiqueté partagé $U$ (par exemple, des données extraites du web ou synthétiques) est disponible.

COSMOS se déroule en deux phases distinctes :

1. Phase de Pré-entraînement et de Regroupement (PTC)

Cette phase établit la correspondance initiale client-serveur et les modèles spécifiques aux clusters :

• Pré-entraînement local : Les clients entraînent leurs modèles locaux $f_i$ sur leurs données privées $D_i$ et génèrent des étiquettes pseudo sur l'ensemble de données public $U$.
• Regroupement glouton : Le serveur regroupe les clients en fonction de la similarité de leurs étiquettes pseudo sur $U$. Contrairement au K-means, COSMOS emploie un algorithme de regroupement glouton qui ne nécessite pas que $K$ soit spécifié à l'avance. Il sélectionne itérativement un client ayant les voisins les plus « proches » (dans un seuil de distance $B_0$) pour former un cluster jusqu'à ce que tous les clients soient assignés. Ce processus est endogène à la procédure d'entraînement.
• Distillation côté serveur : Pour chaque cluster $G_k$, le serveur agrège les étiquettes pseudo de ses clients membres pour créer une distribution moyenne d'étiquettes pseudo. Il entraîne ensuite un modèle « enseignant » dédié $h_k$ pour ce cluster en utilisant ces étiquettes agrégées et un terme de régularisation spécifique pour assurer la régularité locale et la robustesse.
• Distillation côté client : Les clients reçoivent les étiquettes pseudo du modèle enseignant de leur cluster assigné et affinent leurs modèles locaux en utilisant ces étiquettes.

2. Phase de Finetuning Fédéré Itératif (IFFT)

Après la phase PTC, le système entre dans une boucle itérative :

• Affinement local : Les clients affinent leurs modèles sur leurs données étiquetées privées.
• Distillation serveur : Les clients envoient des étiquettes pseudo mises à jour au serveur. Le serveur met à jour les modèles enseignants spécifiques aux clusters en fonction des nouvelles étiquettes pseudo agrégées.
• Distillation client : Les clients affinent davantage leurs modèles locaux en utilisant les étiquettes pseudo des enseignants mises à jour.

Crucialement, les modèles du serveur sont entraînés en utilisant un terme de régularisation dérivé de la théorie de l'apprentissage semi-supervisé (SSL), qui pénalise les incohérences de prédiction sous des transformations de données (par exemple, rotations, translations) pour assurer la robustesse.

Contributions Théoriques

L'article fournit la première analyse théorique de bout en bout de la contraction du risque de personnalisation pour le PFL agnostique au modèle. En tirant parti d'outils issus de la théorie SSL (spécifiquement la connectivité basée sur l'expansion et l'erreur d'étiquette pseudo bornée), les auteurs dérivent des conditions suffisantes pour une contraction exponentielle du risque.

Les composants théoriques clés incluent :

• Hypothèses : L'analyse suppose des poids d'importance bornés entre les distributions publique et privée, des objectifs localement bien conditionnés, et une « marge de confiance » où les clients sont confiants dans leurs prédictions sur les données publiques. Elle repose également sur une propriété d'expansion de la distribution des données, garantissant que les étiquettes pseudo portent des informations non triviales sur la vérité terrain.
• Contraction du risque : Les auteurs prouvent que sous ces conditions, le risque de personnalisation pour les modèles du serveur et les modèles des clients se contracte exponentiellement avec le nombre d'itérations. Plus précisément, la borne de risque diminue à un taux déterminé par la qualité du cluster et la couverture de l'ensemble de données public, finissant par tomber en dessous du risque d'un entraînement local seul.

Résultats Expérimentaux

Les auteurs ont évalué COSMOS sur quatre benchmarks de classification d'images : EMNIST-balanced, CIFAR-10, CIFAR-100 et Tiny ImageNet. Les expériences ont utilisé un mélange d'architectures de clients (AlexNet, ResNet, MobileNet, SqueezeNet) et un modèle serveur VGG16.

• Performance : COSMOS a systématiquement surpassé toutes les références agnostiques au modèle (FedMD, FedCT, COMET) sur tous les benchmarks. L'écart de performance s'est élargi à mesure que la complexité de la tâche augmentait (par exemple, d'EMNIST à Tiny ImageNet).
• Hétérogénéité : COSMOS est resté robuste face à des degrés variables d'hétérogénéité des données (contrôlés par le paramètre Dirichlet $\alpha$) et à différents mélanges d'architectures de clients. Il a maintenu un avantage de performance significatif même lorsque les clients utilisaient des modèles plus faibles.
• Efficacité de communication : En communiquant uniquement des étiquettes pseudo (dures ou douces) au lieu de paramètres de modèle ou de gradients, COSMOS a réduit la bande passante de communication d'un à deux ordres de grandeur par rapport aux méthodes de partage de paramètres.
• Paramètres homogènes : Même dans des paramètres homogènes où tous les clients partagent la même architecture, COSMOS est resté compétitif par rapport aux méthodes PFL de l'état de l'art, démontrant que sa conception agnostique au modèle n'entraîne pas de pénalité de performance lorsque l'homogénéité est présente.

Importance et Revendications

L'article revendique que COSMOS répond à un goulot d'étranglement critique dans le FL en permettant une personnalisation haute performance sans nécessiter de connaissances structurelles des modèles des clients. Son importance réside dans trois domaines principaux :

1. Déploiement pratique : Il élimine les barrières structurelles limitant le PFL dans des écosystèmes réels diversifiés où les clients utilisent des architectures propriétaires ou arbitraires.
2. Avancement théorique : Il fournit les premières garanties générales de contraction du risque pour le PFL agnostique au modèle, allant au-delà des résultats de convergence standards pour prouver une réduction exponentielle du risque sous des hypothèses réalistes.
3. Évolutivité : Il démontre que l'apprentissage personnalisé côté serveur utilisant des étiquettes pseudo est un paradigme viable et évolutif pour des environnements hautement hétérogènes, offrant des gains substantiels en efficacité de communication.

Les auteurs concluent que bien que le travail actuel se concentre sur une structure client-serveur, le cadre pose les bases pour de futurs systèmes d'apprentissage décentralisés, renforçant davantage la pertinence pratique de l'apprentissage fédéré.