Robust Segmented Analog Broadcast Design to Accelerate Wireless Federated Learning

TL;DR

The paper tackles the downlink bottleneck in wireless federated learning under imperfect CSI by proposing SegAB, a segmented analog downlink broadcast that transmits multiple model segments in parallel to K devices using beamforming at a multi-antenna BS. It derives an upper bound on the expected model optimality gap that is separable across rounds, enabling online per-round robust beamforming via an epigraph formulation and an ADMM-based algorithm with monotone convergence. The key contributions are (i) a robust SegAB design that reduces transmission latency, (ii) a per-round online optimization framework that accounts for CSI errors, and (iii) a low-complexity, convergent beamforming algorithm with strong empirical performance, demonstrating significant gains over full-model broadcasting across multiple datasets and system settings.

Abstract

We consider downlink broadcast design for federated learning (FL) in a wireless network with imperfect channel state information (CSI). Aiming to reduce transmission latency, we propose a segmented analog broadcast (SegAB) scheme, where the parameter server, hosted by a multi-antenna base station, partitions the global model parameter vector into segments and transmits multiple parameters from these segments simultaneously over a common downlink channel. We formulate the SegAB transmission and reception processes to characterize FL training convergence, capturing the effects of downlink beamforming and imperfect CSI. To maximize the FL training convergence rate, we establish an upper bound on the expected model optimality gap and show that it can be minimized separately over the training rounds in online optimization, without requiring knowledge of the future channel states. We solve the per-round problem to achieve robust downlink beamforming, by minimizing the worst-case objective via an epigraph representation and a feasibility subproblem that ensures monotone convergence. Simulation with standard classification tasks under typical wireless network setting shows that the proposed SegAB substantially outperforms conventional full-model per-parameter broadcast and other alternatives.

Figures (5)

• Figure 1: Downlink analog broadcast for wireless FL. Left: traditional full-model per-parameter broadcast; Right: proposed SegAB (an example of $S_t=3$ segments).
• Figure 2: Test accuracy vs. number of channel uses for training Model A using MNIST for $S_t=3,K=5$. Left: ($(N,\gamma) = (16,0.01)$). Middle: ($(N,\gamma) = (16,0.1)$). Right: ($(N,\gamma) = (32,0.1)$).
• Figure 3: Test accuracy vs. number of channel uses for training Model C using CIFAR-10 for $S_t=3,K=5$. Left: ($(N,\gamma) = (16,0.01)$). Middle: ($(N,\gamma) = (16,0.1)$). Right: ($(N,\gamma) = (32,0.1)$).
• Figure 4: Test accuracy vs. number of model segments $S_t$ for training Model B using Fashion MNIST for $(N,\gamma) = (16,0.1)$ and $K=5$.
• Figure 5: Test accuracy vs. normalized error bound $\gamma$ for training Model B using Fashion MNIST for $(N,S_t) = (16,5)$ and $K=5$.

Theorems & Definitions (1)

• Proposition 1