Technische Samenvatting: COSMOS

Probleemstelling

Federated Learning (FL) staat voor aanzienlijke uitdagingen in heterogene omgevingen waar clients niet alleen verschillen in data-distributie (statistische heterogeniteit), maar ook in modelarchitectuur (architecturale heterogeniteit). Hoewel er methoden voor Gepersonaliseerd Federated Learning (PFL) zijn ontwikkeld om modellen op maat te maken voor individuele clients, gaan de meeste bestaande benaderingen uit van structurele compatibiliteit of vereisen dat clients hun eigen modelarchitectuur openbaar maken. Dit creëert een praktische barrière, aangezien clients vaak willekeurige of propriëtaire modellen willen gebruiken zonder parameters of gradiënten te onthullen.

Bestaande model-agnostische FL-benaderingen, zoals FedMD en COMET, vertrouwen op communicatie op output-niveau (bijvoorbeeld zachte of harde labels) met behulp van een gedeelde, niet-gelabelde dataset. Deze methoden kampen echter met specifieke beperkingen:

1. Passieve Servercoördinatie: Ze behandelen de server vaak als een passieve coördinator, waardoor ze de rekenkracht van de server niet benutten voor actieve personalisatie.
2. Heuristische Clustering: Methoden zoals COMET vertrouwen op K-means clustering, wat vereist dat het aantal clusters ($K$) vooraf bekend is en geen formele mechanismen biedt voor personalisatiegaranties wanneer de clientdiversiteit hoog is.
3. Gebrek aan Theoretische Garanties: De bestaande literatuur over model-agnostische PFL mist doorgaans rigoureuze theoretische garanties met betrekking tot risicovermindering, en vertrouwt vaak op vereenvoudigde aannames (bijvoorbeeld lineaire modellen, Gaussische verdelingen) of biedt slechts standaard convergentieresultaten naar een stationair punt.

Methodologie: COSMOS Framework

De auteurs stellen COSMOS (Clustered Output-based Server Models) voor, een model-agnostisch PFL-framework dat personalisatie aan de serverzijde mogelijk maakt met uitsluitend communicatie van pseudo-labels. Het framework werkt onder de aanname dat clients privé gelabelde data bezitten en dat een gedeelde, niet-gelabelde publieke dataset $U$ (bijvoorbeeld web-gescrapte of synthetische data) beschikbaar is.

COSMOS verloopt in twee distincte fasen:

1. Pre-training en Clustering (PTC) Fase

Deze fase legt de initiële client-server mapping en cluster-specifieke modellen vast:

• Lokale Pre-training: Clients trainen hun lokale modellen $f_i$ op hun privédata $D_i$ en genereren pseudo-labels op de publieke dataset $U$.
• Gierige Clustering: De server clusteren clients op basis van de gelijkenis van hun pseudo-labels op $U. In tegenstelling tot K-means, maakt COSMOS gebruik van een gierig clustering-algoritme dat niet vereist dat $K$ vooraf wordt gespecificeerd. Het selecteert iteratief een client met de meeste "dichtstbijzijnde" buren (binnen een afstandsdrempel $B_0$) om een cluster te vormen totdat alle clients zijn toegewezen. Dit proces is endogeen aan de trainingsprocedure.
• Serverzijdige Distillatie: Voor elke cluster $G_k$ aggregeert de server de pseudo-labels van de bijbehorende clienten om een gemiddelde pseudo-labelverdeling te creëren. Vervolgens traint het een dedicated "teacher"-model $h_k$ voor die cluster met behulp van deze geaggregeerde labels en een specifieke regularisatieterm om lokale gladheid en robuustheid te waarborgen.
• Clientzijdige Distillatie: Clients ontvangen de pseudo-labels van het teacher-model van hun toegewezen cluster en fine-tunen hun lokale modellen met behulp van deze labels.

2. Iteratieve Federated Fine-tuning (IFFT) Fase

Na de PTC-fase betreedt het systeem een iteratieve lus:

• Lokale Fine-tuning: Clients verfijnen hun modellen op hun privé gelabelde data.
• Server Distillatie: Clients sturen bijgewerkte pseudo-labels naar de server. De server update de cluster-specifieke teacher-modellen op basis van de nieuwe geaggregeerde pseudo-labels.
• Client Distillatie: Clients verfijnen hun lokale modellen verder met behulp van de bijgewerkte teacher-pseudo-labels.

Cruciaal is dat de servermodellen worden getraind met een regularisatieterm die is afgeleid van semi-supervised learning (SSL) theorie, die predictie-inconsistenties onder datatransformaties (bijvoorbeeld rotaties, translaties) bestraft om robuustheid te waarborgen.

Theoretische Bijdragen

Het artikel biedt de eerste end-to-end theoretische analyse van personalisatie-risicovermindering voor model-agnostische PFL. Door gebruik te maken van tools uit SSL-theorie (specifiek connectiviteit op basis van expansie en gebonden pseudo-labelfout), leiden de auteurs voldoende voorwaarden af voor exponentiële risicovermindering.

Belangrijke theoretische componenten zijn:

• Aannames: De analyse gaat uit van gebonden belangsgewichten tussen de publieke en privéverdelingen, lokaal goed gestelde doelfuncties en een "confidence margin" waarbij clients vertrouwen hebben in hun voorspellingen op de publieke data. Het vertrouwt ook op een expansie-eigenschap van de data-distributie, die garandeert dat pseudo-labels niet-triviale informatie dragen over de ground truth.
• Risicovermindering: De auteurs bewijzen dat onder deze voorwaarden het personalisatierisico voor zowel servermodellen als clientmodellen exponentieel afneemt met het aantal iteraties. Specifiek neemt de risicobound af met een snelheid die wordt bepaald door de clusterkwaliteit en de dekking van de publieke dataset, en daalt uiteindelijk onder het risico van lokaal trainen alleen.

Experimentele Resultaten

De auteurs evalueerden COSMOS op vier beeldclassificatie-benchmarks: EMNIST-balanced, CIFAR-10, CIFAR-100 en Tiny ImageNet. De experimenten maakten gebruik van een mix van clientarchitecturen (AlexNet, ResNet, MobileNet, SqueezeNet) en een VGG16 servermodel.

• Prestaties: COSMOS presteerde consistent beter dan alle model-agnostische baselines (FedMD, FedCT, COMET) op alle benchmarks. De prestatiekloof werd groter naarmate de taakcomplexiteit toenam (bijvoorbeeld van EMNIST naar Tiny ImageNet).
• Heterogeniteit: COSMOS bleef robuust voor variërende graden van data-heterogeniteit (gecontroleerd door de Dirichlet-parameter $\alpha$) en verschillende mengsels van clientarchitecturen. Het behield een aanzienlijk prestatievoordeel, zelfs wanneer clients zwakkere modellen gebruikten.
• Communicatie-efficiëntie: Door uitsluitend pseudo-labels (hard of zacht) te communiceren in plaats van modelparameters of gradiënten, verlaagde COSMOS de communicatiebandbreedte met 1–2 ordes van grootte ten opzichte van methoden die parameters delen.
• Homogene Instellingen: Zelfs in homogene instellingen waar alle clients dezelfde architectuur delen, bleef COSMOS concurrerend met state-of-the-art PFL-methoden, wat aantoont dat het model-agnostische ontwerp geen prestatieboete met zich meebrengt wanneer homogeniteit aanwezig is.

Betekenis en Claims

Het artikel beweert dat COSMOS een kritieke knelpunt in FL aanpakt door hoge prestaties voor personalisatie mogelijk te maken zonder structurele kennis van clientmodellen te vereisen. De betekenis ligt op drie hoofdgebieden:

1. Praktische Implementatie: Het verwijdert de structurele barrières die PFL beperken in diverse, real-world ecosystemen waar clients propriëtaire of willekeurige architecturen gebruiken.
2. Theoretische Vooruitgang: Het biedt de eerste algemene garanties voor risicovermindering voor model-agnostische PFL, en gaat voorbij standaard convergentieresultaten om exponentiële risicovermindering te bewijzen onder realistische aannames.
3. Schaalbaarheid: Het toont aan dat gepersonaliseerd leren aan de serverzijde met behulp van pseudo-labels een levensvatbaar en schaalbaar paradigma is voor sterk heterogene omgevingen, met aanzienlijke winsten in communicatie-efficiëntie.

De auteurs concluderen dat, hoewel het huidige werk zich richt op een client-server structuur, het framework de grondslag legt voor toekomstige gedecentraliseerde leersystemen, waardoor de praktische relevantie van federated learning verder wordt versterkt.