FedRE: A Representation Entanglement Framework for Model-Heterogeneous Federated Learning

TL;DR

FedRE tackles model-heterogeneous federated learning by introducing entangled representations that blend local cross-category information with normalized random weights, enabling a shared global classifier without sharing heterogeneous architectures. The framework uploads a single entangled representation plus an entangled-label encoding per client, providing cross-category supervision while mitigating privacy risks and reducing communication overhead. Empirical results demonstrate competitive accuracy and favorable privacy/communication trade-offs across diverse datasets and heterogeneity settings, including partial participation. The work highlights entangled representations as a lightweight, privacy-preserving alternative to sharing full representations or prototypes, with robust performance in both heterogeneous and homogeneous scenarios.

Abstract

Federated learning (FL) enables collaborative training across clients without compromising privacy. While most existing FL methods assume homogeneous model architectures, client heterogeneity in data and resources renders this assumption impractical, motivating model-heterogeneous FL. To address this problem, we propose Federated Representation Entanglement (FedRE), a framework built upon a novel form of client knowledge termed entangled representation. In FedRE, each client aggregates its local representations into a single entangled representation using normalized random weights and applies the same weights to integrate the corresponding one-hot label encodings into the entangled-label encoding. Those are then uploaded to the server to train a global classifier. During training, each entangled representation is supervised across categories via its entangled-label encoding, while random weights are resampled each round to introduce diversity, mitigating the global classifier's overconfidence and promoting smoother decision boundaries. Furthermore, each client uploads a single cross-category entangled representation along with its entangled-label encoding, mitigating the risk of representation inversion attacks and reducing communication overhead. Extensive experiments demonstrate that FedRE achieves an effective trade-off among model performance, privacy protection, and communication overhead. The codes are available at https://github.com/AIResearch-Group/FedRE.

Figures (9)

• Figure 1: FedRE framework. Each client maintains a local model consisting of a representation extractor and a classifier. The client’s local representations and their corresponding one-hot label encodings are integrated into a single entangled representation and entangled-label encoding, respectively, which are then uploaded to the server for training the global classifier.
• Figure 2: A toy experiment is conducted with 300 training and 200 test two-dimensional samples distributed across two clients. FedAllRep uploads all 300 representations, achieving the best performance (63.50%). FedGH uploads 4 prototypes, which may lead to increased focus on the prototypes, yielding sharper decision boundaries and slightly lower performance (60.50%). FedRE uploads 2 entangled representations, which provide cross-category supervision, resulting in smoother decision boundaries and competitive performance (62.00%).
• Figure 3: Accuracy (%) comparison between distinct communication rounds on the TinyImageNet dataset in the model-heterogeneous setting.
• Figure 4: Comparison of privacy protection in restructuring results from representations, prototypes, and entangled representations on the TinyImageNet dataset.
• Figure 5: Comparison of privacy protection in restructuring results from representations, prototypes, and entangled representations on the TinyImageNet dataset.