Source-linked AI summary

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

Yuan Yao, Lixu Wang, Jiaqi Wu, Jin Song, Simin Chen, Zehua Wang, Zijian Tian, Wei Chen, Huixia Li, Xiaoxiao Li

arXiv:2511.22265v2cs.LG

TL;DR

Model-heterogeneous federated learning is motivated by impractical assumptions of shared client architectures and trade-offs among existing forms of client knowledge. FedRE entangles representations and labels across categories for global-classifier training; the authors report strong benchmark comparisons and a balance among performance, privacy, and communication overhead.

  • Problem

    Model-heterogeneous FL is motivated by unrealistic shared-architecture assumptions and client-knowledge options that trade off communication, privacy, and representation of local information.

    Source: §1 · pp. 1–2
  • Method

    FedRE combines each client’s representations across categories and their one-hot labels with the same normalized random weights, uploading both entangled encodings to train a global classifier.

    Source: §1 · p. 2, §3.1 · p. 3
  • Results

    Across varied model-heterogeneous scenarios, FedRE outperforms all baselines; on TinyImageNet under PAT, accuracy surpasses LG-FedAvg by 6.26%.

    Source: §4.2 · p. 5
  • Takeaways & Limitations

    The authors conclude FedRE offers a balanced trade-off among model performance, privacy protection, and communication overhead.

    Source: §5 · p. 8

    A fully rigorous non-convex convergence analysis is left for future work because local classifier replacement and representation entanglement add perturbations.

    Source: §3.3.3 · p. 5

Abstract

from arXiv · show

Federated learning (FL) enables collaborative training across clients while preserving privacy. While most existing FL methods assume homogeneous model architectures, client heterogeneity in both data and resources makes this assumption impractical, thus 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. Specifically, each client aggregates its local representations into a single entangled representation using normalized random weights, and then applies the same weights to integrate the corresponding one-hot label encodings into an entangled-label encoding. Both are subsequently 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 re-sampled at each round to introduce diversity, alleviating overconfidence in the global classifier and yielding smoother decision boundaries. Moreover, each client uploads a single 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.

1. Introduction

Model heterogeneity makes the homogeneous-architecture assumption impractical and direct parameter aggregation infeasible, while existing client-knowledge options involve trade-offs. FedRE introduces cross-category entangled representations and reports a balanced trade-off among performance, privacy, and communication overhead.

1. Introduction

  • Different client data, hardware, and computational capabilities make homogeneous architectures unrealistic, and heterogeneous representation extractors make direct parameter aggregation infeasible.The server or other clients may also lack access to the clients’ architectures.Source: §1 · p. 1
  • Representations, logits, and small models can encode high-level knowledge but may raise communication costs and inversion risks, whereas classifiers and prototypes have other limitations.Classifiers may inherit local-data biases, while prototypes capture category-level information with limited intra-class variability.Source: §1 · p. 2
  • FedRE combines local representations across categories into one entangled representation per client, using the same normalized random weights to combine their one-hot labels.Both encodings are uploaded to train a global classifier; resampling weights each round introduces diversity and supports cross-category supervision.Source: §1 · p. 2
  • FedRE is reported to achieve an effective trade-off among model performance, privacy protection, and communication overhead.Blending cross-category information is described as mitigating representation-inversion risk, while one representation per client reduces communication overhead.Source: §1 · p. 2

2. Related Work

Related work distinguishes model-heterogeneous FL, where direct parameter aggregation is infeasible, from model-homogeneous FL. These approaches aggregate or otherwise use different forms of client knowledge and model components.

2. Related Work

  • The related work groups FL methods into model-heterogeneous and model-homogeneous approaches.The distinction is whether methods address local-model heterogeneity as well as sample-distribution heterogeneity.Source: §2 · p. 2
  • Model-heterogeneous FL handles differences in local models and sample distributions, making direct parameter aggregation infeasible and motivating alternative forms of client knowledge.Examples include logits, aligned representations, classifiers, prototypes, knowledge distillation, and homogeneous small models.Source: §2 · p. 2
  • Model-homogeneous FL maintains homogeneous local models while handling sample-distribution differences through aggregation of all or partial model parameters and other techniques.Examples include FedAvg’s all-parameter aggregation and FedRep’s aggregation of representation extractors.Source: §2 · p. 2

3. Methodology

FedRE addresses model-heterogeneous federated learning with client-specific representation extractors and a shared classifier architecture. Clients entangle representations and labels into pairs for server-side global-classifier training, targeting accuracy, privacy, and communication overhead.

3. Methodology

  • On the toy experiment, FedAllRep scores 63.50%, FedRE 62.00%, and FedGH 60.50%.The figure depicts FedRE's gradual color transitions alongside its smoother decision boundaries.Source: Figure 2 · p. 3

3.1. Motivation

  • FedAllRep uploads all representations, risking inversion and communication overhead, while FedGH uses category prototypes that may bias the global classifier toward sharper boundaries.FedRE instead blends representations across categories, providing cross-category supervision and potentially smoother boundaries while reducing uploads and increasing inversion difficulty.Source: §3.1 · p. 3

3.2.1. Local Model Update

  • Each client updates its local model by minimizing cross-entropy loss over its private dataset.The local model comprises a representation extractor and classifier; extractors may vary across clients while classifier architectures are shared.Source: §3 · p. 3, §3.2.1 · p. 4

3.2. Representation Entanglement and Upload

  • Clients map representations to a consistent dimensionality, then use normalized random weights to upload one entangled representation and its corresponding entangled-label encoding for global-classifier training.Weights are re-sampled each round to introduce diversity; the server broadcasts its updated classifier to clients.Source: §3.2.2 · p. 4, §3.2.3 · p. 4

3.2.3. Global Classifier Update and Broadcast

  • The server trains the global classifier on uploaded representation-label pairs and broadcasts the updated classifier to clients.Its objective minimizes cross-entropy between the classifier's output on each entangled representation and the corresponding entangled-label encoding.Source: §3.2.3 · p. 4

3.3.1. Computational Complexity of RE

  • For representation matrix R and one-hot label matrix Y, the entangled vectors are computed as er = R⊤w and ey = Y⊤w using normalized weights w.The passage introduces n as the sample count and d as representation dimensionality but does not provide a complete complexity expression.Source: §3.3.1 · p. 4

3.3.2. Comparison with Related Studies

  • FedRE uploads one entangled representation per client, and random weighting across rounds is described as improving generalization under client statistical heterogeneity.The paper distinguishes this all-representation aggregation from Mixup's pairwise interpolation.Source: §3.3.2 · p. 4

3.3.3. Convergence Analysis of FedRE

  • FedRE's performance improves and stabilizes across communication rounds, while a fully rigorous non-convex convergence analysis is left for future work.The analysis is described as non-trivial because local classifier replacement and representation entanglement add perturbations.Source: §3.3.3 · pp. 4–5

4. Experiments

Experiments evaluate FedRE on CIFAR-10, CIFAR-100, and TinyImageNet against model-heterogeneous FL baselines, including under varied client participation and statistical heterogeneity. FedRE is reported to outperform baselines, protect against representation inversion, and reduce upload overhead; analyses examine its design choices and costs.

4.2. Main Experiments

  • FedRE outperforms all baselines across the reported model-heterogeneous scenarios, exceeding LG-FedAvg, FedGH, and FedKD by 6.26%, 6.54%, and 6.79% on TinyImageNet under PAT.The comparison covers three datasets; Table 1 reports accuracy under the model-heterogeneous setting.Source: §4.2 · p. 5, Table 1 · p. 6
  • Reconstructed entangled representations reveal no identifiable information; their PSNR is 9.66 and MSE is 7781.87, compared with 12.89 and 4514.91 for representations and 10.25 and 6992.04 for prototypes.Lower PSNR and higher MSE indicate stronger privacy protection in the reported inversion-attack evaluation.Source: §4.2 · p. 6, Figure 4 · p. 6
  • On CIFAR-100, FedRE has the lowest upload overhead and broadcast overhead comparable to classifier- and prototype-based methods.FedRE uploads one entangled representation and its entangled-label encoding per client; Table 2 compares communication overhead in thousands of scalars.Source: §4.2 · p. 6, Table 2 · p. 7

4.3. Analysis

  • With 100 clients and partial participation on CIFAR-10, FedRE achieves the best performance in most tested scenarios with highly heterogeneous distributions.The tests vary participation rates, Dirichlet α, and long-tail imbalance factors.Source: §4.3 · pp. 6–7, Table 3 · p. 7
  • On TinyImageNet under PRA and PAT, FedRE achieves performance comparable to FedAllRep while reducing communication overhead relative to uploading all local representations.Table 4 compares accuracy and communication overhead.Source: §4.3 · p. 7, Table 4 · p. 7
  • Average Pooling is selected as the default RM operation, while prototype-based entanglement outperforms representation-based entanglement and random weights outperform equal weights.The RM comparison evaluates average pooling, max pooling, and a fully connected layer; the RE comparison selects Random Average Prototype (RAP) as the default.Source: §4.3 · pp. 7–8, Table 5 · p. 7, Table 6 · p. 8
  • RAP performs comparably with Uniform, Laplace, and Gaussian weight distributions, with a minor advantage for Uniform; per-round resampling achieves higher accuracy than fixed sampling on both tested datasets.The resampling comparison uses a synthetic two-client dataset and CIFAR-100; Table 8 reports the accuracy comparison.Source: §4.3 · p. 8, Table 7 · p. 8, Table 8 · p. 8
  • Adding RE to local training incurs only slight additional per-round training cost because representation and label aggregation requires no additional gradient computations.The cost comparison averages across 10 clients and 6 communication rounds on CIFAR-10 and CIFAR-100.Source: §4.3 · p. 8, Table 9 · p. 8

5. Conclusion

FedRE presents entangled representations as lightweight, privacy-aware client knowledge, with multiple mechanisms for forming a single client representation.

5. Conclusion

  • FedRE sends one cross-category entangled representation and its associated label encoding for global classifier training, balancing performance, privacy, and communication.The conclusion identifies broader machine-learning tasks as a future direction.Source: §5 · p. 8

A. Mathematical Details of Various Representation Entanglement Mechanisms

  • The general mechanism combines a client's representations and corresponding labels with the same weights to form paired entangled outputs.The representation set is defined as the client's collection of representations and labels.Source: A. Mathematical Details of Various… · p. 11
  • RSR selects one representation per client per round, VAR averages all representations equally, and RAR uses normalized random weights drawn from a uniform distribution.Under RAR, the random weights are normalized to sum to one.Source: A. Mathematical Details of Various… · p. 11
  • VAP averages each client's prototypes equally, RAP combines prototypes with normalized random weights, and RSP selects one prototype per client each round.For selected-prototype weighting, representations in the selected prototype's category receive weight 1/n_c.Source: A. Mathematical Details of Various… · p. 11

B. Detailed Experimental Setup

The setup describes a semi-honest server's representation-inversion threat and standardizes the reconstruction model across experiments.

B. Detailed Experimental Setup

  • The experimental setup covers statistical heterogeneity, model training and configurations, and representation-inversion attacks.The setup assumes a semi-honest server with access to the representation extractor that attempts to reconstruct client samples.Source: B. Detailed Experimental Setup · p. 11, Table 10 · p. 12
  • The server's inversion network reconstructs a target representation, prototype, or entangled representation by optimizing its parameters with fixed noise.The optimization minimizes squared distance between the extractor's output on the reconstructed sample and the target vector.Source: B. Detailed Experimental Setup · p. 12
  • All reconstruction experiments use the same four-layer CNN extractor and the inversion network used in [37].Source: B. Detailed Experimental Setup · p. 12

C.1. Model-homogeneous FL Evaluation

FedRE performs strongly in model-homogeneous FL, including with few clients, and shows stable convergence on TinyImageNet.

C.1. Model-homogeneous FL Evaluation

  • 63.21% average accuracy, 2.58% above LG-FedAvg, is FedRE's best overall result across the three model-homogeneous datasets.Source: Table 11 · p. 12, §C.1 · p. 12
  • FedRE shows rapid initial improvement followed by gradual stabilization on TinyImageNet, indicating stable convergence.Source: §C.1 · p. 12, Figure 5 · p. 13

C.2. Impact of a Small Number of Clients

  • FedGH and FedRE both lose performance as client count falls, which the authors attribute to reduced diversity in prototypes or entangled representations.Source: §C.2 · p. 13
  • FedProto, FedGH, and FedRE consistently outperform FedAvg across the tested client counts.Source: §C.2 · p. 13
  • 41.57% on CIFAR-100 with 2 clients, FedRE exceeds FedAvg by 8.44%, FedProto by 3.70%, and FedGH by 1.96%.Source: Table 12 · p. 13, §C.2 · p. 13

C.3. Privacy Protection Evaluation

Reconstructed TinyImageNet images from entangled representations show less discernible content or category information than reconstructions from vanilla representations or prototypes.

C.3. Privacy Protection Evaluation

  • Entangled-representation reconstructions contain less discernible content or category information than vanilla-representation or prototype reconstructions, suggesting privacy protection against inversion attacks.Source: §C.3 · p. 13

C.4. Communication Overhead Evaluation

In the PRA setting, Table 13 compares FedRE and baselines’ communication overhead across three datasets and both model-heterogeneous and model-homogeneous scenarios; FedRE generally has lower overhead.

  • Table 13 compares communication overhead on CIFAR-10, CIFAR-100, and TinyImageNet under both model-heterogeneous and model-homogeneous scenarios in the PRA setting.Source: §C.4 · p. 13
  • FedRE generally achieves lower communication overhead than the baselines.Source: §C.4 · p. 13
  • The authors interpret the overhead results as suggesting potential effectiveness in reducing communication costs.Source: §C.4 · p. 13

C.5. Statistical Heterogeneity Analysis

FedRE achieves competitive accuracy across the evaluated levels of statistical heterogeneity on CIFAR-10 and CIFAR-100. The analysis varies sample skewness and presents sample distributions and comparisons across scenarios.

  • The analysis varies the Dirichlet parameter α in PRA and client participation rates in PAT to control sample skewness.Source: §C.5 · pp. 13–14
  • Figure 6 compares privacy protection in restructuring results from representations, prototypes, and entangled representations on TinyImageNet.Source: Figure 6 · p. 14
  • FedRE achieves competitive accuracy across different levels of statistical heterogeneity on CIFAR-10 and CIFAR-100 under the model-heterogeneous setting.Source: §C.5 · p. 14, Figure 9 · p. 15
  • Figures 7–8 show client sample distributions for CIFAR-10 and CIFAR-100 under varying α in PRA and varying client numbers in PAT; circle size indicates sample count.Source: §C.5 · p. 14, Figure 7 · p. 15, Figure 8 · p. 15
Loading 2511.22265v2…