FedMCSA: Personalized Federated Learning via Model Components Self-Attention

Considering that $f(U)=\frac{1}{N}\sum\limits_{i=1}^{N}{{f_{i}}}\left({{\Theta_{i}}}\right)$ and $\phi(U)=\frac{1}{{2N}}\sum\limits_{l=1}^{L}{\sum\limits_{i\neq j}^{N}d\left({{{\left\|{{\theta_{i,l}}-{\theta_{j,l}}}\right\|}^{2}}}\right)}$, we can rewrite the optimization problem in (2) to

A natural way to tackle the optimization problem in (3) is to adopt the framework of incremental-type optimization [35, 11], which can be realized by iteratively optimize $F(U)$ by alternatively optimizing $f(U)$ and $\phi(U)$ until convergence. Specifically, the general method first optimize ${\phi(U)}$ by applying a gradient descent step, and then optimize ${f(U)}$ by applying a proximal point step. In the $t$-th iteration, the update is formulated as

where ${V^{t}}$ denotes the prox-center, and ${\tau_{t}}>0$ is the step size of gradient descent. The general method has been proven in prior work  [11] to converge to the optimal solution when $F(U)$ is convex and to a stationary point when $F(U)$ is nonconvex.

The general method can be easy to deploy by gathering all clinets’ data together. In this subsection, we propose FedMCSA to optimize the personalized FL problem while protecting the data privacy of each client. Specifically, FedMCSA maintains a personalized cloud model for each client on a cloud server to implement the optimization step (4) of the general method and deploys the optimization step (5) privately in each client. The workflow of FedMCSA is illustrated in Figure 2.

Following the optimization steps of the general method, FedMCSA firstly aims to optimizes ${\phi(U)}$ and implements the optimization step in (4) by computing matrix ${V^{t}}$ on the cloud server. Let ${V^{t}}=[W_{1}^{t},...,W_{N}^{t}]$ and $W_{i}^{t}=[w_{i,1}^{t}\parallel w_{i,2}^{t}\parallel...\parallel w_{i,L}^{t}]$, where $W_{1}^{t},W_{2}^{t},...,W_{N}^{t}$ are the columns of ${V^{t}}$ and $w_{i,1}^{t},w_{i,2}^{t},...,w_{i,L}^{t}$ are the layers of $W_{i}^{t}$, the update of the $i$-th column and $l$-th layer $w_{i,l}^{t}$ of ${V^{t}}$ computed in (4) can be rewritten into a linear combination of the model parameter sets ${\theta_{1,l}^{t-1}},\ldots,{\theta_{N,l}^{t-1}}$ as follows

where $d^{\prime}\left(\left\|\theta_{i,l}^{t-1}-\theta_{j,l}^{t-1}\right\|^{2}\right)$ is the derivative of $d\left(\left\|\theta_{i,l}^{t-1}-\theta_{j,l}^{t-1}\right\|^{2}\right)$ and ${\psi_{i,l,1}^{t}},\ldots,{\psi_{i,l,N}^{t}}$ are the linear combination weights of the model parameter ${\theta_{1,l}^{t-1}},\ldots,{\theta_{N,l}^{t-1}}$, respectively. Since ${\psi_{i,l,1}^{t}}+\ldots+{\psi_{i,l,N}^{t}}=1$, $w_{i,l}^{t}$ is actually a convex combination [11] of the model parameter ${\theta_{1,l}^{t-1}},\ldots,{\theta_{N,l}^{t-1}}$.

We have argued for finding a linear combination of the model components to update the model component. The main question that follows is how to compute the effective weight coefficients in such a way that FedMCSA can seamlessly be incorporated into any existing FL paradigm. It is vital to determine how to calculate each model component’s weight coefficient as it has a crucial effect on the output of each model component and the performance of personalized FL.

A native method is to refer to HeurFedAMP to fix the weight coefficient for the model component that needs to be updated currently, set it as a hyperparameter, and then assign the remaining weights to other model components according to the similarity between the current model component and other model components. The setting of this hyperparameter needs to consider the clients’ data distribution, for example, it is set to $1/(M_{i}+1)$, where $M_{i}$ is the number of same distribution clients for client $i$. Unfortunately, in practice, it is often difficult to provide this knowledge about the clients’ data distribution. Moreover, in the actual model components update process, the interaction between different model components is dynamic, therefore fixing the weight coefficient limits the dynamic change in the weight coefficients, which artificially damages the performance of the personalized models.

To avoid the requirement for clients’ data distribution and inappropriate restrictions on the weight coefficients of model components, FedMCSA introduces self-attention into the update of model components, which performs $L$ self-attention processes in parallel to produce the updated personalized model for each client. FedMCSA makes no assumptions about the knowledge of underlying data distributions or client similarities to provide greater flexibility in personalization. Specifically, for each self-attention process, $N$ model components $[{\theta_{1,l}^{t-1}},\ldots,{\theta_{N,l}^{t-1}}]$ corresponding to $l$-th layer from clients are used as input, and the output is the updated $N$ model components $[{w_{1,l}^{t}},\ldots,{w_{N,l}^{t}}]$ for $l$-th layer. The $query$ vector, $key$ vector, and $value$ vector are denoted by $[{\theta_{1,l}^{t-1}},\ldots,{\theta_{N,l}^{t-1}}]$ for $l$-th layer in each model components self-attention process.

Meanwhile, FedMCSA adopts the simple but very efficient cosine similarity function $sim(query,key)=\sigma\cos(query,key)$ with a scale hyperparameter $\sigma$ as the metric function between different model components. Therefore, ${\psi_{i,l,1}^{t}},\ldots,{\psi_{i,l,N}^{t}}$ in (6) can be expressed as the following

For simplicity, we denote the operation of FedMCSA on the server as the function $\mathscr{F}_{M\!C\!S\!A}$, where $U^{t-1}$ and $V^{t}$ are the input and output of $\mathscr{F}_{M\!C\!S\!A}$, respectively. Therefore, (4) can be reformulated as the following

After optimizing $\phi(U)$ using the model components self-attention mechanism on the server, we further implement the optimization of $f(U)$ in the client. Let ${U^{t}}=[\Theta_{1}^{t},...,\Theta_{N}^{t}]$, where $\Theta_{1}^{t},...,\Theta_{N}^{t}$ are the columns of ${U^{t}}$, the update of the $i$-th column $\Theta_{i}^{t}$ of ${U^{t}}$ computed in (5) can be rewritten as the following

where $\lambda=\frac{\gamma}{{{\tau_{t}}}}$. We optimize $\phi(U)$ by (8) on the server and $f(U)$ by (9) in the client, which together constitute the complete personalized FL framework FedMCSA. In the continuous iterative process of personalized FL, the entire process can be considered as a continuous federated model components self-attention network, as shown in Figure 2.

The complete algorithm description of FedMCSA is presented in Algorithm 1. FedMCSA implements a client-server personalized FL framework to boost the performance of personalized FL dramatically, which alternately optimizes $\phi(U)$ and $f(U)$ through the model components self-attention mechanism on the server and the proximal gradient descent method in the client until a maximum number $T$ of iterations is reached.

To demonstrate the empirical superiority of FedMCSA, several comparisons between FedMCSA, HeurFedAMP, pFedMe, Per-FedAvg, Fedprox, and FedAvg are conducted. First of all, the same hyperparameters for all algorithms are used as a basic comparison. Considering that the performance of algorithms will behave differently when the hyperparameters change, we further perform a grid search of hyperparameters to find the highest performance of combination fine-tuning hyperparameters.

To further analyze the model components self-attention mechanism of FedMCSA, we conduct ablation studies to verify the effectiveness of the mechanism from two different perspectives. On the one hand, we construct a variant of FedMCSA by replacing the model components self-attention mechanism with the average aggregation method. On the other hand, we integrate this mechanism into the most basic FL algorithm, FedAvg, as well as the representative personalized FL with average aggregation, pFedMe. Considering that the performance of the personalized model pFedMe(PM) is always better than the global model pFedMe(GM) in the same setting, here we use pFedMe(PM) to represent the pFedMe. Note that in each group of experiments, the settings are identical except for the variables of interest.

Specifically, we conduct three variants of FedMCSA, FedAvg, and pFedMe as follows.

• 1.

  FedMCSA${}^{-M\!C\!S\!A}$: The variant model that replaces the model components self-attention mechanism with the average aggregation method.

• 2.

  FedAvg${}^{+M\!C\!S\!A}$: The variant model that replaces FedAvg’s aggregation method with the model components self-attention mechanism.

• 3.

  pFedMe${}^{+M\!C\!S\!A}$: The variant model that replaces pFedMe(PM)’s aggregation method with the model components self-attention mechanism.

The experimental results of ablation studies are shown in Table 2. Compared with FedMCSA, FedMCSA${}^{-M\!C\!S\!A}$ has a significant drop in accuracy on the four benchmark datasets whether using the MLR model or the DNN model. FedMCSA${}^{-M\!C\!S\!A}$ achieves relatively good performance on Mnist dataset, with 6.62% and 2.64% drop in accuracy compared to FedMCSA using the MLR and DNN models, respectively. This is because the Mnist dataset is relatively simple, and the basic average aggregation method can mine sufficient effective knowledge. When the dataset is relatively complex, e.g., the Synthetic dataset, the impact of the absence of the model components self-attention mechanism is significantly enlarged. It is noted that compared with FedMCSA, FedMCSA${}^{-M\!C\!S\!A}$ drops by 38.13% with the MLR model and 37.98% with the DNN model on Cifar10 dataset. The consistent accuracy drop across different datasets illustrates the effectiveness of the model components self-attention mechanism of FedMCSA, and the greater accuracy impact on more complex datasets further illustrates that the role of the model components self-attention mechanism is more pronounced on relatively complex datasets.

We also compare the performance of FedAvg, FedAvg${}^{+M\!C\!S\!A}$, pFedMe, pFedMe${}^{+M\!C\!S\!A}$ on the four benchmark datasets. From Table 2, we can observe that the model components self-attention mechanism of FedMCSA not only achieves the overall improvement on FedAvg, but also achieves the improvement on pFedMe. FedAvg${}^{+M\!C\!S\!A}$ has a significant accuracy improvement over FedAvg on the four benchmark datasets. Especially on Cifar10 dataset, compared with FedAvg, FedAvg${}^{+M\!C\!S\!A}$ has improved the accuracy by 37.32% and 38.77% under the MLR model and DNN model, respectively. Although pFedMe has achieved good performance on different datasets, pFedMe${}^{+M\!C\!S\!A}$ still achieves performance improvements under a variety of different settings. On Synthetic dataset, pFedMe${}^{+M\!C\!S\!A}$ improves the accuracy by 10.3% and 7.45% under the MLR model and DNN model, respectively.

The results of Table 2 demonstrate the effectiveness of the model components self-attention mechanism of FedMCSA from the consistent accuracy decrease of FedMCSA${}^{-M\!C\!S\!A}$ and the consistent accuracy improvement of FedAvg${}^{+M\!C\!S\!A}$ and pFedMe${}^{+M\!C\!S\!A}$. Furthermore, the performance improvement across different datasets, different network models, and different algorithms further illustrates the generality and usability of the mechanism.

To investigate the impact of the important hyperparameter $\sigma$ on the convergence of FedMCSA, we conduct diverse experiments on the Synthetic dataset using both MLR and DNN models.

As $\sigma$ is the scale of model components similarity in the FedMCSA, $\sigma$ is considered as a hyperparameter of FedMCSA. As illustrated in Figure 5, all the parameters except $\sigma$ are fixed during the experiments. For MLR, once $\sigma$ increases to 30, the upper limit of FedMCSA’s accuracy is reached under the current conditions, while the upper limit of accuracy for DNN is only realized when $\sigma=70$. We observe that FedMCSA requires a moderate value to compute the personalized model. When the value of $\sigma$ is too small, FedMCSA will converge slowly and reduce the accuracy. However, when the value of $\sigma$ exceeds a certain threshold, the accuracy of FedMCSA is not further improved, while an excessive value of $\sigma$ will cause the computational cost of the process to increase sharply. As a result, the value of $\sigma$ = 50 is selected as the default in the case of both MLR and DNN for the Synthetic dataset. Similarly, the value of $\sigma$ = 50 is selected as the default in the case of both MLR and DNN for the Mnist, FMnist, and Cifar10 datasets.