PaperlayerView exact revision on arXiv

Source-linked AI summary

FedBABU: Towards Enhanced Representation for Federated Image Classification

Jaehoon Oh, Sangmook Kim, Se-Young Yun

arXiv:2106.06042v3cs.LG

TL;DR

Federated learning has studied global-model improvement and personalization largely separately, leaving their relationship underexplored. FedBABU trains and aggregates only the body, then fine-tunes a randomly initialized head for personalization, yielding consistent improvements and efficient personalization, particularly under greater data heterogeneity.

  • Problem

    Research has largely pursued global-model improvement and personalized models separately, leaving limited study of how the two objectives relate.

  • Method

    FedBABU updates and aggregates only the body during federated training, then fine-tunes a randomly initialized head for client personalization.

  • Results

    FedBABU consistently improves performance and personalizes efficiently, particularly under greater data heterogeneity.

  • Takeaways & Limitations

    Freezing the head supports learning a shared representation that rapidly adapts to heterogeneous client data for personalization.

  • Takeaways & Limitations

    The study mainly considers balanced federated environments in which clients have equally sized datasets.

Abstract

from arXiv · show

Federated learning has evolved to improve a single global model under data heterogeneity (as a curse) or to develop multiple personalized models using data heterogeneity (as a blessing). However, little research has considered both directions simultaneously. In this paper, we first investigate the relationship between them by analyzing Federated Averaging at the client level and determine that a better federated global model performance does not constantly improve personalization. To elucidate the cause of this personalization performance degradation problem, we decompose the entire network into the body (extractor), which is related to universality, and the head (classifier), which is related to personalization. We then point out that this problem stems from training the head. Based on this observation, we propose a novel federated learning algorithm, coined FedBABU, which only updates the body of the model during federated training (i.e., the head is randomly initialized and never updated), and the head is fine-tuned for personalization during the evaluation process. Extensive experiments show consistent performance improvements and an efficient personalization of FedBABU. The code is available at https://github.com/jhoon-oh/FedBABU.

1 INTRODUCTION

Federated learning has expanded from optimizing one global model under data heterogeneity to learning multiple personalized models, but the introduction highlights a connection between these goals through client-level FedAvg analysis. It attributes personalization degradation to training the classifier head and motivates a fixed-random-head approach that updates only the body for efficient personalization.

  • Model decomposition: The model is decomposed into a body for representation learning and a head for learning the linear decision boundary.For MobileNet and ResNet, the head is the final linear layer and the body comprises all preceding layers; the head can be biased under class imbalance.
  • FedAvg: FedAvg updates and aggregates the entire network across clients, including their decision-boundary heads.Figure 1 contrasts this process with FedBABU and depicts heads as the last linear classifiers defining client decision boundaries.
  • FedBABU: The proposed method modifies FedAvg so clients share a fixed randomly initialized head while federated training updates only the body, enabling efficient personalization by changing the head.The introduction also reports that a fixed random head can perform comparably to a learned head in centralized learning.
  • Motivation: The paper connects global-model learning with personalized models by analyzing FedAvg at the client level and identifying shared-head training as harmful to personalization.The stated contribution is that training the head using shared server data negatively impacts personalization.

2 RELATED WORKS

Prior work studies federated learning for either a single global model or personalized local models, while network body–head decoupling has also been explored across several learning settings. The paper situates its approach within these lines of research and focuses on decoupling the entire network into body and head.

  • FL for a Single Global Model: FedAvg learns one global model without centrally storing clients’ raw data and reduces communication costs through local updates, but non-IID client data makes a globally optimal model difficult.
  • Personalized FL: Personalized federated learning develops client-specific local models while retaining federation’s benefits, using clustering, multi-task, transfer, regularization, and meta-learning methods.
  • Decoupling the Body and the Head for Personalized FL: Body–head network decoupling has been applied to long-tail recognition, noisy-label learning, and meta-learning, and has also been attempted for personalized federated learning.
  • Decoupling the Body and the Head for Personalized FL: The related-work discussion focuses on decoupling the entire network into a body and head, although other parameter-decoupling studies also exist.

3 PRELIMINARIES

This section formalizes the federated learning procedure, including client sampling, local updates, and server aggregation, and defines the body–head parameter decomposition used by FedBABU. It also specifies the experimental setup and client-level evaluation protocol for initial and personalized accuracy.

  • FL training procedure: Each communication round samples clients by fraction ratio f, locally trains their models for τ epochs, and aggregates the resulting parameters at the server.
  • FL training procedure: The study mainly considers a balanced environment in which all clients have equally sized datasets, making each client’s data-size weighting constant.
  • Experimental setup: Experiments use MobileNet on CIFAR100 with 100 clients, each holding 500 training samples and 100 test samples, under a heterogeneous shard-based distribution without overlap between shards.
  • Evaluation: Evaluation reports initial accuracy before personalization and personalized accuracy after fine-tuning each global model on every client’s training data, with table values given as mean±standard deviation across clients.

4 PERSONALIZATION OF A SINGLE GLOBAL MODEL

This section evaluates how a single FedAvg global model connects to personalized client models under decentralized test data. It identifies head training as a source of personalization degradation and motivates updating only the model body on the server.

  • FedAvg personalization: FedAvg personalization is evaluated using both initial accuracy without fine-tuning and personalized accuracy after five client-side fine-tuning epochs.The evaluation assumes test datasets remain scattered across clients rather than being collected by the server.
  • FedAvg personalization: More realistic federated settings—with smaller f, larger τ, and smaller s—produce lower initial accuracy.The comparison concerns FedAvg on CIFAR100 with 100 clients across different federated-learning settings.
  • Personalization degradation: The study constructs an experiment where initial accuracy increases while personalized accuracy decreases, showing that stronger global performance need not improve personalization.This experiment uses a server-held portion p of clients’ non-private data; the authors note that the random sampling of this data can violate federated-learning assumptions and serves only as motivation.
  • Personalization degradation: The authors hypothesize that server training injects client-irrelevant information about absent similar classes into the head, harming personalization.They test this hypothesis by zeroing the head’s learning rate and updating only the body on the server.
  • Personalization degradation: Training the head also negatively affects personalization for centrally trained models, according to the appendix experiment.This observation is reported alongside related work attributing head bias to its proximity to clients’ label distributions, including Zhuang et al. (2021), Luo et al. (2021), and Achituve et al. (2021).

5 FEDBABU: FEDERATED AVERAGING WITH BODY AGGREGATION AND BODY UPDATE

FedBABU decouples the model into a generalizable body and personalized head, training and aggregating only the body during federated learning. The randomly initialized fixed head provides shared criteria for representation learning, while personalization occurs by fine-tuning the head during evaluation.

  • Design motivation: The method decouples generalization and specialization by training the body for representation learning and the head for personalization.This design is motivated by prior work on long-tailed recognition, fine-tuning, and self-supervised learning, together with data-sharing results.
  • FedBABU algorithm: FedBABU trains and aggregates only body parameters, leaving a randomly initialized head fixed throughout federated training.The head’s learning rate is zeroed, so clients learn representations under the same fixed criterion and no head aggregation is needed.
  • Representation power: FedBABU improves “w/o head” initial accuracy over FedAvg in every tested CIFAR100 setting, especially under large data heterogeneity.The evaluation replaces each trained head with the nearest template; MobileNet is evaluated with 100 clients.
  • Ablation analysis: Freezing the head completely is critical, while freezing it also causes less damage to out-of-class representations during local updates.Appendix O compares different head-versus-body learning rates, and Appendix P provides class-wise analysis.
  • Personalization: For FedBABU, fine-tuning the head or full network personalizes better than fine-tuning only the body.Global models are fine-tuned for five epochs using each client’s training data; fine-tuning including the head is also computationally more practical than body-partial fine-tuning.

6 CONCLUSION

The study frames personalization as the problem of pre-training a better federated backbone for downstream personal tasks. It finds that methods improving the global model can reduce personalization ability, motivating network decoupling into body and head.

  • 6 CONCLUSION: Personalization depends on pre-training a better backbone federatedly for downstream personal tasks.The conclusion describes this as the central training problem for a useful federated global model.
  • 6 CONCLUSION: Existing global-model improvements, including data sharing and regularization, can reduce the ability to personalize.
  • 6 CONCLUSION: The proposed mitigation decouples the entire network into a body associated with the backbone and a head for personalization.The supplied conclusion passage introduces this decoupling as the response to personalization performance degradation, but the passage ends before describing the full procedure.

REPRODUCIBILITY STATEMENT

FedBABU’s implementation is publicly available, with shell files included for each algorithm to facilitate reproducibility.

  • The code is available at https://github.com/jhoon-oh/FedBABU, and shell files for each algorithm are included for reproducibility.

A IMPLEMENTATION DETAIL … F.1 ACCURACY CURVES IN THE CENTRALIZED SETTING

The paper details its experimental implementations, datasets, evaluation procedures, initialization analysis, and additional studies of personalization, momentum, client scale, and CIFAR10 accuracy curves. Across these analyses, representation learning, head initialization, and evaluation conditions materially shape personalized performance.

  • A IMPLEMENTATION DETAIL: The implementation uses public code, reports final-model results, and requires approximately 2 hours for FedBABU on CIFAR100 with MobileNet when f=1.0 and τ=1.Training time decreases as f decreases and τ increases.
  • A.1 ARCHITECTURES: The implementations use 3convNet for EMNIST, 4convNet for CIFAR10, and MobileNet or ResNet for CIFAR100.MobileNet uses depthwise and pointwise convolutions, while ResNet uses skip connections.
  • A.2 DATASETS: Experiments evaluate CIFAR and EMNIST with horizontal flipping and random cropping, with dataset composition summarized in Table 8.The datasets are identified as CIFAR and EMNIST.
  • A.4 EVALUATION OF FL ALGORITHMS: The evaluation procedure fine-tunes all algorithms and reports mean and standard deviation as initial and personalized accuracy.FedAvg and FedBABU share the full model after federated training, whereas FedPer and FedRep share the body and LG-FedAvg shares the head.
  • B ORTHOGONALITY OF RANDOM INITIALIZATION: When only the body is trained centrally, similar classifier-row initialization creates a convergence gap, motivating He (uniform) initialization as the default.The paper attributes this to the role of classifier-row directions in separating extractor outputs by class.
  • D EFFECT OF MOMENTUM DURING LOCAL UPDATES ON PERFORMANCE: Appropriate local-update momentum improves FedAvg and FedBABU, especially personalized accuracy, whereas extreme momentum m=0.99 causes FedAvg to completely lose the reported advantage.The comparison uses N=100, f=0.1, and τ=10, with results summarized in Table 10.
  • E EFFECT OF MASSIVENESS ON PERFORMANCE: With 500 clients and 100 training examples per client, FedAvg is slightly better than FedBABU in many CIFAR100 settings, while the passage reports an exceptional extreme setting.The extreme setting is s=10, f=0.1, and τ=10.
  • F RESULTS OF 4CONVNET ON CIFAR10: 4convNet experiments on CIFAR10 similarly emphasize representation learning for difficult tasks and show that orthogonality is necessary when only the body is trained.Figure 5 also examines how random initialization approximates orthogonal initialization with 4convNet.

F.2 EXPERIMENTAL RESULTS … G.1 ACCURACY CURVES IN THE CENTRALIZED SETTING

Across additional 4convNet and ResNet experiments, FedBABU retains its performance superiority, while ablations and centralized accuracy curves support fully training the body and the previously observed trends.

  • F.2 EXPERIMENTAL RESULTS: On CIFAR10 with 4convNet, FedBABU shows the same performance-comparison tendency and superiority as the corresponding main-paper experiments.Tables 12 and 13 report initial and personalized accuracy under various settings with 100 clients; Figure 5 examines initialization when only the body is trained.
  • F.3 ABLATION STUDY ACCORDING TO THE LOCAL UPDATE PARTS: Across all local-update-part ablations, FedBABU achieves the best performance, indicating that the body must be totally trained.The 4convNet CIFAR10 ablation evaluates initial and personalized accuracy while fine-tuning the entire network during personalization.
  • G.1 ACCURACY CURVES IN THE CENTRALIZED SETTING: In centralized CIFAR100 experiments, ResNet18 and ResNet50 accuracy curves show the same trends previously observed for update-part and initialization comparisons.Figures 6–9 cover both architectures and both comparison dimensions.

G.2 EXPERIMENTAL RESULTS

On CIFAR100 with ResNet architectures, FedBABU gains a larger performance advantage over FedAvg than with MobileNet. As model complexity increases from ResNet18 to ResNet50, FedAvg performance declines while FedBABU performance remains stable.

  • G.2 EXPERIMENTAL RESULTS: FedBABU’s performance gap over FedAvg increases with ResNet compared with MobileNet.This comparison is reported for ResNet18 and ResNet50 results.
  • G.2 EXPERIMENTAL RESULTS: As the network grows from ResNet18 to ResNet50, FedAvg performance decreases, whereas FedBABU performance does not.

H RESULTS OF RESNET10 ON DIRICHLET DISTRIBUTION-BASED CIFAR100

The section evaluates the proposed and existing algorithms on unbalanced, non-IID CIFAR100 data generated from Dirichlet distributions. It examines heterogeneity across β values and reports personalized accuracy under a realistic federated-learning setting.

  • Results on Dirichlet distribution-based CIFAR100: The evaluation covers unbalanced and non-IID settings derived from the Dirichlet distribution, where lower β indicates greater heterogeneity.Figures 10 and 11 show data distributions for β = 1.0 and β = 0.5, respectively.
  • Results on Dirichlet distribution-based CIFAR100: The experiments use a realistic federated-learning setting with f=0.1 and τ=10.These settings are specified for the Table 18 comparison.
  • Results on Dirichlet distribution-based CIFAR100: Personalized accuracy is compared on Dirichlet distribution-based non-IID CIFAR100 with 100 clients, using f=0.1 and τ=10.Table 18 identifies the used network as MobileNet.

I RESULTS OF 3CONVNET ON EMNIST … O FEDAVG WITH DIFFERENT LEARNING RATES

Across EMNIST and CIFAR100 experiments, FedBABU consistently improves personalization by preserving a shared classifier and strengthening the body representation. Additional analyses attribute this benefit to classifier alignment, server-side body updates, and reduced head learning rates, while FedProx+BABU further improves performance when all clients participate.

  • I RESULTS OF 3CONVNET ON EMNIST: FedBABU achieves the best personalized performance on EMNIST with a 3convNet under the realistic federated setting.The experiment uses 1488 clients, approximately 450 data points per client, fraction ratio 0.1, and 10 local epochs.
  • J RESULTS ON IN-DISTRIBUTION (ID) AND OUT-OF-DISTRIBUTION (OOD) CLASSES: FedAvg and FedBABU are compared on in-distribution and out-of-distribution class accuracy before and after personalization under varied settings.In-distribution classes occur in clients’ training data, whereas out-of-distribution classes do not; the experiment scatters an IID-like test set to all clients.
  • K SIMILARITY BETWEEN CLIENTS DURING FINE-TUNING: During personalization fine-tuning, clients retain similar extractors but develop different classifiers, showing that the classifier is closely related to personalization.Both FedAvg and FedBABU update the entire network during evaluation, and cosine similarity is measured across 20 fine-tuning epochs.
  • L FEDBABU WITH BODY UPDATE ON THE SERVER: With server-side body updates, FedBABU improves personalization over FedAvg while maintaining the initial accuracy.The body is partially updated on the server using a small portion p of clients’ nonprivate data, as in the corresponding experiment in Table 2.
  • M BODY AGGREGATION AND BODY UPDATE ON THE FEDPROX: When every client participates in each communication round, FedProx+BABU outperforms both FedAvg and FedBABU in personalized performance.This comparison uses CIFAR100 with 100 clients and fraction ratio f=1.0; FedProx uses µ=0.01.
  • N DISCUSSION ON THE EFFECTIVENESS OF FEDAVG: MOCHA [Smith et al., 2017], pFedMe [T Dinh et al., 2020], and Ditto [Li et al., 2021] span regularized personalized federated learning between local-only and FedAvg, depending on regularization.FedAvg+Fine-tuning also lies between these extremes, but is more FedAvg-based because it first develops a shared global model before local personalization.
  • N DISCUSSION ON THE EFFECTIVENESS OF FEDAVG: Cheng et al. (2021) theoretically analyzed FedAvg+Fine-tuning alongside local-only, FedAvg, Per-FedAvg [Fallah et al., 2020], and pFedMe [T Dinh et al., 2020], motivating further real-data study.The paper notes that existing real-dataset evidence is limited beyond test accuracies, despite theoretical bias-variance analysis in a stylized linear-regression setting.
  • O FEDAVG WITH DIFFERENT LEARNING RATES: FedAvg representation power increases as the head learning rate decreases, supporting a shared classifier that gives all clients the same class boundary.The experiment fixes the body initial learning rate αb=0.1, fraction ratio 0.1, and local epochs 10, and evaluates accuracy without a classifier to measure representation power.

P CLASS-WISE ANALYSIS DURING FEDERATED TRAINING … T QUALITATIVE COMPARISON BETWEEN FEDAVG AND FEDBABU

The merged sections analyze FedAvg and FedBABU under realistic heterogeneous federated settings, covering class-wise accuracy, classifier orthogonality, shard-based client distributions, longer training, and representation visualizations. They report that orthogonal head initialization is necessary for desirable FedBABU performance and provide experimental analyses of the remaining factors without specifying unseen numerical or visual outcomes.

  • P CLASS-WISE ANALYSIS DURING FEDERATED TRAINING: The class-wise analysis distinguishes in-class test examples from out-of-class examples according to whether their classes appear in a client’s training data.For example, a client trained on classes 0–9 treats classes 0–9 as in-class and classes 10–99 as out-of-class.
  • P CLASS-WISE ANALYSIS DURING FEDERATED TRAINING: FedBABU’s in-class aggregated-model accuracy is significantly low but increases dramatically during local updates under realistic federated settings.The analysis uses f=0.1, τ=10, and separates each client’s CIFAR-100 test data into in-class and out-of-class groups; Figure 13 reports the corresponding accuracy curves.
  • Q FEDBABU WITH THE NON-ORTHOGONAL CLASSIFIERS: Without orthogonal row vectors in the randomly initialized head, FedBABU cannot achieve desirable performance under federated learning settings.Table 24 compares FedBABU with and without the proposed orthogonal head across various settings using 100 clients and MobileNet on CIFAR100.
  • R DESCRIPTION OF DATA DISTRIBUTION ACCORDING TO THE SHARDS PER USER s: Shard-based client distributions are illustrated for s=10 and s=2 in a CIFAR10 example with 10 clients.When s=10, the dataset is divided into 100 same-class shards, each containing 500 samples.
  • S PERFORMANCE WITH LARGER TOTAL EPOCHS: The study fixes total training epochs at 320 and additionally evaluates 640 epochs because algorithms may not converge, including shard settings s=20 and s=5.Table 25 reports personalized accuracy comparisons for 640 total epochs with f=0.1, 100 clients, and MobileNet on CIFAR100.
  • T QUALITATIVE COMPARISON BETWEEN FEDAVG AND FEDBABU: Figure 15 visualizes representations learned by FedAvg and FedBABU with t-SNE on CIFAR10 and CIFAR100 under realistic federated settings.The visualizations use s=2 for CIFAR10 and s=10 for CIFAR100; CIFAR100 sparse classes are also converted into coarse classes for visualization.
  • T QUALITATIVE COMPARISON BETWEEN FEDAVG AND FEDBABU: Figure 16 provides an additional t-SNE comparison using only subclasses of the CIFAR100 vehicle super-class.The models use the same realistic settings, with s=2 for CIFAR10 and s=10 for CIFAR100.
Loading 2106.06042v3…