FedTGP: Trainable Global Prototypes with Adaptive-Margin-Enhanced Contrastive Learning for Data and Model Heterogeneity in Federated Learning

Jianqing Zhang; Yang Liu; Yang Hua; Jian Cao

Source-linked AI summary

FedTGP: Trainable Global Prototypes with Adaptive-Margin-Enhanced Contrastive Learning for Data and Model Heterogeneity in Federated Learning

Jianqing Zhang, Yang Liu, Yang Hua, Jian Cao

arXiv:2401.03230v1Jan 6, 2024cs.LGcs.CRcs.DC

TL;DR

Prototype-based HtFL lowers communication and protects client models, but weighted prototype aggregation can produce suboptimal global knowledge under heterogeneous data and models. FedTGP learns trainable global prototypes on the server with adaptive-margin-enhanced contrastive learning, and experiments report improved performance across heterogeneous settings, including up to 18.96% over FedProto.

Problem
Prototype-based HtFL reduces communication and protects client models, but naive weighted averaging can produce suboptimal global prototypes under statistical and model heterogeneity.
Method
FedTGP trains server-side global prototypes with adaptive-margin-enhanced contrastive learning to improve class separability while preserving semantics.
Results
FedTGP outperforms FedProto by up to 18.96% across four datasets, twelve heterogeneous models, and two statistically heterogeneous settings.
Takeaways & Limitations
FedTGP shares class-wise prototypes while enhancing inter-class separability without accessing client models or features.
Takeaways & Limitations
The formulation assumes clients collaboratively train heterogeneous models on private and heterogeneous datasets.

Abstract

from arXiv · show

Recently, Heterogeneous Federated Learning (HtFL) has attracted attention due to its ability to support heterogeneous models and data. To reduce the high communication cost of transmitting model parameters, a major challenge in HtFL, prototype-based HtFL methods are proposed to solely share class representatives, a.k.a, prototypes, among heterogeneous clients while maintaining the privacy of clients' models. However, these prototypes are naively aggregated into global prototypes on the server using weighted averaging, resulting in suboptimal global knowledge which negatively impacts the performance of clients. To overcome this challenge, we introduce a novel HtFL approach called FedTGP, which leverages our Adaptive-margin-enhanced Contrastive Learning (ACL) to learn Trainable Global Prototypes (TGP) on the server. By incorporating ACL, our approach enhances prototype separability while preserving semantic meaning. Extensive experiments with twelve heterogeneous models demonstrate that our FedTGP surpasses state-of-the-art methods by up to 9.08% in accuracy while maintaining the communication and privacy advantages of prototype-based HtFL. Our code is available at https://github.com/TsingZ0/FedTGP.

Introduction

HtFL reduces communication and privacy concerns by sharing global knowledge instead of client models, but prototype averaging can shrink class-separation margins. FedTGP addresses this with trainable global prototypes and adaptive-margin-enhanced contrastive learning, achieving stronger reported performance across heterogeneous settings.

Prototype-based HtFL shares class representatives instead of private model parameters to reduce communication and protect client models.
Weighted averaging can produce uninformative global prototypes because heterogeneous client prototypes differ in scale and separation margins.
FedTGP trains global prototypes with adaptive-margin-enhanced contrastive learning to improve separability while preserving semantic meaning.
Adaptive margins retain the maximum prototype margin among clients at each iteration, avoiding overly large early margins and preserving the best available separability.
FedTGP outperforms FedProto by up to 18.96% across four datasets, twelve heterogeneous models, and two statistically heterogeneous settings.
The method’s contributions include identifying margin shrinkage from naive averaging and demonstrating superiority over FedProto and other HtFL methods.

Related Work

Related HtFL methods accommodate heterogeneous clients through shared architectures, partial model sharing, global-data knowledge transfer, auxiliary models, or prototypes. FedTGP instead performs server-side prototype learning without accessing client models or features, addressing limitations of these alternatives.

Submodel-sharing methods accommodate diverse capabilities but still raise concerns about sharing client model architectures.
Split-model methods share selected layers but require clients to retain codependent model architectures.
Knowledge-distillation methods transfer predictions through a global dataset, which can be difficult to obtain.
FedProto and FedPCL share prototypes, but naive weighted averaging yields subpar global prototypes under statistical and model heterogeneity.
FedPCL relies on pretrained models for projection-network training, which is difficult when clients have private architectures and scarce data.
FedTGP learns prototypes on the server solely from client prototypes, without accessing client models or features.

Method

FedTGP replaces weighted averaging of heterogeneous client prototypes with trainable server-side global prototypes optimized by adaptive-margin-enhanced contrastive learning. The method preserves class semantics while improving inter-class separability and guides local client training using the learned prototypes.

Prototype-based HtFL: FedTGP splits each heterogeneous client model into a feature extractor and classifier, then exchanges class prototypes rather than model parameters.The feature extractor maps inputs to a shared K-dimensional feature space, while the classifier maps features to classes.
Prototype-based HtFL: Weighted averaging can produce ineffective global prototypes because heterogeneous clients generate representations with different scales, separability, and margins.The averaging weights depend on private class-distribution information and may shrink separation margins.
Trainable Global Prototypes: FedTGP learns trainable global prototypes on the server and uses contrastive objectives to align each prototype with same-class client prototypes while separating other classes.A shared server network transforms trainable class vectors into global prototypes, and all classes interact during gradient updates even with partial client participation.
Training Procedure: The server repeatedly samples clients, sends global prototypes, receives client prototypes, updates the adaptive margin, and optimizes the global prototypes.The resulting separable prototypes are sent in the next iteration to guide local feature representations toward greater inter-class separation.
Adaptive-Margin-Enhanced Contrastive Learning: Adaptive-margin-enhanced contrastive learning sets a positive margin between classes and chooses the maximum observed cross-class client-prototype margin subject to threshold τ.The adaptive margin avoids overly large margins early in training, when heterogeneous feature extractors are still weak.
Communication and Privacy: FedTGP retains prototype-based communication and privacy benefits by transmitting compact class prototypes without sharing model parameters or private class-distribution counts.The paper states that the low-dimensional prototypes are irreversible and that clients need not upload private class-distribution information.

Experiments

FedTGP is evaluated across heterogeneous datasets, model architectures, classifiers, and statistical settings, where it consistently outperforms the compared methods.

Experimental setup: Experiments cover four image datasets, six HtFL baselines, and heterogeneous feature extractors under pathological and practical statistical settings.The main model group uses eight architectures and feature dimension K = 512.
Main results: 9.08% is FedTGP’s maximum accuracy improvement over all baselines across four datasets.Using TGP with ACL improves over FedProto by up to 13.85%.
Model heterogeneity: 5.64% is FedTGP’s maximum advantage across increasing model heterogeneity, irrespective of model sizes.FedTGP’s accuracy decreases by 1.77%, compared with 3.53%–15.04% for counterparts as heterogeneity increases.
Heterogeneous classifiers: 18.96% is FedTGP’s accuracy advantage over FedProto in the HtFE8-HtC4 setting with heterogeneous feature extractors and classifiers.LG-FedAvg and FedGen are not applicable in these heterogeneous-classifier scenarios.

Partial Participation with More Clients

With partial client participation and more clients, all methods lose accuracy as each client receives less data, but FedTGP’s relative superiority becomes more pronounced.

Experimental setting: 50% of clients participate in each training iteration, while experiments compare 50-client and 100-client Cifar100 settings.The evaluation uses HtFE8 and the practical statistical setting.
Results: 4.66% and 5.48% are FedTGP’s advantages over other methods with 50 and 100 clients, respectively.FedTGP’s superiority is more obvious with more clients.
Results: Larger client counts reduce every method’s performance because the data amount on each client decreases.The passage reports this trend for Cifar100 with partial client participation.
Client training epochs: Increasing client training epochs reduces accuracy for auxiliary-model methods such as FML and FedKD.The decrease is attributed to increased heterogeneity in the shared auxiliary model parameters.

Impact of Feature Dimensions

Feature dimension affects performance non-monotonically: most methods improve through K = 256 but degrade at K = 1024, while FedTGP remains competitive at K = 64.

Feature dimension: 6.45% is FedProto’s performance deficit at K = 64 compared with K = 256.FedTGP achieves competitive performance with K = 64.
Communication cost: Prototype-based methods communicate class prototypes, whereas auxiliary-model methods transmit substantially larger model-related information.FML and FedKD have the highest communication overhead because they additionally transmit an auxiliary model.
Communication cost: FedKD’s communication cost remains much larger than prototype-based methods despite SVD compression of its auxiliary-model parameters.FedGen also incurs download cost for its server-side generator.

Ablation Study

Ablations indicate that margins, adaptive margins, and trainable prototype parameters each contribute to FedTGP’s performance, while hyperparameters trade accuracy against stability and computation.

Margin and contrastive learning: 12.47% is the improvement from adaptive-margin ACL for FedProto on Cifar100, exceeding the 10.82% gain from a fixed margin.Standard contrastive learning improves FedProto by 5.60%.
Trainable prototype parameters: 6.57% is the maximum performance decrease when the further processing model F is removed from TGP.The reduced version still outperforms FedProto by a large gap.
Hyperparameters: An excessively large threshold τ = 1000 slightly reduces accuracy because prototype guidance becomes unstable and δ(t) may keep growing.Larger τ values from 1 to 100 generally improve performance.
Hyperparameters: Increasing server training epochs S improves accuracy, but the gain from S = 100 to S = 1000 is negligible.The default S = 100 saves computation, while even S = 1 achieves at least 43.23% accuracy.

Conclusion

FedTGP shares class-wise prototypes between server and clients and uses TGP with ACL to enhance class separability. Experiments across two heterogeneous settings and twelve models show superiority over baseline methods.

FedTGP shares class-wise prototypes between the server and clients while enhancing separation between different classes through TGP and ACL.

Additional Experimental Details

The experiments use fixed hardware and baseline-specific hyperparameters, with FedTGP configured using shared settings across tasks. Model heterogeneity is evaluated through heterogeneous feature extractors and classifier configurations.

Experimental environment: Experiments use 64 Intel Xeon CPUs, 256G memory, eight NVIDIA 3090 GPUs, and Ubuntu 20.04.4 LTS.
Hyperparameter settings: FedTGP uses λ = 0.1, margin threshold τ = 100, and 100 server learning epochs across all tasks.
Model heterogeneity: Model heterogeneity is evaluated by separating classifier and feature-extractor roles, including heterogeneous classifiers in additional settings.
Model architectures: The HtFE8 architectures include single- and two-layer fully connected configurations with output dimensions ending in 100-way classification layers.
Computation measurement: FLOPs estimates include trainable forward and backward operations, with the backward pass approximated as twice the forward-pass cost.

Additional Experimental Results

Additional experiments evaluate FedTGP across datasets, model settings, accuracy, and computation. The results report continued performance advantages and a substantial iteration-time efficiency advantage over FedProto.

Fashion-MNIST: FedTGP is additionally evaluated on Fashion-MNIST using the HtCNN8 model group.The evaluation covers both pathological and practical statistical heterogeneity settings.
Model architectures: The HtCNN8 model group contains eight architectures designed for grayscale Fashion-MNIST images.
Homogeneous models: FedTGP also outperforms other methods when clients use homogeneous models on Cifar100.
Model scale: Larger models reduce performance for all methods when feature extractors are not shared, because local data scarcity makes deeper feature-extractor training challenging.
Computation cost: FedTGP matches FedProto’s client computation overhead per iteration while requiring fewer iterations to reach the reported accuracy.FedTGP reaches 36.34% accuracy in 17 iterations, whereas FedProto requires 489 iterations.

Visualizations

Visualizations indicate that FedTGP converges stably, preserves larger prototype margins, and produces more separated feature representations than FedProto. The figures also document the client data distributions used across datasets and settings.

Training Error Curve: FedTGP optimizes quickly during the initial 50 iterations and maintains stable performance after converging around the 100th iteration.
Prototype Visualization: FedProto’s weighted averaging produces smaller prototype margins, whereas FedTGP pushes global prototypes apart while retaining prototype semantics.
Feature Representations: FedTGP yields more discriminative and more closely clustered cross-client feature representations than FedProto.FedProto’s representations show class overlap or mixing, while FedTGP produces closer clustering for the same class across clients.
Data Distributions: The visualizations include client data distributions for pathological and practical settings across FMNIST, Cifar10, Flowers102, Cifar100, and Tiny-ImageNet.Circle size represents the number of samples in the distribution figures.