Personalizing Federated Learning with Over-the-Air Computations

TL;DR

This paper introduces a personalized federated learning framework using over-the-air computation to improve training efficiency and model robustness in wireless edge networks, addressing communication and data heterogeneity challenges.

Contribution

It proposes a novel combination of analog over-the-air computation with bi-level optimization for personalized federated learning, enhancing efficiency and robustness.

Findings

Improved training efficiency demonstrated through convergence analysis.

Enhanced model personalization and robustness validated by extensive experiments.

Effective handling of communication bottlenecks in wireless networks.

Abstract

Federated edge learning is a promising technology to deploy intelligence at the edge of wireless networks in a privacy-preserving manner. Under such a setting, multiple clients collaboratively train a global generic model under the coordination of an edge server. But the training efficiency is often throttled by challenges arising from limited communication and data heterogeneity. This paper presents a distributed training paradigm that employs analog over-the-air computation to address the communication bottleneck. Additionally, we leverage a bi-level optimization framework to personalize the federated learning model so as to cope with the data heterogeneity issue. As a result, it enhances the generalization and robustness of each client's local model. We elaborate on the model training procedure and its advantages over conventional frameworks. We provide a convergence analysis that theoretically demonstrates the training efficiency. We also conduct extensive experiments to validate the efficacy of the proposed framework.

Equations

$$\underset{{\mathbf{v}}_k}{\min }{f}_k\left({\mathbf{v}}_k;{\mathbf{w}}^{\ast }\right)=F_k\left({\mathbf{v}}_k\right)+\frac{\lambda }{2}{\Vert {\mathbf{v}}_k-{\mathbf{w}}^{\ast }\Vert }^2$$

$$\text{s.t.}{\mathbf{w}}^{\ast }\in \underset{\mathbf{w}}{\mathrm{argmin}}\frac{1}{K}\sum _{i=1}^K{F}_i\left(\mathbf{w}\right)$$

$$x_k(s)=⟨\mathbf{u}(s),\nabla F_k({\mathbf{w}}^t)⟩$$

$${\int }_0^{\tau }{u}_i^2(s)ds=1,i=1,2,...,d$$

$${\int }_0^{\tau }{u}_i(s)u_j(s)ds=0,i\ne j.$$

$$y(s)=\sum _{k=1}^K{h}_{k,t}{P}_k{x}_k(s)+\xi (s),$$

$${\mathbf{g}}^t=\frac{1}{K}\sum _{k=1}^K{h}_{k,t}\nabla F_k\left({\mathbf{w}}^t\right)+{\mathbf{\xi }}_t,$$

$${\mathbf{w}}^{t+1}\leftarrow {\mathbf{w}}^t-{\eta }_g{\mathbf{g}}^t,$$

$$0<{\eta }_g<\min \left\{\frac{2}{{\mu }_h(\mu +L)},\frac{2{\mu }_h\mu LK}{{\sigma }_h^2{\bar{L}}^2(1+2\delta )(\mu +L)}\right\},$$

$$E\left[||{\mathbf{v}}_k^{t+1}-{\mathbf{v}}_k^{\ast }|{|}^2\right]\le \left(1-\mu {\eta }_l\right)E\left[||{\mathbf{v}}_k^t-{\mathbf{v}}_k^{\ast }|{|}^2\right]+{\eta }_l^2{\lambda }^2{M}^2$$

$$+{\eta }_l^2{\lambda }^{2}E\left[||{\mathbf{w}}^t-{\mathbf{w}}^{\ast }|{|}^2\right]+2{\eta }_l^2{\lambda }^{2}M\sqrt{E\left[||{\mathbf{w}}^t-{\mathbf{w}}^{\ast }|{|}^2\right]}$$

$$+2{\eta }_l\lambda \sqrt{E\left[||{\mathbf{v}}_k^t-{\mathbf{v}}_k^{\ast }|{|}^2\right]E\left[||{\mathbf{w}}^t-{\mathbf{w}}^{\ast }|{|}^2\right]}.$$

Full text

**Index Terms— ** Federated learning, personalization, wireless edge network, over-the-air computation, robustness.

1 Introduction

With the increasing concerns on data privacy as well as the rapid growing capability of edge devices, deploying the federated learning (FL) [1] at the edge of wireless network, commonly coined as federated edge learning (FEEL), is attracting arising attentions [2, 3], where the computation tasks could be decoupled from the cloud to the edge of the network in a privacy-preserving paradigm.

However, in real-world implementations of the FEEL system, a typical training process of a generic global model requires hundreds of communication rounds among the massively distributed clients. The iterative gradient exchange would bring hefty communication overhead [1, 4]. Hence, for a digital communication based-FEEL system run over the resource-constrained network, the limited communication bandwidth would inevitably constrain the scalability, since every selected client in each round requires an assigned orthogonal sub-channel to perform the update [5, 6].

To combat the communication bottleneck, an array of recent studies [6, 7, 8, 9, 10, 11] suggest incorporate analog over-the-air (A-OTA) computations into the design of FEEL systems, exploiting the superposition property of the multi-access channels for fast and scalable model aggregations. The adoption of the A-OTA computations with FEEL, termed as A-OTA-FEEL, have been demonstrated to have high spectral efficiency, low access latency, enhanced privacy protection, and reduced communication costs [6, 12, 5], all benefiting from the automatic “one-shot” gradient aggregation for model update [13, 7]. Nevertheless, A-OTA computations inevitably introduce the random channel fading and interference into the aggregated gradients, leading to performance degradations such as the slower convergence and instability [5, 10]. Hence, robust training techniques could be adopted to enhance the performance with channel imperfections.

In addition to the inherent channel fading and interference, current approaches for A-OTA-FEEL system design have not addressed the existing discrepancies in both local data statistics and qualities (i.e., data heterogeneity and label noise) due to the diverse preferences, bias, and hardware capabilities of different clients [14, 11]. Such discrepancies in clients’ datasets can significantly degrade the FL performance. More crucially, these discrepancies would even make the single generic global model fail to achieve good generalization and robustness performance on diverse local data [15, 16, 17]. On the other hand, the future intelligent network is envisioned to be able to provide customized services to the clients [18, 19]. It is necessary to address the individuality of the clients in the design of the A-OTA-FEEL system with personalized intelligent services.

In view of the above challenges, we propose a personalized training framework in the context of the A-OTA-FEEL. The proposed framework provides personalized model training services while still enjoying the benefits of analog over-the-air computations, in which each client would maintain two models (i.e., generic and personalized models) at the local via two different global and local objectives. We also provide a convergence analysis of the proposed personalized A-OTA-FEEL framework. Both the theoretical and numerical results validate the gain from the personalization design.

2 System model

We consider a wireless system consisting of one edge server that is attached to an access point and $K$ clients, where the $i$-th device has a local dataset $\mathcal{D}_{i}$. In this system, communications between the clients and the server are taken place over the spectrum. Each client’s goal is to ($a$) learn a statistical model based on its own dataset and ($b$) exploit information from the dataset of other clients and, aided by the orchestration of the server, attain an improvement toward its locally learned model while preserving privacy. Such tasks can be achieved via a bi-level optimization based PFL framework. More precisely, every client $k$ aims to find a local model $\boldsymbol{v}_{k}\in\mathbb{R}^{d}$ that solves the following personal objective funciton

[TABLE]

in which $F_{i}(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}$ denotes the loss function of client $i$, $\boldsymbol{w}\in\mathbb{R}^{d}$ is a globally trained generic model, and $\lambda$ is a hyper-parameter that controls the level of personalization of the clients’ locally trained personal models. We use $\eta_{l}$ to denote the learning rate in the optimization of personal objective. Notably, a large value of $\lambda$ indicates that the clients’ local models $\{\boldsymbol{v}_{i}\}_{i=1}^{K}$ need to well align with the global model $\boldsymbol{w}^{*}$, promoting commonality across the local models. In contrast, a small $\lambda$ improves personalization. Moreover, benefiting from such a bi-level optimization design, the personalized local models $\{\boldsymbol{v}_{i}\}_{i=1}^{K}$ would have better generalization and robustness performance on the limited local data.

To solve the above optimization problem, the clients need to not just train their local models through (1), but more importantly, jointly minimize a global objective function as per (2). Due to privacy concerns, the clients will carry out the minimization problem (2) without sharing data in an FL manner. The following section presents a model training approach that capitalizes on the properties of analog transmissions for low-latency and high-privacy federated computing.

3 Model Training Procedure

This section details the PFL model training process based on over-the-air computing schemes. (See Fig. 1 for an overview.) More precisely, we employ A-OTA computations for fast (and highly scalable) gradient aggregation that significantly improves the training efficiency of the global model. The detailed training procedure is elaborated on below.

1) Local Model Training: Without loss of generality, we assume the system has progressed to the $t$-th round of global training, where the clients just received the global model parameters $\boldsymbol{w}^{t}$ from the edge server. 111Because of the high transmit power of the access point, we assume the global model can be successfully received by all the clients. Then, each client $k$ updates its personalized local model $\boldsymbol{v}^{t}_{k}$ by optimizing the local personal objective function $f_{k}(\boldsymbol{v}_{k};\boldsymbol{w}^{t})$. (For simplicity, we use $\boldsymbol{v}_{k}$ to denote personal model.) Each client $k$ also computes its local gradient $\nabla F_{k}(\boldsymbol{w}^{t})$ for global model update.

2) Analog Gradient Aggregation: We consider the clients adopt analog transmissions to upload their locally trained parameters. Specifically, once $\nabla F_{k}(\boldsymbol{w}^{t})$ is computed, client $k$ modulates it entry-by-entry onto the magnitudes of a common set of orthogonal baseband waveforms [5], forming the following analog signal

[TABLE]

where $\langle\cdot,\cdot\rangle$ denotes the inner product between two vectors and $\boldsymbol{u}(s)=({u}_{1}(s),...,{u}_{d}(s)),s\in[0,\tau]$ has its entries satisfying

[TABLE]

$\tau$ represents the total time of signal duration. Once the analog waveforms $\{x_{k}(s)\}_{k=1}^{K}$ are available, the clients transmit them concurrently to the access point. Owing to the superposition property of electromagnetic waves, the signal received at the radio front end of the access point can be expressed as:

[TABLE]

where $h_{k,t}$ is the channel fading experienced by client $k$, $P_{k}$ the corresponding transmit power, and $\xi(s)$ denotes the additive noise. In this work, we assume the channel fading is i.i.d. across clients, with mean $\mu_{h}$ and variance $\sigma_{h}^{2}$. Besides, the transmit power of each client is set to compensate for the large-scale path loss and we use $P$ to denote the average power for all clients. This received signal will be passed through a bank of match filters, with each branch tuning to $u_{i}(s),i=1,2,..,d$. On the output side, the server obtains:

[TABLE]

in which $\boldsymbol{\xi}_{t}$ is a $d$-dimensional random vector with each entry being i.i.d. and follows a zero-mean Gaussian distribution with variance $\sigma^{2}$. It is noteworthy that the vector given in (7) is a distorted version of the globally aggregated gradient.

3) Global Generic Model Update: Using $\boldsymbol{g}^{t}$, the server updates the global model as follows:

[TABLE]

where $\eta_{g}$ is the learning rate for generic global model udpate. After this, the server broadcasts the $\boldsymbol{w}^{t+1}$ to all clients for the next round local computing. Such a process will be iterated through multiple rounds until the global model converges.

Notably, the bi-level optimization in the personal model mitigates impacts from the random channel fading and noise introduced by analog over-the-air computations to the globally aggregated gradient, thus improving the robustness of the analog over-the-air federated edge learning system. Consequently, personalization enhances both the generalization and robustness of the FL system in the presence of data heterogeneity and noisy model aggregation. 222This paper does not consider the architecture-based PFL methods in which each client maintains a personal model with unique architecture via techniques such as sparsification or model weight decoupling [15]. It would increase the cost of the synchronizations for signal transmission to achieve automatic signal aggregations in the context of A-OTA computations.

We summarize the proposed framework in Algorithm 1. It is worthwhile to highlight several advantages of the presented framework, including high scalability, low access latency, enhanced privacy, better generalization as well as robustness, brought together by analog over-the-air computations and personalized training. We would also like to address that we make no assumption about the generic model training, as well as the OTA communication, which indicates that the performance could be further enhanced by advanced federated optimization [4] and OTA techniques [8, 12].

4 Convergence analysis

This section provides the convergence analysis of our proposed framework from the perspective of both the global model and the local personalized model.

To facilitate the analysis, we assume that each client’s loss function is $\mu$-strongly convex and the local gradient $\nabla F_{k}(\boldsymbol{w}^{t})$ is Lipschitz continuous with constant $L_{k}>0$. We use $\bar{L}$ to denote the maximal constant among all clients and $L$ is the Lipschitz gradient constant of global objective. $\delta$ is the diameter of the compact convex parameter set that all model parameters lie in. We consider that if the global model converges, its convergence rate is denoted by $g\left(t\right)$, i.e., there exists $g\left(t\right)$ that $\underset{t\rightarrow\infty}{\lim}g\left(t\right)=0$ and $E\left[||\boldsymbol{w}^{t}-\boldsymbol{w}^{*}||^{2}\right]\leq g\left(t\right)$. In this work, we denote by $\boldsymbol{v}_{k}^{*}$ and $\boldsymbol{z}_{k}^{*}$ as $\boldsymbol{v}_{k}^{*}={\mathrm{arg}\min}_{\boldsymbol{v}}f_{k}\left(\boldsymbol{v};\boldsymbol{w}^{*}\right)$ and $\boldsymbol{z}_{k}^{*}={\arg\min}_{\boldsymbol{z}}F_{k}\left(\boldsymbol{z}\right)$, respectively. We assume the $l_{2}$ distance between the optimal local and global model is bounded, i.e., for any $k\in\left[K\right]$, $\left\|\boldsymbol{z}_{k}^{*}-\boldsymbol{w}^{*}\right\|\leq M$.

We now present the main theoretical finding of this paper. First of all, the following theorem provides the convergence rate of the global generic model.

Theorem 4.1**.**

Under the considered A-OTA FEEL system, let $r_{0}^{2}\triangleq\left\|\boldsymbol{w}^{0}-\boldsymbol{w}^{*}\right\|^{2}$ be the squared distance between the initial estimate $\boldsymbol{w}_{0}$ and $\boldsymbol{w}^{*}$. If the learning rate $\eta_{g}$ satisfies

[TABLE]

then the error of $\boldsymbol{w}^{t}$ can be bounded as:

[TABLE]

where $0<c\triangleq 1-\frac{2\eta_{g}\mu_{h}\mu L}{\mu+L}+\frac{\eta_{g}^{2}\sigma_{h}^{2}\bar{L}^{2}(1+2\delta)}{K}<1$.

Proof.

Please refer to [10] for a detailed proof. ∎

Next, we employ the following lemma to characterize the convergence rate of the local personalized models.

Lemma 4.2**.**

Under the considered system, let local learning rate satisfy condition (9), then the local model of client $k$ converges as:

[TABLE]

Proof.

Please refer to [16] for detailed proof. ∎

Aided by the above result, we obtain the convergence rate of the global model as the following.

Theorem 4.3**.**

Under the considered A-OTA FEEL system, if there exists a variable $A$ satisfying $\frac{g(t+1)}{g(t)}\geq 1-\frac{g(t)}{A}$, then, there is a constant $C<\infty$ such that for any client $k$, $E\left[\left\|\boldsymbol{v}_{k}^{t}-\boldsymbol{v}_{k}^{*}\right\|^{2}\right]\leq Cg(t)$ with a local learning rate given by $\eta_{l}=\frac{2g\left(t\right)}{A\mu}$.

Proof.

We omit the proof due to the space limit. ∎

To this end, we can see that via A-OTA computing, both the global and local personalized models attain linear convergence rates, while addressing the non-ideal gradient updates.

5 Numerical results

This section evaluates the performance of our proposed framework. Particularly, we examine the performance of the personalized local training in terms of generalization power and robustness compared to conventional settings and baselines. We also explore the robustness performance of the framework in the context of the noisy local data (i.e., part of the local training data are annotated with wrong labels).

5.1 Experiment setup

We evaluate our framework on image classification tasks on CIFAR-10/100 [20] with ResNet-18 and ResNet-34 [21], respectively. Both IID and non-IID data settings are considered, in which the non-IID data partitions are implemented with Dirichlet distribution and the identicalness of the distributions could be controlled by the parameter $\alpha$. Unless otherwise specified, we use $K=100$ for CIFAR-10, $K=50$ for CIFAR-100, and Rayleigh fading with average channel gain $\mu_{h}=1$. We select $\lambda$ from comparison experiments. The federated label noise setting is the same as the [14]333For all label noise settings, we use lower bound 0.5 for local label noise level. Details can be found in [14]..

5.2 Performance evaluation

We first compare the personalized models performance of our proposed framework with generic global model from conventional FL setup in Fig. 2 with IID local data partition, using the same configurations of A-OTA-FEEL system. The two sub-figures demonstrate the consistent outperformance of personalization training scheme. Specifically, increasing total number of clients in the system (i.e., a larger $K$) would improve the system performance for both two settings in A-OTA, and personalized training presents a more robust generalization with diverse local data quality.

To further demonstrate the outperformance of the proposed framework, we provide the detailed best test accuracy comparison in Tab. 1 on CIFAR-10/100 with non-IID data, compared with FedAvg [1], FedProx [4] and FedRep [22] with same OTA setup. In such context, our proposed personalzied training method achieves best test accuracies across all settings, which shows the superiority with respect to the generalization and robustness.

6 Conclusion

In this paper, we proposed a personalized A-OTA-FEEL framework that utilizes bi-level optimization and analog transmissions to address the data heterogeneity and communication efficiency challenges. Both the theoretical and empirical results were provided to demonstrate the effectiveness of the proposed framework. We highlighted the robustness performance of the PFL in edge learning. To the best of our knowledge, this is the first work that explores the PFL model in A-OTA FEEL systems. We envision that PFL could be a potential technique to provide customized services in future intelligent networks.