Heterogeneous Ensemble Knowledge Transfer for Training Large Models in Federated Learning

Yae Jee Cho; Andre Manoel; Gauri Joshi; Robert Sim; Dimitrios Dimitriadis

Source-linked AI summary

Heterogeneous Ensemble Knowledge Transfer for Training Large Models in Federated Learning

Yae Jee Cho, Andre Manoel, Gauri Joshi, Robert Sim, Dimitrios Dimitriadis

arXiv:2204.12703v1Apr 27, 2022cs.LG

TL;DR

Federated learning needs to train large server models despite clients’ limited resources, heterogeneous data, and differing model architectures. Fed-ET transfers knowledge from heterogeneous small client models to a larger server model using weighted consensus distillation and diversity regularization. Across image and language tasks, it reports higher accuracy, fewer communicated parameters, and robustness to high data heterogeneity.

Problem
Federated learning commonly requires identical client and server architectures, while limited client resources and heterogeneous data make large-model training and naive aggregation impractical.
Method
Fed-ET trains heterogeneous small client models and transfers their knowledge to a larger server model through unlabeled-data-based weighted consensus distillation with diversity regularization.
Results
Across image and language classification tasks, Fed-ET achieves higher test accuracy, fewer communication rounds, and greater robustness to data heterogeneity than other FL algorithms.
Takeaways & Limitations
Fed-ET enables large server models while clients train smaller heterogeneous models without additional client computation or access to data beyond their private datasets.
Takeaways & Limitations
Fed-ET assumes that only the server has access to an unlabeled public dataset used for knowledge transfer.

Abstract

from arXiv · show

Federated learning (FL) enables edge-devices to collaboratively learn a model without disclosing their private data to a central aggregating server. Most existing FL algorithms require models of identical architecture to be deployed across the clients and server, making it infeasible to train large models due to clients' limited system resources. In this work, we propose a novel ensemble knowledge transfer method named Fed-ET in which small models (different in architecture) are trained on clients, and used to train a larger model at the server. Unlike in conventional ensemble learning, in FL the ensemble can be trained on clients' highly heterogeneous data. Cognizant of this property, Fed-ET uses a weighted consensus distillation scheme with diversity regularization that efficiently extracts reliable consensus from the ensemble while improving generalization by exploiting the diversity within the ensemble. We show the generalization bound for the ensemble of weighted models trained on heterogeneous datasets that supports the intuition of Fed-ET. Our experiments on image and language tasks show that Fed-ET significantly outperforms other state-of-the-art FL algorithms with fewer communicated parameters, and is also robust against high data-heterogeneity.

1 Introduction

Federated learning must reconcile clients’ limited resources and heterogeneous data, which make identical large models and naive aggregation impractical. Fed-ET addresses this by transferring knowledge from heterogeneous small client models to a larger server model and reports stronger accuracy with fewer communication rounds.

Motivation: Clients’ limited system resources restrict them to smaller models, while heterogeneous datasets can make naive model aggregation hinder convergence.These constraints can cause the globally trained model to perform poorly in practice.
Research challenge: Fed-ET asks how an ensemble of differently architected models trained on heterogeneous datasets can train a larger server model.The approach targets heterogeneous client and server architectures rather than requiring one shared model architecture.
Method: Fed-ET uses unlabeled data for ensemble knowledge transfer between smaller client models and a larger server model.The method is designed for resource-constrained, data-heterogeneous clients.
Method: Its weighted consensus distillation filters heterogeneous experts while diversity regularization preserves useful information from models that differ from the consensus.The method is accompanied by corresponding generalization bounds.
Results: Fed-ET achieves higher test accuracy than other FL algorithms on image and language classification tasks, with greater robustness to data heterogeneity and fewer communication rounds.The reported comparison covers both model-homogeneous and model-heterogeneous FL baselines.

2 Background and Related Work

Prior ensemble-transfer methods assume similarly trained experts, whereas federated learning combines heterogeneous data and partial client participation. Fed-ET adapts ensemble transfer with a server-to-client feedback loop, confidence-weighted consensus, and diversity regularization.

Ensemble Knowledge Transfer: Earlier ensemble knowledge-transfer methods typically use experts trained on similar datasets, often produced through augmentation or data shuffling.Their setting differs from federated learning’s heterogeneous local data distributions.
Fed-ET: Fed-ET establishes a feedback loop that transfers knowledge from the target server model back to the ensemble of client models.This complements data-aware weighted consensus from the ensemble to the server.
Related FL Methods: Existing FL knowledge-distillation methods can be limited by labeled public data, image-only tasks, or low data heterogeneity.Fed-ET is presented as not restricted to these scenarios.
Fed-ET: Fed-ET weights clients with higher inference confidence more heavily while allowing non-consensus clients to transfer useful representations through diversity regularization.This combines consensus filtering with exploitation of ensemble diversity.

3 Federated Ensemble Transfer: Fed-ET

Fed-ET trains a large server model from small, heterogeneous client models through three ensemble-transfer steps, using weighted consensus and diversity regularization to handle heterogeneous data.

Fed-ET uses an unlabeled public dataset and three steps: client training and representation transfer, weighted consensus distillation with diversity regularization, and server representation transfer.The server performs the additional ensemble-transfer computations while clients conduct local training as in standard federated learning.
Step 1: Client Local Training & Representation Transfer: Clients receive designated small models according to their resource capabilities, train them locally on private data, and return updated models and representations to the server.The client models may differ in architecture, while their representation layers are shared with the large server model.
Step 3: Server’s Representation Transfer: The server uses the ensemble-derived labels and representations from unlabeled public data to train and update its larger model.The large server model is assumed to have substantially more parameters than the small client models.
Step 2: Ensemble Loss by Weighted Consensus with Diversity Regularization: The server forms a weighted consensus from client logits on public data, assigning greater weight to models with higher confidence measured through logit variance.This consensus is intended to reduce the influence of low-confidence non-experts and overly confident erroneous outliers.
Step 2: Ensemble Loss by Weighted Consensus with Diversity Regularization: Diversity regularization preserves useful representation information from clients whose predictions are less confident on particular public samples.The method combines consensus knowledge with diversity from models trained on different data distributions.
3.3 Generalization Bound for Ensemble Transfer: The generalization bound identifies local model quality, client dataset size, and client–server distribution discrepancy as factors that worsen ensemble-transfer performance.Fed-ET addresses these factors through confidence-based weighting and client sampling proportional to dataset size.

4 Experiments

Fed-ET is evaluated against homogeneous and heterogeneous FL baselines on image and language tasks under varying data heterogeneity. It achieves strong accuracy and communication efficiency while using smaller client models and benefiting from diversity regularization.

Experimental setup: Experiments compare Fed-ET with model-homogeneous and model-heterogeneous baselines using large-server-model evaluation where applicable.The study uses partial participation and three random seeds; image and language model configurations differ between small and large server models.
Accuracy: 78% and 35% test accuracy are achieved by Fed-ET on CIFAR10 and CIFAR100 at α = 0.1, exceeding MOON’s 75% and 33%.Fed-ET uses smaller client models than the large VGG19 communicated directly by MOON.
Accuracy: Fed-ET achieves higher test accuracy on CIFAR10 with α = 0.1 and Sent140 than both model-homogeneous and model-heterogeneous baselines.At α = 0.5, MOON slightly out-performs Fed-ET by around 1% while training larger client models.
Communication efficiency: Approximately 3× fewer communicated parameters are required by Fed-ET to reach target test accuracy than the model-homogeneous baselines.The baselines communicate the large server model, whereas Fed-ET communicates smaller client models each round.
Diversity regularization: λ = 0.05 gives the best image-task performance, while λ = 0.5 gives the best language-task performance in the reported diversity-regularization experiment.λ controls the diversity regularization term; the best value depends on the task.

5 Conclusion

Fed-ET trains large server models from smaller, heterogeneous client models through data-aware weighted consensus distillation without adding client overhead. It achieves high accuracy, lower communication overhead, and robustness to data heterogeneity.

Conclusion: Fed-ET trains a large server model with smaller client models using data-aware weighted consensus distillation.The framework is designed for heterogeneous client architectures and data distributions.
Conclusion: Fed-ET achieves high test accuracy with significantly lower communication overhead and robustness against data heterogeneity.The conclusion identifies heterogeneous-model deployment strategies and broader ensemble knowledge transfer as future directions.

A Proof for Theorem 1

The proof establishes a generalization bound by combining domain-adaptation and finite-sample arguments. It relates population and empirical losses through bounded-loss concentration and distribution discrepancies.

Supporting lemmas: Lemma 1 bounds domain-adaptation error using distribution discrepancy between two true distributions and the joint optimal error term ν.The bound holds with probability at least 1 − δ over the sample choice.
Supporting lemmas: Lemma 2 shows that limited client sample size increases generalization error through discrepancy between Dk and its empirical distribution.The relevant sample-size variable mk is the number of training samples of client k.
Proof strategy: Hoeffding’s inequality bounds the gap between population loss LDk(ĥDk) and empirical loss L̂Dk(ĥDk) because the losses lie within [0, 1].The proof treats the empirical loss as an average of bounded random variables and rearranges the resulting inequality.
Proof strategy: Theorem 1’s proof concludes after applying the preceding bounds to the generalization analysis.The supplied proof passages show the final use of Lemma 2 and completion of the argument.

B Comparisons with FedGKT

FedGKT is adapted for comparison under partial participation and heterogeneous CIFAR10 partitions. In this setting, FedGKT is less robust to high data heterogeneity than Fed-ET, especially at α = 0.1.

Comparison setting: FedGKT differs from the studied setting because its original experiments use 16 clients with every client participating in every round.The comparison instead samples a few clients from a larger pool each round.
Adaptation: Adapting FedGKT requires on-the-fly server-logit computation and storing client models for later loading because clients are updated at random intervals.These changes arise from partial participation and differ substantially from the original FedGKT setup.
Results: FedGKT fails to be robust against high data heterogeneity and partial-client participation compared with Fed-ET, with the largest performance gap at α = 0.1.The comparison uses CIFAR10 partitions formed with α ∈ {0.1, 0.5}, sampling 10 clients from 100.
Results: The CIFAR10 comparison reports test accuracy for Fed-ET and FedGKT under partitions created with α = 0.1 and α = 0.5.The table caption identifies the two heterogeneity settings; the reported runs use three different partitions.

C Details of Experiment Setup

The experiments evaluate heterogeneous federated learning under controlled data splits, varied model architectures, tuned training settings, and a fixed GPU-cluster platform.

Dataset: Image datasets are split into training, public, and test sets at 7:1:2, then distributed across 100 clients using DirK(α) with α = 0.1 and α = 0.5.Smaller α produces greater client data-size imbalance and label skew.
Model configuration: The setup uses CNN, modified pretrained VGG and ResNets, pretrained Tiny-BERT and Mini-BERT, and a two-layer LSTM binary classifier.The CNN has two convolutional and four fully connected layers, while the LSTM uses 256 hidden units.
Model configuration: The experiment reports trainable parameter counts for the configured models in Table 6, using M = 10^6 as the unit.The table is identified as the parameter-count summary for the experimental models.
Hyperparameters and training: Image-task hyperparameters are selected by grid search over learning rate, batch size, and local iterations, without learning-rate decay.The search selects settings maximizing test accuracy for each benchmark.
Hyperparameters and training: Server-side settings are tuned over learning rate, public batch size, SGD iterations, and regularization weight, with image-task values η = 0.005, bs = 64, τs = 128, and λ = 0.05.The reported image-task local settings are η = 0.1, b = 64, and τ = 30.
Platform: Experiments run on clusters equipped with one NVIDIA TitanX GPU per cluster, using Ethernet communication and PyTorch implementations.Each machine is treated as one federated-learning client, and the number of clusters matches the selected-client fraction.