Quantized Analog Beamforming Enabled Multi-task Federated Learning Over-the-air

Let $\mathcal{D}_{n,l}$, $\forall n\in[N],l\in[L]$, denote the local dataset of the $l$-th device in the $n$-th FL task and $D_{n,l}$ represents the number of training samples in dataset $\mathcal{D}_{n,l}$. Then, the local loss function of the $l$-th device with the $n$-th FL task is given by $f_{n,l}(\mathbf{v}_{n},\mathcal{D}_{n,l})=\frac{1}{D_{n,l}}\sum_{\mathbf{u}\in\mathcal{D}_{n,l}}\mathcal{L}(\mathbf{v}_{n},\mathbf{u})$, where $\mathbf{u}$ is the data sample selected from the local dataset and $\mathcal{L}(\mathbf{v}_{n},\mathbf{u})$ represents the sample-wise loss function quantifying the prediction error of model parameter $\mathbf{v}_{n}$ on data sample $\mathbf{u}$. The learning process aims to optimize the specific model parameter $\mathbf{v}_{n}\in\mathbb{R}^{d}$ for the $n$-th task by minimizing the global loss function defined as $F_{n}(\mathbf{v}_{n})=\sum_{l=1}^{L}\alpha_{n,l}f_{n,l}(\mathbf{v}_{n},\mathcal{D}_{n,l})$, where $\alpha_{n,l}\triangleq\frac{D_{n,l}}{\sum_{\ell=1}^{L}D_{n,\ell}}$ represents the aggregation weight for the $l$-th device. For simplicity, we assume that for all tasks, the models contain the same number of parameters, denoted by $d$. For models with different sizes, zero padding can be employed to ensure a uniform size, thereby applicable to general scenarios.

For model training, we employ the typical FedAvg algorithm to iteratively minimize $F_{n}(\mathbf{v}_{n})$. The main steps of the $t$-th round of the FedAvg algorithm are listed as follows.

• 1)

  Model broadcasting: The PS broadcasts the latest global model $\mathbf{v}_{n,t}$ to all devices with the $n$-th FL task.

• 2)

  Local computing: Based on the received global model $\mathbf{v}_{n,t}$, each device updates its local model via mini-batch stochastic gradient descent (SGD), i.e., $\mathbf{v}_{n,t,i+1}^{l}=\mathbf{v}_{n,t,i}^{l}-\eta_{n,t}\nabla f_{n,l}\left(\mathbf{v}_{n,t,i}^{l},\bm{\xi}_{n,t,i}^{l}\right)$, where $\mathbf{v}_{n,t,i}^{l}$ denote the local model of the $l$-th device obtained after the $i$-th local iteration, $\mathbf{v}_{n,t,0}^{l}$ is initialized as $\mathbf{v}_{n,t}$, $\eta_{n,t}$ denotes the learning rate, and $\bm{\xi}_{n,t,i}^{l}$ is a local mini-batch selected from $\mathcal{D}_{n,l}$. Each device performs $E$ local iterations in a specific round and updates its local model as $\mathbf{v}_{n,t,E}^{l}$.

• 3)

  Local updates uploading: After local computing, each device transfers its local updates $\mathbf{g}_{n,l,t}$ to the PS, where $\mathbf{g}_{n,l,t}\!\triangleq\!\frac{\mathbf{v}_{n,t,0}^{l}\!-\!\mathbf{v}_{n,t,E}^{l}}{\eta_{n,t}}\!=\!\sum_{i=0}^{E-1}\!\nabla f_{n,l}\!\left(\mathbf{v}_{n,t,i}^{l},\bm{\xi}_{n,t,i}^{l}\right)$.

• 4)

  Model aggregation: Upon receiving all the local updates, the PS calculates the global update as $\mathbf{g}_{n,t}=\sum_{l=1}^{L}\alpha_{n,l}\mathbf{g}_{n,l,t}$, and then updates the global model according to $\mathbf{v}_{n,t+1}=\mathbf{v}_{n,t}-\eta_{n,t}\mathbf{g}_{n,t}$.

The downlink model broadcasting is usually assumed error-free due to abundant resource at the PS [4]. For the uplink, we adopt the AirComp technique for simultaneous transmission and model aggregation. Moreover, to enable parallel model aggregation for multiple FL tasks, the PS exploits hybrid analog and digital beamforming with a typical sub-connected structure [12]. In particular, the PS is equipped with $N_{\text{RF}}$ RF chains and $N_{r}$ antennas. Each RF chain connects to an exclusive set of $M$ antennas through a dedicated phase shifter, where $M=\frac{N_{r}}{N_{\text{RF}}}$. We consider the scenario with $N_{\text{RF}}=N$, where each RF chain and its connected analog array is dedicated to serving a specific FL task.

Now, it is ready to introduce the considered multi-task AirComp. Before transmission, local updates are first normalized as $\mathbf{s}_{n,l,t}=\frac{\mathbf{g}_{n,l,t}-\mu_{n,l,t}\mathbf{1}}{v_{n,l,t}}$, where $\mu_{n,l,t}$ and $v_{n,l,t}$ respectively represent the mean and standard deviation of all entries in $\mathbf{g}_{n,l,t}$, and $\mathbf{1}$ is all one vector. Unlike traditional AirComp relying on perfect CSIT, we assume that no CSIT is available at the device, and consequently corresponding preprocessing, such as the typical channel inversion scheme in [13], is not performed. Hence, the received signal at the PS, $\mathbf{Y}_{t}\in\mathbb{C}^{N\times d}$, is given as follows:

where $\mathbf{h}_{i,l,t}\sim\mathcal{CN}\left(\mathbf{0}_{N_{r}\times 1},\mathbf{I}_{N_{r}}\right)$ denotes the channel between the PS and the $l$-th device in the $i$-th cluster, $p_{i,l,t}$ represents its transmit power, $\mathbf{W}_{t}=[\mathbf{w}_{1,t},\mathbf{w}_{2,t},\cdots,\mathbf{w}_{N,t}]^{H}\in\mathbb{C}^{N\times N}$ and $\mathbf{A}_{t}\in\mathbb{C}^{N\times N_{r}}$ represent the digital combiner and analog beamformer, respectively, and $\mathbf{N}_{t}\in\mathbb{C}^{N\times d}$ denotes the additive Gaussian noise matrix and its entries are independent and identically Gaussian distributed with zero mean and variance of $\sigma^{2}$. With the sub-connected structure, the analog beamformer can be expressed as $\mathbf{A}_{t}=\text{diag}\left\{\mathbf{a}_{1,t}^{H},\cdots,\mathbf{a}_{N,t}^{H}\right\}$, where $\mathbf{a}_{i,t}$ denotes the analog beamforming vector of the $i$-th sub-array satisfying $\left|[\mathbf{a}_{i,t}]_{j}\right|^{2}=\frac{1}{M}$, $\forall j\in[M]$. Since the $n$-th RF chain is dedicated to serving the $n$-th FL task, we exploit the processed signal at the $n$-th RF chain as an estimate of the desired weighted aggregation, $\sum_{l=1}^{L}\alpha_{n,l}v_{n,l,t}\mathbf{s}_{n,l,t}$. It follows

The global update for the $n$-th task is then calculated as

where the statistics of the local updates are transmitted to PS in advance with ignorable communication overhead [14]. By comparing (3) with the ideal local updates, we express the error of global update estimation, $\mathbf{e}_{n,t}\triangleq\hat{\mathbf{g}}_{n,t}-\mathbf{g}_{n,t}$, as

Note that the error in (II-B) comprises of two main components. First, there is distortion in the aggregation weights of the desired signal. The second component consists mainly of interference from other tasks and the additive Gaussian noise. To tackle with the inter-task interference, existing works [8, 9] rely on sufficient spatial degrees of freedom (DoF) to apply ZF receivers for complete interference elimination, which requires a large number of RF chains, i.e., $N_{\text{RF}}>K$. However, in the scenario considered in this paper, the number of RF chains is limited, making complete interference elimination unattainable. In the subsequent section, we introduce how analog beamforming design is utilized to statistically eliminate interference and support the multi-task FL over-the-air.

Favorable propagation is a unique characteristic of large-scale MIMO systems. When the number of antennas, $N_{r}\to\infty$, the channels of different devices are asymptotically orthogonal, i.e., $\mathbf{h}_{i,l,t}^{H}\mathbf{h}_{i^{\prime},l^{\prime},t}\to 0$, $\forall(i,l)\neq(i^{\prime},l^{\prime})$. Inspired by the nature of favorable propagation, the inter-task interference can be statistically eliminated with fully-digital beamforming at the receiver by applying a linear projection [16], referred to as random orthogonalization (RO). To be specific, we set the fully-digital beamformer for the $n$-th task as $\mathbf{f}_{n,t}=\sum_{l=1}^{L}\mathbf{h}_{n,l,t}\in\mathbb{C}^{N_{r}}$ and the received signal in (2) becomes

where $\overset{\text{a.s}}{\to}$ represents “almost surely converge to” and it is due to the favorable propagation property with $N_{r}\to\infty$, and the inter-device interference is asymptotically eliminated. Also, from the statistical perspective, we have

Hence, it is straightforward to verify that the inter-task interference is statistically eliminated. The receive beamforming acts as a filter, which filters out interference components orthogonal to it. However, achieving better interference elimination requires a substantial deployment of extra RF chains, significantly increasing hardware cost. Alternatively, large-scale antenna arrays enable the direct use of ZF receivers to obtain individual signals from different devices without interference. The RO approach may not necessarily yield better performance compared to the ZF receiver.

The core idea of (III-A) is to include the target channel components, i.e., $\mathbf{h}_{n,l,t}$, $\forall l$, in the $n$-th digital beamformer, thus achieving statistical interference elimination via the favorable propagation property. However, the included amplitude components of $\mathbf{h}_{n,l,t}$, $\forall l$, do not exert a significant impact. Hence, we posit that utilizing analog phase shifters alone can still achieve this statistical interference elimination. In particular, considering that the $n$-th sub-array only serves the $n$-th FL task, we set its analog beamforming as

where $\mathbf{h}_{n,l,t,n}\in\mathbb{C}^{M}$ denotes the channel between the $n$-th sub-array and the $l$-th device in the $n$-th cluster in round $t$ and $\mathbf{h}_{n,l,t}=\left[\mathbf{h}_{n,l,t,1}^{H},\cdots,\mathbf{h}_{n,l,t,N}^{H}\right]^{H}$, $\forall n\in[N],l\in[L]$. In (7), we only incorporate the phases of channels related to devices in the $n$-th cluster, aiming at filtering out interference from other clusters. With the analog beamforming in (7), the effective channel for the $l$-th device in the $n$-th cluster at the $t$-th round is given as $\bar{\mathbf{h}}_{n,l,t}=\mathbf{A}_{t}\mathbf{h}_{n,l,t}$, whose statistics are derived in the following lemma.

Building upon the analytical results in Lemma 1, the statistical interference elimination can be achieved only with analog beamforming. To be concrete, the digital combiner $\mathbf{W}_{t}$ does not cope with the inter-task interference. We set $\mathbf{W}_{t}=\zeta\mathbf{I}_{N}$, where $\zeta>0$ is a scaling factor. Then, it follows

which no longer includes the inter-task interference terms.

To further demonstrate the statistical interference elimination method, we characterize the scaling law of the average power of the interference with respect to the number of antennas, $N_{r}$, in the following theorem. The average power of the interference for the $n$-th task is given by

Due to hardware limitations, phase shifts are usually limited to a finite number of discrete values. We further extend the proposed analog beamforming design to the scenarios with discrete phase control. Specifically, we denote the set of discrete phase shifts by $\mathcal{A}\triangleq\left\{0,\frac{2\pi}{2^{b}},\cdots,\frac{(2^{b}-1)2\pi}{2^{b}}\right\}$, where $b$ is the number of quantization bits. Then, we rewrite (7) as

where $\bm{\phi}$ is an $M$-dimensional vector with elements selected from $\mathcal{A}$. Comparing with the perfect case in (7), the configured phase is disturbed by a quantization error, which is modelled as uniformly distributed RVs following $\mathcal{U}\left(-2^{-b}\pi,2^{-b}\pi\right)$ [17]. Then, the discrete phase control in (12) is rewritten as

where $\bm{\psi}_{n,t}$ denotes the quantization error and $[\bm{\psi}_{n,t}]_{m}\sim\mathcal{U}\left(-2^{-b}\pi,2^{-b}\pi\right)$, $\forall m$. Accordingly, we denote the effective channel for the $l$-th device in the $n$-th cluster at the $t$-th round as $\tilde{\mathbf{h}}_{n,l,t}=\tilde{\mathbf{A}}_{t}\mathbf{h}_{n,l,t}$, where $\tilde{\mathbf{A}}_{t}=\text{blkdiag}\left\{\tilde{\mathbf{a}}_{1,t}^{H},\cdots,\tilde{\mathbf{a}}_{N,t}^{H}\right\}$.

Next, regarding the analysis of the average power of the interference, we derive the following lemma.

As observed in Lemma 3, the performance gap between the continuous and discrete phase control methods vanishes with increasing $b$. Moreover, $P_{\text{I},n}^{\text{D}}$ also decreases by $\frac{1}{N_{r}}$ as $N_{r}$ increases and we conclude the following corollary.

The proposed analog beamforming method, designed to prioritize favorable propagation over precise signal processing, demonstrates resilience to the impacts of non-ideal hardware. Hence, it suggests us to use a large number of low-precision phase shifters to achieve a satisfactory performance of interference elimination.