Technische Samenvatting: PEP-FedPT

Probleemstelling

Federated Learning (FL) staat voor aanzienlijke uitdagingen bij het aanpassen van grote vooraf getrainde Vision Transformers (ViTs) aan downstream taken onder data-heterogeniteit. Hoewel Visual Prompt Tuning (VPT) een parameter-efficiënte oplossing biedt voor fijnafstemming, worstelen bestaande benaderingen met het vinden van een balans tussen generalisatie en personalisatie:

• Globale Prompt Tuning: Methoden zoals FedPR gebruiken gedeelde prompts voor alle clients. Hoewel ze goed generaliseren, slagen ze er niet in zich aan te passen aan sterk heterogene client-dataverdelingen (non-iid), wat leidt tot "client drift" en suboptimale prestaties.
• Gepersonaliseerde Prompt Tuning: Methoden zoals pFedPG genereren client-specifieke prompts. Hoewel ze zich goed aanpassen aan lokale data, leiden ze vaak tot overfitting, wat resulteert in slechte generalisatie naar onzichtbare of niet-deelnemende clients. Bovendien gaan veel gepersonaliseerde benaderingen uit van volledige client-participatie in elke ronde, wat in de praktijk onuitvoerbaar is voor echte FL-systemen.
• Bestaande Hybride Benaderingen: Methoden zoals SGPT proberen deze kloof te overbruggen met groeps-specifieke prompts, maar introduceren optimalisatiecomplexiteit door niet-differentieerbare selectiemechanismen en twee-trainingsprocessen.

De centrale vraag die wordt behandeld is: Kunnen we effectieve personalisatie bereiken terwijl we uitsluitend vertrouwen op globaal gedeelde prompts?

Methodologie: PEP-FedPT

De auteurs stellen PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning) voor, een unificerend kader dat personalisatie bereikt zonder client-afhankelijke trainbare parameters op te slaan. De kerninnovatie is de Class-Contextualized Mixed Prompt (CCMP).

1. Prompt Architectuur

Het kader maakt gebruik van twee sets trainbare parameters die globaal worden gedeeld over alle clients:

• Gedeelde Prompts ($P_S$): Ingevoegd in de allerlaatste laag van de ViT. Deze vangen laag-niveau, class-agnostische representaties op die gemeenschappelijk zijn over het manifold.
• Class-specifieke Prompts ($P_C$): Een set prompts die overeenkomt met elke klasse, globaal gedeeld.

2. Class-Contextualized Mixed Prompt (CCMP)

In plaats van een enkele statische prompt in te voegen, genereert PEP-FedPT dynamisch een gemengde prompt ($m$) voor elke invoersteekproef op tussenliggende lagen (specifiek lagen 5, 6 en 7 in de experimenten). Deze gemengde prompt is een zachte combinatie van de globale class-specifieke prompts, gewogen door scores afgeleid van:

• Globale Class Prototypes ($\mu$): Geaggregeerde centroiden van CLS-token representaties voor elke klasse, berekend door het middelen van CLS tokens van datapunten binnen een klasse over clients. Deze worden periodiek bijgewerkt via federated averaging met momentum.
• Client Class Priors ($\delta$): De empirische labelverdeling (frequentie van klassen) specifiek voor de lokale data van elke client.

3. Prompt Mengmechanisme

Voor een gegeven invoer $x$ bij client $k$, berekent het systeem een zachte gewichtsvector $s$ voor de class-specifieke prompts:

1. Similariteitscore: De cosinus-similariteit tussen de CLS-token representatie van de invoer ($cls_{l-1}$) op laag $l$ en de globale class prototypes ($\mu$) wordt berekend.
2. Gewogen Scoren: Deze similariteiten worden geschaald met de client's lokale class priors ($\delta$) en een temperatuurparameter ($\tau$).
3. Normalisatie: De scores worden genormaliseerd om een kansverdeling te vormen.
4. Mengen: De uiteindelijke CCMP wordt berekend als de gewogen som van de globale class-specifieke prompts: $m = P_C \times s$.

Dit stelt elke client in staat om dynamisch de prompt te personaliseren op basis van hun lokale dataverdeling en de huidige invoer, uitsluitend gebruikmakend van globale informatie (prompts en prototypes) en lokale statistieken (priors), zonder ruwe data of client-specifieke modelgewichten te verzenden.

4. Theoretische Grondslagen

Het paper biedt twee theoretische rechtvaardigingen voor het CCMP-ontwerp:

• Minimalisatie van een Quadratische Bovengrens: De auteurs bewijzen dat de CCMP een quadratische bovengrens op de globale verliesfunctie minimaliseert, ervan uitgaande dat het verlies Lipschitz-glad is en de prompt-ruimte compact is. De optimale mengstrategie komt overeen met het gewogen gemiddelde van class prompts gebaseerd op class priors.
• MMSE-Optimaliteit: De CCMP blijkt de Minimum Mean Squared Error (MMSE) schatter te zijn van de ware class prompt gegeven de waargenomen CLS-token representatie.

Belangrijkste Bijdragen

1. Unificerend Kader: Introductie van PEP-FedPT, dat gezamenlijk gedeelde en class-specifieke prompts optimaliseert om data-heterogeniteit in ViT federated learning aan te pakken, en een balans slaat tussen generaliseerbaarheid en personalisatie.
2. Nieuwe Mengstrategie: Ontwerp van de CCMP, die per-steekproef prompts genereert als een functie van globale class-specifieke prompts, gemoduleerd door lokale class priors en globale prototypes. Dit elimineert de noodzaak voor lokale prompt-opslag of server-side generatie van client-specifieke prompts.
3. Theoretische Inzichten: Aantonen dat CCMP een quadratische bovengrens op het verlies minimaliseert en optimaal is in de MMSE-zin.
4. Empirische Superioriteit: Uitgebreide validatie die aantoont dat de methode state-of-the-art baselines overtreft in diverse heterogeniteitsscenario's.

Experimentele Resultaten

De methode werd geëvalueerd op CIFAR-100, TinyImageNet, DomainNet en iNaturalist onder verschillende non-iid instellingen (pathologische splitsen en Dirichlet-verdelingen).

• Prestaties: PEP-FedPT presteerde consequent beter dan baselines (inclusief FedVPT, FedPR, SGPT en pFedPG) in zowel Gemiddelde Nauwkeurigheid als Slechtste Lokale Nauwkeurigheid.
  • Op CIFAR-100 (Pathologisch) behaalde het 95,46% gemiddelde nauwkeurigheid, wat pFedPG (92,96%) en FedVPT-D (85,15%) overtrof.
  • Op TinyImageNet verbeterde het de gemiddelde nauwkeurigheid met 8,59% ten opzichte van pFedPG.
• Generalisatie: In vastgehouden evaluaties (testen op clients die niet deelnamen aan training) behield PEP-FedPT hoge nauwkeurigheid (93,71% op CIFAR-100), terwijl gepersonaliseerde methoden zoals pFedPG faalden om te generaliseren (gemarkeerd als NA of aanzienlijk lager).
• Efficiëntie: De methode vereiste aanzienlijk minder communicatierondes (12 rondes om 83% nauwkeurigheid op CIFAR-100 te bereiken) en lagere communicatie-overhead (4,6M parameters gecommuniceerd) in vergelijking met SGPT (13M) en FedVPT-D.
• Ablatiestudies: Experimenten bevestigden dat het combineren van Gedeelde Prompts met CCMP aanzienlijke winsten oplevert ten opzichte van het gebruik van Gedeelde Prompts alleen. De opname van Class Priors verbeterde de prestaties verder, vooral in pathologische instellingen.

Betekenis en Claims

Het paper beweert dat PEP-FedPT een sterke basis legt voor efficiënte en generaliseerbare federated prompt tuning van Vision Transformers. De primaire betekenis ligt in het oplossen van de afweging tussen personalisatie en generalisatie in FL:

• Het bereikt personalisatie door lokale dataverdelingen (via priors) te benutten om prompts dynamisch te mengen.
• Het behoudt generalisatie door uitsluitend te vertrouwen op globaal gedeelde prompts en prototypes, waardoor overfittingproblemen van volledig gepersonaliseerde modellen worden vermeden.
• Het biedt schaalbaarheid door geen client-specifieke parameteropslag of volledige client-participatie te vereisen, waardoor het geschikt is voor randapparaten met beperkte middelen.

De auteurs merken beperkingen op, specifiek dat de methode vertrouwen stelt op empirische schattingen van class priors en centroiden, wat onbetrouwbaar kan zijn in semi-supervised of unsupervised instellingen, en dat het voordeel afneemt in strikt IID-dataverdelingen waar de methode effectief reduceert tot standaard Fed-VPT. Echter, binnen de context van non-IID federated learning, biedt de voorgestelde aanpak een robuuste, theoretisch onderbouwde en empirisch superieure oplossing.