# Personalized Federated Multi-Task Learning over Wireless Fading Channels
## 3 Algorithm Description
### FedGradNorm Description  
_FedGradNorm_ adjusts gradient magnitudes to balance training rates between different tasks across clients. _FedGradNorm_ is distributed across the clients and the parameter server. \(\tilde{G}_{\tilde{\omega}}\) is used to have a common scale for the gradient sizes while the gradient norms are adjusted according to the relative inverse training rates \(r_{k}^{(i)}\). A higher value of \(r_{k}^{(i)}\) leads to a larger gradient magnitude for task \(i\), which encourages the task to train faster. Each client \(i\) sends its inverse training rate \(\bar{F}_{k}^{(i)}\) at time \(k\) to the PS, so that the PS can construct \(r_{k}^{(i)}\), \(i\in[N]\). Therefore, given the common scale of gradient magnitudes, and the relative inverse training rate, the desired gradient norm of task \(i\) at iteration \(k\) is calculated as \(\tilde{G}_{\tilde{\omega}}(k)\times\left[r_{k}^{(i)}\right]^{\gamma}\), where \(\gamma\) represents the strength of the restoring force which pulls tasks back to a common training rate, which can also be considered as a metric of task asymmetry across different tasks. In cases where tasks have different learning dynamics, a larger \(\gamma\) would be a better choice for a stronger balancing. In order to shift the gradient norms towards the desired norm, the loss weights \(p_{k}^{(i)}\) are updated by minimizing an auxiliary loss function \(F_{\text{grad}}\left(k;\{p_{k}^{(i)}\}_{i=1}^{N}\right)\) which is the sum of \(\ell_{2}\) distances between the actual gradient norm and the desired gradient norm across all tasks for each iteration \(k\), i.e.,  
\[F_{\text{grad}}\left(k;\{p_{k}^{(i)}\}_{i=1}^{N}\right) =\sum_{i=1}^{N}F_{\text{grad}}^{(i)}\left(k;p_{k}^{(i)}\right) \tag{11}\] \[=\sum_{i=1}^{N}\left\|p_{k}^{(i)}\|\
abla_{\tilde{\omega}_{k}^{(i )}}F_{k}^{(i)}\|-\tilde{G}_{\tilde{\omega}}(k)\times[r_{k}^{(i)}]^{\gamma}\right\| \tag{12}\]  
The auxiliary loss function \(F_{\text{grad}}\left(k;\{p_{k}^{(i)}\}_{i=1}^{N}\right)\) is constructed by the parameter server at each global iteration \(k\) by using \(\
abla_{\tilde{\omega}_{k}^{(i)}}F_{k}^{(i)}\), which is a subset of the whole gradient of the global shared network sent by client \(i\) at iteration \(k\) for the global aggregation. Next, the parameter server performs the differentiation of \(F_{\text{grad}}\left(k;\{p_{k}^{(i)}\}_{i=1}^{N}\right)\) with respect to each element of \(\{p^{(i)}\}_{i=1}^{N}\) so that \(\
abla_{p^{(i)}}F_{\text{grad}}\) is applied via gradient descent to update \(p^{(i)}\). The desired gradient norm terms, \(\tilde{G}_{\tilde{\omega}}(k)\times\left[r_{k}^{(i)}\right]^{\alpha}\), are treated as constant to prevent loss weights \(\{p^{(i)}\}_{i=1}^{N}\) from drifting towards zero while differentiating \(F_{\text{grad}}\left(k;\{p_{k}^{(i)}\}_{i=1}^{N}\right)\) with respect to each loss weight \(p_{k}^{(i)}\). The weights are updated as,  
\[p^{(i)}\gets p^{(i)}-\alpha\
abla_{p^{(i)}}F_{\text{grad}},\quad i\in[N] \tag{13}\]  
The updated \(\{p^{(i)}\}_{i=1}^{N}\) are normalized such that \(\sum_{i=1}^{N}p^{(i)}=N\). Finally, the parameter server obtains the global aggregated gradient \(g_{k}=\frac{1}{N}\sum_{i=1}^{N}p_{k}^{(i)}g_{k}^{(i)}\) to update the global shared network parameters \(\omega\) via \(\omega_{k+1}=\omega_{k}-\beta g_{k}\) and broadcasts the updated parameters to the clients for the next iteration. The overall _FedGradNorm_ algorithm is summarized in Algorithm 2. In _FedGradNorm_, \(\text{Update}(f,h)\) represents the generic notation for the update of the variable \(h\) by using the gradient of \(f\) function with respect to the variable \(h\). ```
1:Initialize \(\{p_{k}^{(i)}\}_{i=1}


```
Initialize \(\omega_{0}\), \(\{p_{0}^{(i)}\}_{i=1}^{N}\) for\(k\)=1 to\(K\)do  The parameter server sends the current global shared network parameters \(\omega_{k}\) to the clients.