PaperlayerView exact revision on arXiv

Source-linked AI summary

FedTP: Federated Learning by Transformer Personalization

Hongxia Li, Zhongyi Cai, Jingya Wang, Jiangnan Tang, Weiping Ding, Chin-Teng Lin, Ye Shi

arXiv:2211.01572v2cs.LGcs.AI

TL;DR

Federated learning has not established how aggregation affects Transformer self-attention, despite its use in heterogeneous client settings. FedTP personalizes attention through a server hypernetwork and reports state-of-the-art performance in non-IID image and language tasks. The paper also reports improved adaptation to novel clients.

  • Problem

    The effects of federated aggregation on self-attention in Transformer-based federated learning had not been studied, leaving a gap in understanding its use with heterogeneous client data.

    Source: §I · pp. 1–2

  • Method

    FedTP uses a server hypernetwork to generate personalized self-attention projection matrices for clients while aggregating the other model parameters.

    Source: §I · p. 2, §V · p. 12

  • Results

    FedTP achieves state-of-the-art performance over personalized federated learning benchmarks in non-IID image and language tasks.

    Source: §I · p. 2, §V · p. 12

  • Takeaways & Limitations

    The authors report that FedTP adapts better to novel clients through its central hypernetwork.

    Source: §V · p. 12

Abstract

from arXiv · show

Federated learning is an emerging learning paradigm where multiple clients collaboratively train a machine learning model in a privacy-preserving manner. Personalized federated learning extends this paradigm to overcome heterogeneity across clients by learning personalized models. Recently, there have been some initial attempts to apply Transformers to federated learning. However, the impacts of federated learning algorithms on self-attention have not yet been studied. This paper investigates this relationship and reveals that federated averaging algorithms actually have a negative impact on self-attention where there is data heterogeneity. These impacts limit the capabilities of the Transformer model in federated learning settings. Based on this, we propose FedTP, a novel Transformer-based federated learning framework that learns personalized self-attention for each client while aggregating the other parameters among the clients. Instead of using a vanilla personalization mechanism that maintains personalized self-attention layers of each client locally, we develop a learn-to-personalize mechanism to further encourage the cooperation among clients and to increase the scablability and generalization of FedTP. Specifically, the learn-to-personalize is realized by learning a hypernetwork on the server that outputs the personalized projection matrices of self-attention layers to generate client-wise queries, keys and values. Furthermore, we present the generalization bound for FedTP with the learn-to-personalize mechanism. Notably, FedTP offers a convenient environment for performing a range of image and language tasks using the same federated network architecture - all of which benefit from Transformer personalization. Extensive experiments verify that FedTP with the learn-to-personalize mechanism yields state-of-the-art performance in non-IID scenarios. Our code is available online.

I. INTRODUCTION

Personalized federated learning addresses client heterogeneity by combining personalized models with joint training, but how federated aggregation affects Transformer self-attention had not been studied. FedTP personalizes self-attention and reports strong performance across non-IID image and language tasks.

I. INTRODUCTION

  • Data heterogeneity and system heterogeneity can make learning a single global model fail, motivating personalized models that still benefit from joint training.Client datasets can vary in both size and distribution –.Source: §I · p. 1
  • Whether federated algorithms affect Transformer self-attention was unstudied, despite self-attention's role in modeling global input interactions.The introduction notes that Transformers may be more robust to distribution shifts and data heterogeneity than CNNs.Source: §I · p. 1
  • In a five-client CIFAR-100 illustration, Local-T and Vanilla Personalized-T found task-relevant image regions, while FedAvg-T failed to produce meaningful attention maps.The paper uses Attention Rollout to visualize the maps and suggests aggregating attention across heterogeneous clients may ruin client-specific representations.Source: Figure 1 · p. 1, §I · p. 2
  • FedTP personalizes each client's self-attention while aggregating and sharing the other model parameters across clients.Its learn-to-personalize mechanism uses a server hypernetwork to generate client-specific self-attention projection matrices.Source: §I · p. 2
  • FedTP with learn-to-personalize achieves state-of-the-art performance across image and language tasks in non-IID settings.The paper reports experiments on three benchmark datasets under different non-IID settings.Source: §I · p. 2

A. Personalized Federated Learning

  • Prior personalized federated learning methods include global-model fine-tuning, proximal regularization, model interpolation, distillation, clustering, and locally personalized layers.The cited approaches include FedProx, pFedMe, Ditto, FedMD [30], FedGen, FedPer [34], and FedRep.Source: A. Personalized Federated Learning · pp. 2–3

B. Transformers

Transformers use self-attention to model global interactions, while the effects of federated aggregation on that mechanism remain a concern in heterogeneous settings. FedTP uses a server hypernetwork to personalize Transformer attention, with the model formulated over client-specific datasets and losses.

B. Transformers

  • Transformers model global interactions between inputs through self-attention, and prior work found them more suitable than CNNs for heterogeneous federated data.The paper notes that many federated Transformers still use FedAvg, which may limit their capabilities in non-IID scenarios.Source: B. Transformers · p. 3

C. Hypernetworks

  • Unlike prior federated hypernetworks that generate CNN weights or layer aggregations, FedTP generates Transformer attention projection matrices for client-wise queries, keys, and values.The cited comparison includes pFedHN, pFedLA, and Fed-RoD.Source: C. Hypernetworks · p. 3

III. FEDERATED LEARNING BY TRANSFORMER PERSONALIZATION

  • FedTP learns personalized self-attention for each client while sharing other parameters, using a server hypernetwork to generate each client's projection matrices.The framework uses ViT for image tasks and Transformer for language tasks; attention is computed from projected queries, keys, and values.Source: §III · p. 3, A. Problem Formulation · p. 3, Figure 2 · p. 4

A. Problem Formulation

  • The formulation assigns each of N clients a local dataset drawn from its own distribution and optimizes client-personalized models using a shared per-sample loss.The loss can be mean squared error or cross-entropy, selected for the task.Source: A. Problem Formulation · p. 3

B. Vanilla personalization for self-attention

Vanilla personalization keeps each client’s self-attention parameters local while aggregating the Transformer’s other parameters. This produces client-specific attention, but can limit cooperation, scalability, and generalization.

B. Vanilla personalization for self-attention

  • Each client retains its self-attention parameters locally, while the other Transformer parameters are aggregated across clients.The model separates personalized projection parameters from shared parameters during training.Source: B. Vanilla personalization for… · p. 4
  • The procedure locally trains each personalized model for several steps and aggregates the shared parameters on the server.The client-specific attention parameters remain local after training.Source: B. Vanilla personalization for… · p. 4
  • Because it ignores relationships between clients, vanilla personalization may yield suboptimal attention, scales linearly with client count, and requires retraining for novel clients.The passage also characterizes its generalization to novel clients as poor.Source: B. Vanilla personalization for… · p. 4

C. Learn-to-personalize for self-attention

FedTP uses a server-side hypernetwork to generate personalized self-attention projection parameters for each client, while sharing parameters across clients. The approach is designed to improve scalability and generalization to unseen clients compared with local-only personalization.

C. Learn-to-personalize for self-attention

  • FedTP learns a server-side hypernetwork to generate each client’s personalized self-attention projection matrices while sharing parameters across clients.The generated matrices provide client-specific self-attention while retaining flexibility in the Transformer.Source: C. Learn-to-personalize for… · p. 4
  • Given a client embedding, the hypernetwork outputs self-attention projection parameters that are decomposed into query, key, and value projections.The hypernetwork uses fully connected layers, with its last layer unique to each Transformer block.Source: C. Learn-to-personalize for… · p. 4
  • The method updates the hypernetwork and client embeddings using gradients calculated through the chain rule after local training.The stated update rules use the projection-parameter changes after local epochs.Source: C. Learn-to-personalize for… · p. 5
  • When the number of clients N exceeds the hypernetwork hidden dimension Dh, learn-to-personalize uses fewer parameters than vanilla personalization.The hypernetwork parameter count is approximately Dh × Ds, compared with N × Ds projection parameters for vanilla personalization.Source: C. Learn-to-personalize for… · p. 5
  • The shared hypernetwork is presented as improving generalization to novel clients whose data was unseen during training.It generates client-specific attention using client-wise embedding vectors.Source: C. Learn-to-personalize for… · p. 5

D. Generalization Bound

The generalization analysis bounds empirical performance under bounded-parameter and Lipschitz assumptions. It relates the bound to hypothesis-class complexity, client and dataset sizes, and indicates that learn-to-personalize has a smaller hypothesis-class VC-dimension than vanilla personalization, especially with many clients.

D. Generalization Bound

  • The analysis assumes bounded client embeddings, hypernetwork weights, and other Transformer parameters, along with Lipschitz conditions on expected loss.The parameter bounds are expressed using radii Rh and Rt.Source: D. Generalization Bound · p. 5
  • Theorem 1 gives a high-probability generalization bound for the personalized hypothesis under those assumptions.The theorem specifies probability at least 1 − δ and uses the hypothesis class’s VC-dimension d.Source: D. Generalization Bound · p. 5
  • The bound depends on hypothesis-class complexity, client count, total dataset size, and Lipschitz constants.The analysis identifies these as factors affecting empirical-distribution performance.Source: D. Generalization Bound · p. 6
  • Learn-to-personalize has lower VC-dimension than vanilla personalization, especially with many clients, because one hypernetwork generates attention layers for all clients.The analysis says the resulting reduction in d decreases the bound’s M/dN term and gives FedTP better generalization, particularly as client numbers grow.Source: D. Generalization Bound · p. 6

IV. EXPERIMENTS

The experiments compare FedTP with baselines across multiple learning setups and examine model components, generalization, robustness, and attention maps.

IV. EXPERIMENTS

  • The experiments evaluate performance across different network backbones and personalized Transformer components.The section also compares FedTP with baseline methods in varied learning setups.Source: §IV · p. 6
  • The evaluation examines generalization to novel clients, compatibility with other methods, and attention-map visualizations.Source: §IV · p. 6
  • The ablation study varies label-distribution heterogeneity, noise-based feature imbalance, and model parameters.Source: §IV · p. 6

1) Benchmarks:

The experiments compare FedTP with federated and personalization baselines across image and language tasks under several non-IID settings. They use Transformer architectures for FedTP and specify client participation and training procedures.

1) Benchmarks:

  • FedTP is compared with FedAvg, FedProx, FedPer, pFedMe, FedBN, pFedHN, pFedGP, and FedRoD.The listed comparison methods include basic federated algorithms and advanced personalization algorithms.Source: 1) Benchmarks · p. 6

2) Non-IID Settings of Datasets:

  • CIFAR-10, CIFAR-100, and Shakespeare are evaluated using pathological or Dirichlet image splits and Shakespeare’s original 80/20 train-test split.The pathological split assigns two or ten classes per client for CIFAR-10 or CIFAR-100; CIFAR-100 Dirichlet partitions use coarse and fine labels.Source: 2) Non-IID Settings of Datasets · p. 6

3) Model Architectures:

  • FedTP uses a tiny ViT for image classification and a two-block Transformer for Shakespeare, while benchmark language methods use a two-layer LSTM.The image ViT has eight blocks, eight attention layers per block, eight heads, patch size 4, and hidden size 128.Source: 3) Model Architectures · p. 7

4) Implementation Details:

  • Image experiments use 50 or 100 clients with 10% or 5% participation, and train for 1,500 communication rounds; Shakespeare methods train for 300 rounds.Training uses SGD for five local epochs, learning rate 0.01, and batch size 64; pFedHN uses 5,000 global rounds for equivalent communication cost.Source: 4) Implementation Details · p. 7

B. Performance Evaluation

FedTP is evaluated against Transformer and other federated methods using test accuracy and communication-round curves. It outperforms baselines in most reported cases, with a reported average improvement over pFedHN on CIFAR-100.

1) Evaluation Protocol:

  • Evaluation accuracy is averaged across client test predictions, with each algorithm tested every five rounds during its final 200 communication rounds.The reported mean and standard deviation are computed over those evaluation steps.Source: 1) Evaluation Protocol · p. 7

3) Analysis of Network Backbone:

  • On CIFAR-100, FedTP improves 11.31% on average over pFedHN and outperforms baselines in most image and language cases.The reported results also show that changing CNN or LSTM backbones to Transformers can improve performance under data heterogeneity.Source: Table III · p. 8, 3) Analysis of Network Backbone · p. 8, Table IV
  • Test accuracy is 84.40±0.10 for FedTP, 27.56±0.65 for FedPer, 21.34±1.04 for FedAvg, and 20.48±1.09 for FedProx.These are the values reported for the test-accuracy row.Source: 3) Analysis of Network Backbone · p. 8

2) Performance Analysis:

  • FedTP’s CIFAR-100 Dirichlet accuracy curve is smoother and higher than those of the other Transformer methods.The figure compares test accuracy and convergence behavior across 50 clients.Source: 2) Performance Analysis · p. 7, Figure 4 · p. 8

4) Analysis of Different Personalized Parts:

Among the tested Transformer components, personalizing self-attention performs best. Attention-map comparisons and extension experiments further examine client-specific attention and combinations with local memory or classifier heads.

4) Analysis of Different Personalized Parts:

  • Personalizing self-attention achieves the best results among the tested components; normalization personalization beats MLP-layer and whole-encoder personalization.The comparison uses the same hypernetwork to generate parameters for each component.Source: 4) Analysis of Different Personalized… · p. 8, Table V · p. 8

5) Visualization of Attention Maps:

6) Extension of FedTP:

  • Local memory can improve performance, especially when each class has few samples; a personalized classifier head may hurt under Dirichlet data, while retaining the global head mitigates this.FedTP is also described as compatible with methods that add personalized heads or local memory.Source: 6) Extension of FedTP · pp. 9–10, Figure 7 · p. 10

7) Generalization to Novel Clients:

FedTP’s hypernetwork-based personalization is reported to adapt well to novel clients and outperform local personalized self-attention. Across label-distribution settings, it is more robust than the compared methods.

7) Generalization to Novel Clients:

  • FedTP (Hypernetwork) is more robust and adapts well to novel clients after only a few fine-tuning epochs.The evaluation sampled 20% of CIFAR-100 clients as novel clients whose data was unseen during training.Source: 7) Generalization to Novel Clients · p. 10

8) Analysis of Hypernetworks:

  • FedTP leads Vanilla Personalized-T, supporting the role of hypernetworks in its performance.Vanilla Personalized-T restores each client’s projection parameters locally without hypernetworks.Source: 8) Analysis of Hypernetworks · p. 10
  • FedTP’s learned client embeddings cluster clients with the same coarse labels apart from those with different coarse labels.The authors interpret this as evidence that the hypernetwork encodes personalized information in client embeddings.Source: Figure 9 · p. 10, 8) Analysis of Hypernetworks · p. 10

1) Effect of Heterogeneity in Label Distribution:

  • Across the tested Dirichlet settings, FedTP outperforms the compared methods and remains robust as label-distribution heterogeneity varies.FedAvg-T performance decreases as heterogeneity rises, while some personalized methods may fall behind FedAvg-T as heterogeneity lessens; FedTP continues to work well.Source: Table VII · p. 11, 1) Effect of Heterogeneity in Label… · p. 11

2) Effect of Heterogeneity in Noise-based Feature Imbalance:

Experiments report that FedTP performs well under client-specific noise and is more stable than FedAvg-T as client sample rates vary. The conclusion reports strong non-IID performance alongside slower training than CNN-based federated models.

2) Effect of Heterogeneity in Noise-based Feature Imbalance:

  • FedTP leads the compared personalized federated methods across the tested client-specific noise levels.The experiments used 10 clients, a 50% participation rate, and σM values from 0 to 20.Source: 2) Effect of Heterogeneity in… · p. 12

3) Impact of the Self-Attention Block Number:

  • Increasing the number of self-attention blocks can improve the model’s ability to capture data heterogeneity; eight blocks were selected as the default for CIFAR-10 and CIFAR-100.The tested block counts were 1, 2, 4, 6, 8, and 10.Source: 3) Impact of the Self-Attention Block… · p. 12

4) Impact of Client Sample Rate:

V. CONCLUSIONS

  • FedTP outperforms state-of-the-art personalized federated learning methods in non-IID scenarios.The conclusion also reports adaptation to novel clients and better results than benchmarks across tested noise levels.Source: §V · p. 12
  • FedTP can be combined with methods including FedPer, FedRod, and KNN-Per, and the combined models achieved better experimental performance.The conclusion identifies compatibility with advanced personalized federated learning methods.Source: §V · p. 12
  • FedTP’s training speed is slower than CNN-based federated learning models, a limitation the authors identify as common to Transformer-based methods.More computation- and communication-efficient FedTP is future work.Source: §V · p. 12
Loading 2211.01572v2…