PaperlayerView exact revision on arXiv

Source-linked AI summary

FedAUX: Leveraging Unlabeled Auxiliary Data in Federated Learning

Felix Sattler, Tim Korjakow, Roman Rischke, Wojciech Samek

arXiv:2102.02514v1cs.LGcs.DCstat.ML

TL;DR

Federated Distillation uses unlabeled auxiliary data to support heterogeneous client models, but its performance remains limited under client-data heterogeneity and compared with centralized training. FedAUX adds self-supervised pre-training and differentially private certainty-weighted distillation, improving performance across iid and non-iid settings and further closing the centralized-training gap.

  • Problem

    Federated Distillation still falls short of centralized training and deteriorates with heterogeneous client data despite access to unlabeled auxiliary data.

  • Method

    FedAUX pre-trains a feature extractor on auxiliary data and uses (ε, δ)-differentially private certainty scores to weight client predictions during distillation.

  • Results

    FedAUX substantially improves Federated Distillation on non-iid data and improves FL performance across baselines using a wide range of auxiliary datasets.

  • Takeaways & Limitations

    Using unlabeled auxiliary data for pre-training and weighted ensemble distillation can drastically improve FL performance and further close the gap to centralized training.

  • Takeaways & Limitations

    The privacy analysis covers the scoring mechanism, while privacy of the deep classifier training is outside this work; total privacy cost also includes communicating updated models.

Abstract

from arXiv · show

Federated Distillation (FD) is a popular novel algorithmic paradigm for Federated Learning, which achieves training performance competitive to prior parameter averaging based methods, while additionally allowing the clients to train different model architectures, by distilling the client predictions on an unlabeled auxiliary set of data into a student model. In this work we propose FedAUX, an extension to FD, which, under the same set of assumptions, drastically improves performance by deriving maximum utility from the unlabeled auxiliary data. FedAUX modifies the FD training procedure in two ways: First, unsupervised pre-training on the auxiliary data is performed to find a model initialization for the distributed training. Second, $(\varepsilon, δ)$-differentially private certainty scoring is used to weight the ensemble predictions on the auxiliary data according to the certainty of each client model. Experiments on large-scale convolutional neural networks and transformer models demonstrate, that the training performance of FedAUX exceeds SOTA FL baseline methods by a substantial margin in both the iid and non-iid regime, further closing the gap to centralized training performance. Code is available at github.com/fedl-repo/fedaux.

1. Introduction

Federated Distillation can use unlabeled auxiliary data but remains vulnerable to client-data heterogeneity and a gap to centralized training. FedAUX addresses this by combining self-supervised pre-training with privacy-preserving certainty-weighted distillation.

  • Federated Learning trains models across private client data without transferring that data to a centralized location.
  • Unlabeled auxiliary data can roughly match client-data distributions but lacks labels needed for task-specific training.
  • Federated Distillation supports heterogeneous client architectures by distilling client predictions on auxiliary data into a student model.
  • FD performance falls short of centralized training and deteriorates under non-iid client-data distributions.
  • FedAUX uses self-supervised pre-training and certainty-weighted FD to improve performance across baselines and substantially address non-iid data.
  • The method constrains certainty-score privacy loss with an (ε, δ)-differentially private mechanism while adding negligible client computation and quantifiable excess privacy loss.

2. Related Work

Federated Distillation transfers client knowledge through predictions on distillation data and supports heterogeneous architectures, while related ensemble methods improve aggregation through weighted or Bayesian combinations.

  • Distillation transfers knowledge from one or more classifiers to a student model using teacher predictions on a distillation dataset.
  • Federated Distillation aggregates client knowledge through distillation and can outperform conventional parameter-averaging methods.
  • FD methods also allow clients to train heterogeneous model architectures and can reduce communication overhead.
  • Prior work combines client predictions using equally weighted or certainty-weighted ensembles for classification.

3. Federated Learning with Auxiliary Data

FedAUX extends federated distillation by exploiting unlabeled auxiliary data through self-supervised initialization and certainty-weighted, privacy-preserving prediction aggregation. The method addresses performance degradation under non-iid client data while retaining distillation-based aggregation and heterogeneous-model support.

  • Setting: FedAUX assumes clients hold potentially non-iid private labeled data while clients and server share public unlabeled auxiliary data.The auxiliary distribution may differ from the private training distribution, and this setting is shared by FD studies.
  • Federated Ensemble Distillation: Federated ensemble distillation trains a student model on client predictions over public auxiliary data instead of averaging client parameters.The client predictions are aggregated into soft labels using a mechanism such as the mean before student optimization.
  • FedAUX Preparation: FedAUX uses self-supervised pre-training on auxiliary data to obtain a feature extractor for distributed training.The paper motivates contrastive representation learning for image data and next-token prediction for sequence data.
  • Motivation: Existing federated ensemble distillation degrades when client training data is heterogeneous, as does the training speed, stability, and maximum achievable accuracy of existing FL algorithms.The paper identifies this non-iid degradation as the performance gap FedAUX aims to reduce.
  • Weighted Ensemble Distillation: FedAUX replaces the standard prediction mean with a certainty-weighted average to better combine client hypotheses from heterogeneous source distributions.Client certainty is estimated through a local separation task using a pre-trained feature extractor, logistic regression, and normalized features.
  • Privacy Analysis: The certainty-scoring mechanism is made (ε, δ)-differentially private using sensitivity analysis and a Gaussian mechanism, while total privacy also includes model-communication costs.The scoring sensitivity is at most 2(λ(|D_i| + |D−|))^-1, and post-processing preserves the scoring mechanism’s privacy guarantee.

4. Experiments

Experiments evaluate FedAUX across image and text classification, varying architectures, client heterogeneity, participation, auxiliary data, and privacy parameters. FedAUX consistently improves federated performance, especially under non-iid data, while retaining benefits from auxiliary-data pre-training and weighted distillation.

  • Experimental setup: Experiments cover image and text classification with convolutional, MobileNet, ShuffleNet, and TinyBERT models, using multiple local and auxiliary datasets.Client populations reach n = 100, with heterogeneity controlled through Dirichlet parameter α.
  • Image benchmarks: At α = 0.01, FEDAUX raises ResNet8 accuracy from 18.2% to 78.1% over FEDDF with 80 clients after 100 communication rounds.Both unsupervised pre-training and weighted ensemble distillation contribute independently to the improvement.
  • Image benchmarks: FEDAUX outperforms FEDDF across all tested heterogeneity levels and matches or exceeds equally weighted distillation in the mixed-architecture setting.The mixed setting distributes clients across ResNet8, MobileNetv2, and Shufflenet models.
  • Participation and heterogeneity: At α = 0.01, FEDAUX outperforms all other methods by up to 29%, whereas at α = 100.0 its maximum gap from FEDAVG+P and FEDPROX+P is 0.8%.Higher client participation rates benefit FEDAUX, and pre-trained initialization improves methods relative to random initialization.
  • NLP benchmarks: In NLP experiments at α = 0.01, FEDAUX outperforms FEDDF+P and FEDAVG+P and reaches 95% of final accuracy after one communication round on both datasets.At α = 1.0, FEDAUX is mostly comparable to pre-trained FEDAVG and FEDDF, with a maximum test-accuracy gap of 1.1%.
  • One-shot learning and privacy: FEDAUX outperforms baseline methods in one-shot learning at every tested heterogeneity level, while privacy-preserving scoring can approach non-private performance.With λ = 0.01, ε = 0.04, and δ = 10^-6, accuracy decreases from 74.6% to 70.8%.

5. Conclusion

FedAUX leverages unlabeled auxiliary data through unsupervised pre-training and weighted ensemble distillation, substantially improving federated learning performance under the paper’s assumptions.

  • FedAUX combines unsupervised pre-training with weighted ensemble distillation to derive greater utility from unlabeled auxiliary data.The approach uses auxiliary data for both model initialization and client-prediction aggregation.
  • These results show that unlabeled auxiliary data can drastically improve federated learning performance.
  • The findings question direct comparisons between Federated Distillation and parameter-averaging methods because the approaches make different data-availability assumptions.

- SUPPLEMENTARY MATERIALS -

The supplementary material situates FedAUX within federated distillation, describes its preparation and training procedures, and details the heterogeneous-data evaluation setup.

  • Ensemble Distillation in Federated Learning: Federated Distillation aggregates client knowledge independently of model architecture, allowing clients to train structurally different models.This flexibility is especially relevant in hardware-constrained settings.
  • Ensemble Distillation in Federated Learning: Unlike parameter averaging, federated distillation communicates clients’ soft-label predictions on distillation data rather than raw model parameters.The cited passage notes that this changes communication scaling, while the proposed approach communicates full models.
  • Weighted Ensembles: FedAUX weights client predictions with instance-dependent certainty scores rather than static distillation weights.The related-work discussion connects this design to weighted ensembles and data-dependent prediction weighting.
  • Data Generation: Client data heterogeneity is controlled with a Dirichlet splitting strategy whose concentration parameter α adjusts the degree of non-iidness.The procedure assigns clients non-overlapping class-specific data points after repeated row and column normalization.
  • Preparation Phase: FedAUX uses self-supervised pre-training on server-held unlabeled auxiliary data to initialize feature extractors for each model prototype.The pre-trained extractors are communicated to clients and initialize part of the local classifiers.

D. Qualitative Comparison with Baseline Methods

The qualitative comparison characterizes FedAUX’s computational, communication, privacy, and data requirements relative to FEDAVG and FEDDF.

  • Computational Complexity: FedAUX adds a one-time strongly convex ERM solve for clients, while the server performs self-supervised pre-training and certainty-score computation.The client overhead is described as negligible relative to repeated local deep-network training, and server resources are assumed stronger.
  • Communication: FedAUX communicates scoring vectors, a feature extractor, and negative data in addition to its model communication.The scoring-vector overhead is described as negligible, while downstream communication increases by less than a factor of (T + 1)/T.
  • Privacy: The scoring-model communication is made (ε, δ)-differentially private, with reported experiments using ε = 0.1 and δ = 10^-5.The passage states that performance improvements remain substantial under these privacy constraints.
  • Assumptions: FedAUX assumes the server has access to unlabeled auxiliary data, as do Federated Distillation methods including FEDDF.This is identified as an additional assumption relative to the comparison context.

E. Additional Results and Detailed Training Curves

Additional experiments report strong FedAUX performance under severe heterogeneity and across higher-class-count settings, while detailed curves expose convergence behavior at different α values.

  • Detailed Training Curves: At α ∈ {0.01, 0.04}, all methods exhibit convergence issues in the highly non-iid setting.The passage identifies this behavior as known in federated learning.
  • Detailed Training Curves: After one communication round, FedAUX exceeds the maximum performance achieved by every other method over the entire training course in the highly non-iid setting.This comparison is reported for the detailed training-curve experiments.
  • Additional Results: At every CIFAR-100 heterogeneity level, FedAUX outperforms FEDAVG and FEDDF.The comparison uses ResNet-8 after T = 100 communication rounds with STL-10 as auxiliary data.
  • Additional Results: At α = {0.01, 0.04}, FedAUX also surpasses FEDDF+P, whereas FEDDF+P slightly exceeds FedAUX at α = {0.16, 0.64}.The passage links this pattern to the behavior of the differentially private certainty-scoring method.
  • Auxiliary Data: The supplementary experiments vary auxiliary data by generating structured ImageNet subsets defined by top-level WordNet IDs.The subsets are resized to 32 × 32 × 3 and their sample counts are reported.

G. Details on the Implementation and Results of the NLP Benchmarks

The NLP experiments use TinyBERT with BookCorpus as public auxiliary data and evaluate federated learning on AG News and Multilingual Amazon Reviews.

  • TinyBERT was used for the NLP experiments and was pre-trained on BookCorpus.
  • BookCorpus was selected as the public dataset to follow the paper’s methodology.
  • AG News was used as a private topic-classification dataset.
  • English texts from the Multilingual Amazon Reviews Corpus were used to predict review star ratings.

H. Hyperparameter Evaluation

The section evaluates hyperparameters and reports that FedAUX matches the strongest methods in iid training while distinctly outperforming them under strong non-iid heterogeneity.

  • The hyperparameter analysis varies Adam learning rates from 1e−2 to 1e−5 for local training and distillation.
  • For FedPROX, the proximal parameter λprox is varied from 1e−2 to 1e−5.
  • FedAUX matches the performance of the best-performing methods in the iid setting with α = 100.0.
  • FedAUX distinctly outperforms all other methods in the non-iid setting with α = 0.01.

J. Proof of Theorem 1

The section presents a theorem and proof concerning the privacy mechanism, alongside benchmark and toy-example materials related to hyperparameter evaluation and distribution-weighted scoring.

  • J. Proof of Theorem 1: Theorem 2 assumes differentiable, 1-strongly convex R and differentiable l with |l′(z)| ≤ 1 for all z.
  • J. Proof of Theorem 1: The proof adapts a result from Chaudhuri et al. (2011).
  • H. Hyperparameter Evaluation: Figure 13 evaluates ResNet8 over 50 communication rounds with 20 clients and a participation rate of C = 40%.
  • H. Hyperparameter Evaluation: Table 8 reports the best-performing hyperparameter combinations for ResNet8 with 20 clients, 50 rounds, and C = 40%.
  • H. Hyperparameter Evaluation: Figure 14 uses a three-client Gaussian toy example to compare certainty-score weights with robust distribution weights.

K. Empirical Privacy Evaluation

The empirical privacy evaluation studies differentially private certainty scores and their effect on weighted ensemble distillation for non-iid CIFAR-10 clients using ImageNet Dogs auxiliary data.

  • The privacy evaluation compares highest-scoring auxiliary images with nearest neighbors in clients’ local data.
  • Weighted ensemble distillation improves training performance from 48.46% to 75.59% in the non-iid CIFAR-10 example.
  • The evaluation uses ImageNet Dogs as auxiliary data with clients holding non-iid CIFAR-10 subsets at α = 0.01.
  • The scoring models expose fuzzy, indirect information about local training data rather than direct single-data-point exposure.
  • In trustworthy-server settings, privacy loss through certainty scores is less critical because outside adversaries cannot inspect them.
Loading 2102.02514v1…