Personalized Federated Learning with Attention-based Client Selection

Standard federated learning (FL) aims to train a global model by aggregating local models contributed by multiple clients without the need for their raw data to be centralized. In FL, there are $n$ clients communicating with a server to solve the problem: $\min_{w}\frac{1}{n}\sum_{i=1}^{n}\mathbb{F}_{i}(w)$ to find a global model $w$. The function $\mathbb{F}_{i}$ denotes the expected loss over the data distribution of the client $i$. One of the first FL algorithms is FedAvg McMahan et al. (2017), which uses local stochastic gradient descent (SGD) updates and builds a global model from a subset of clients.

However, the performance of the global model tends to degrade when faced with clients whose data distributions differ significantly from the global training data distribution. On the other hand, local models trained on the respective data distributions match the distributions at inference time, but their generalization capabilities are limited due to the data scarcity issue Mansour et al. (2020). Personalized federated learning (PFL) can be viewed as a middle ground between pure local models and global models, as it combines the generalization properties of the global model with the distribution matching characteristic of the local model Li et al. (2021); Fallah et al. (2020); T Dinh et al. (2020); Deng et al. (2020).

In PFL, the goal is to train personalized models $w_{1},w_{2},...,w_{n}$ for individual clients while respecting data privacy and accommodating user-specific requirements. The problem can be framed as:

Instead of solving the traditional PFL problem, Huang et al. Huang et al. (2021) argued that the fundamental bottleneck in personalized cross-silo federated learning with non-IID data is that the misassumption of one global model can fit all clients. Therefore, they proposed a different approach by adding a regularized term with $\ell_{2}$-norm of model weights distances among clients and formulate the PFL problem as:

in which $W=[w_{1},...,w_{n}]$ is a matrix that contains clients’ local models as its columns and $\lambda$ is a regularization parameter. $\mathcal{F}(W)=\sum_{i=1}^{n}\mathbb{F}_{i}(w_{i})$ is known as “client-side personalization" in the context of PFL. $\mathcal{R}$ is a regularization term designed to help clients with similar data distributions collaboratively train their own models.

We propose a novel federated attention-based client selection algorithm (FedACS) (as illustrated in Algorithm1) to solve the optimization problem (1). In this paper, we mainly consider:

where $s_{ij}$ is the normalized score to measure data distribution similarity. It is worth mentioning that the formation of $\mathcal{R}(W)$ is easy to generate into other functions, like $\mathcal{R}(x)=1-e^{-x^{2}/\sigma}$ Huang et al. (2021). The above formulation implies that locally optimizing the objective function requires each client to collect the other clients’ model parameters for computation. However, researchers find that adversaries can infer data information in the training set based on the model parameters, and directly gathering model parameters will violate the principles of federated learning to protect data privacy Fredrikson et al. (2015); Shokri et al. (2017); Hitaj et al. (2017). To address this dilemma, we leverage the idea of incremental optimization Bertsekas et al. (2011). Specifically, we iteratively optimize $\mathcal{F}_{\lambda}(W)$ by alternatively optimizing $\mathcal{R}(W)$ and $\mathcal{F}(W)$ until convergence. Note that FedACS handles the refinement of $\mathcal{R}(W)$ on the server side. We introduce an intermediate model $u_{i}$ for each client, and these intermediate models form the model matrix $U=[u_{1},..,u_{n}]$ in the same way as $W$. In the $k$-th iteration, we first apply a gradient descent step to update $U$:

where $W^{k-1}$ is the collection of local models from the last round and $\alpha_{k}$ is the learning rate. Under the definition (2), the gradient of $\mathcal{R}(W)$’s $i$-th column is $\nabla\mathcal{R}(W)_{i}=w_{i}-\sum_{j}s_{ij}w_{j}$. In this paper, we use the cosine similarity of the model parameters instead of the data similarity to avoid the accessibility of the server to local data, i.e. $s_{ij}=\frac{\langle w_{i},w_{j}\rangle}{||w_{i}||||w_{j}||}$.

As mentioned above, the performance of a local model could be adversely affected if each client averages information from those with different data distributions. Thus, in practice, we also introduce $\delta$ to control the degree of collaboration. We adopt a dynamic strategy to define $\delta$ for each round. After gathering the parameters from clients at the $k$-th round, the server first computes the similarity of the model $s_{ij}^{k}$ to generate the similarity matrix $\mathcal{S}^{k}:(\mathcal{S}^{k})_{ij}=s_{ij}^{k}$, then $\delta$ is calculated by the $p$-quantile of all elements in $\mathcal{S}^{k}$. We can adjust the ratio $p$ to alter the degree of collaboration. Modifying $\delta$ can effectively mitigate the aggregation effect that arises from initial training rounds, where the model parameters fail to accurately represent the underlying data distribution accurately. After having $\delta^{k}$, we only consider the model similarity greater than $\delta^{k}$ and calculate $u_{i}^{k}$ as follows:

Since the similarity of the client model to itself is always 1 and greater than $\delta^{k}$. Like the attention mechanism in deep learning, clients with highly similar data distributions receive higher similarity scores, enabling the server to assign greater weights to their models during the combination Vaswani et al. (2017). By differentiating between models based on their relevance and informativeness, the server can selectively assign the most pertinent models to each client rather than treating all models equally. Furthermore, the server can collect weights from clients and utilize the attention mechanism to optimize $\mathcal{R}(W)$ effectively while still ensuring the preservation of data privacy for all the involved clients.

After optimizing $\mathcal{R}(W)$ on the server, FedACS then optimizes $\mathcal{F}(W)$ on clients’ devices by taking $U^{k}$ as initialization of $W^{k}$ and computes $w_{i}^{k}$ locally by:

The update of $w_{i}^{k}$ draws inspiration from local fine-tuning in FL that takes the global model as initialization and updates the local models with several gradient steps to fit local data distribution. Considering the scarcity of data, it is more efficient to use $u_{i}^{k}$ as an initial value instead of incorporating it into the objective function. . This research is driven by prior studies that have shown difficulties in achieving high training accuracy when the data is either limited or unevenly spread out, thus highlighting the significance of a carefully selected starting pointLi et al. (2019).

By iteratively optimizing $\mathcal{F}_{\lambda}(W)$ through alternating between $\mathcal{R}(W)$ and $\mathcal{F}(W)$ optimization, FedACS continues this process until a predefined maximum number of iterations, $K$, is reached.

We note that the usage of the regularization term in FedACS bears resemblance to FedAMPHuang et al. (2021), a method that also applies attentive message passing to facilitate similar clients to collaborate more. However, FedACS differs from FedAMP in the update steps. First, when computing the intermediate model matrix $U$, FedAMP combines all client models for each client, while FedACS only uses information from clients who share similar data distribution to further reduce the influence from unrelated clients. Meanwhile, if we set the threshold $\delta=\infty$ then the calculation of $U$ is the same as that in FedAMP. Second, for each round, FedAMP updates $W^{k}$ by $W^{k}=\operatorname*{arg\,min}_{W}\mathcal{F}(W)+\frac{\lambda}{2\alpha_{k}}||W-U^{k}||^{2}$, which means the $W^{k}$ needs to reach the local optimal. However, it is not easy to be satisfied when data is insufficient on local devices, and when the learning rate is small enough, the update function in FedAMP will be simplified to $W^{k}=U^{k}$ without using local data information. Additionally, the Euclidean distance version of FedAMP incorporates an $\ell_{2}$-term regularization, which adds the Euclidean distance of the models to the objective function without using weighted scores. However, this approach provides limited assistance to the server in effectively assigning similar clients’ models to each client, thereby underutilizing the potential of the attention mechanism. As discussed above, the limitation of data may adversely influence the performance of FedAMP, while our proposed FedACS will not suffer from this issue. The experiment results also demonstrate our conjectures.

In this subsection, we provide the convergence analysis of FedACS under suitable conditions without the special requirement of the convexity of the objective function. We start with the definition of $\epsilon$-approximate first-order stationary point.

Definition 1.    A point $W$ is considered an $\epsilon$-approximate First-Order Stationary Point (FOSP) for the problem (1) if it satisfies $||\nabla\mathcal{F}_{\lambda}(W)||\leq\epsilon.$

Consistent with the analysis performed in various incremental and stochastic optimization algorithms Bertsekas et al. (2011); Nemirovski et al. (2009); Huang et al. (2021), we introduce the following assumption and present the main result of the paper.

Assumption 1.    There exists a constant $B>0$ such that $max\{||Y||:Y\in\partial\mathcal{F}(W^{k})\}\leq B$ and $||\nabla\mathcal{R}(W_{k})||\leq B/\lambda$ hold for every $k\geq 0$, where $\partial\mathcal{F}$ is the subdifferential of $\mathcal{F}$ and $||\cdot||$ is the Frobenius norm.

Remark.    Theorem 1 implies that for any $\epsilon>0$, FedACS needs at most $\mathcal{O}(\epsilon^{-4})$ iterations to find an $\epsilon$-approximate stationary point $\hat{W}$ of problem (1) such that $||\nabla\mathcal{F}_{\lambda}(\hat{W})||\leq\epsilon$. It also establishes the global convergence of FedACS to a stationary point of problem (1) when $\mathcal{F}_{\lambda}$ is smooth and non-convex. We are unable to provide proof due to page limitations.

We evaluate the performance of FedACS and compare it with the state-of-the-art PFL algorithms, including Ditto Li et al. (2021), PerAvg Fallah et al. (2020), pFedMe T Dinh et al. (2020), APFL Deng et al. (2020), and FedAMP Huang et al. (2021). The first three methods are local fine-tuning methods, using the global model to help fine-tune local models. FedAMP is a method that focuses on adaptively facilitating pairwise collaborations among clients. To ensure the completeness of our experiments, we have also included the results obtained using the standard FL method FedAvg McMahan et al. (2017). The performance of all methods is evaluated by the mean test accuracy across all clients. We perform five experiments for each dataset and partition configuration, wherein we record both the mean and variance of the test accuracy. All methods are implemented in PyTorch 1.8 running on NVIDIA 3090.

We use two public benchmark data sets, FMNIST (FashionMNIST) Xiao et al. (2017) and CIFAR10 Krizhevsky et al. (2009). First, we simulate the non-IID and data-sufficient settings by employing two classical data split methods to distribute the data among 100 clients. These methods include: 1) a pathological non-IID data setting, where the dataset is partitioned in a non-IID manner, assigning two classes of samples to each client; and 2) a practical non-IID data setting, where the dataset is partitioned in a non-IID manner following a Dirichlet 0.5 distribution McMahan et al. (2017). To simulate scenarios where data insufficiency is encountered during training for each dataset, we adopt two distinct data settings: 1) a pathological non-IID data setting that distributes data into 100 clients following Dirichlet 0.5 distribution, and each user only keeps 50 pieces of data samples for training to simulate the scenario where clients lack sufficient training data: 2) a practical non-IID data setting that distributes data into 500 clients following Dirichlet 0.5 distribution to simulate the scenario of a large number of clients and only some clients participate in each round.

Figure 1(a) presents the mean accuracy and standard deviation per method and sufficient data setting. The two histograms on the left display the performance of PFL methods under pathological settings, while the right-side histograms illustrate the results when the data is partitioned according to the Dirichlet 0.5 distribution. It can be observed that FedACS performs comparably to most PFL methods, but with lower standard deviations. This similarity in performance may be attributed to the availability of sufficient data, allowing the global model to possess adequate generalization ability, while its local fine-tuning enables better adaptation to local data distributions. The abundance of training data simplifies the classification task for each client, as evidenced by the strong performance of models trained solely on local datasets. However, non-IID data settings pose challenges for global federated learning methods Huang et al. (2021). FedAvg’s performance on Cifar10 declines noticeably because it introduces instability in the gradient-based optimization process. This instability arises from aggregating personalized models trained on non-IID data from different clients Zhang et al. (2021). APFL achieves the best performance in the pathological setting by adaptively learning the model and leveraging the relationship between local and global models, thus increasing the diversity of local models as learning progresses Deng et al. (2020). On the other hand, we can observe that pFedMe struggles to achieve good performance, as it focuses on solving the bilevel problem by using the aggregation of data from multiple clients to improve the global model rather than optimize local performance.

Figure 1(b) presents the mean accuracy and standard deviation per method and insufficient data setting. The two histograms on the left depict the mean accuracy with only 50 data samples per client, whereas those on the right showcase the results when the data is divided among 500 clients. It is evident that FedACS outperforms all other methods with smaller standard deviations in both the FMNIST and CIFAR10 datasets. The three local fine-tuning PFL methods face challenges across different datasets and settings, highlighting the fundamental bottleneck in PFL with non-IID data. By relying solely on a misleading global model, the significance of local data distribution in PFL is underestimated, and a single global model fails to capture the intricate pairwise collaboration relationships between clients. It is evident that PerAvg outperforms pFedMe in all settings. Both methods incorporate a “meta-model" during training, but PerAvg utilizes it as an initialization to optimize a one-step gradient update for its personalized model. On the other hand, pFedMe simultaneously pursues both personalized and global models. As discussed in Section 2.3, leveraging the meta-model for a single local update step may yield better results when local devices have limited data. The reason behind PerAvg’s suboptimal performance in scenarios where each client possesses only 50 data samples can be attributed to its requirement of three data batches for a single update step. Consequently, the training process becomes insufficient due to the limited amount of data available. As discussed in Section 2.3, FedAMP fails to leverage the potential of the attention mechanism, resulting in underutilization. Additionally, it requires a large amount of local data to achieve local optima in each round, which is often unrealistic in many federated learning scenarios.