An Autoencoder-Based Constellation Design for AirComp in Wireless Federated Learning

TL;DR

This paper addresses the challenge of decoding the sum of local model updates in wireless federated learning when using digital modulation. It introduces an end-to-end autoencoder-based system that jointly optimizes transmitters and receivers to perform AirComp, effectively recovering the sum $\sum_i \mathbf{s}_{i,t+1}$ over fading channels with CSIT/CSIR. Through quantization-aware encoding, one-hot signaling, and a neural decoder, the method supports higher-order modulations and employs a coarse-grained categorization strategy to scale to many clients. Experimental results on CIFAR-10 show that AE_opt can achieve near-perfect communication performance at moderate-to-high SNRs and remains robust under non-IID data, indicating practical potential for scalable wireless FL. The approach advances digital AirComp by enabling end-to-end learning-based constellation design and joint transmitter-receiver optimization for efficient, scalable FL over real-world wireless links.

Abstract

Wireless federated learning (FL) relies on efficient uplink communications to aggregate model updates across distributed edge devices. Over-the-air computation (a.k.a. AirComp) has emerged as a promising approach for addressing the scalability challenge of FL over wireless links with limited communication resources. Unlike conventional methods, AirComp allows multiple edge devices to transmit uplink signals simultaneously, enabling the parameter server to directly decode the average global model. However, existing AirComp solutions are intrinsically analog, while modern wireless systems predominantly adopt digital modulations. Consequently, careful constellation designs are necessary to accurately decode the sum model updates without ambiguity. In this paper, we propose an end-to-end communication system supporting AirComp with digital modulation, aiming to overcome the challenges associated with accurate decoding of the sum signal with constellation designs. We leverage autoencoder network structures and explore the joint optimization of transmitter and receiver components. Our approach fills an important gap in the context of accurately decoding the sum signal in digital modulation-based AirComp, which can advance the deployment of FL in contemporary wireless systems.

Figures (6)

• Figure 1: Federated learning pipeline in the $t$-th global round.
• Figure 2: An example of overlapped constellation points.
• Figure 3: End-to-end communication system realized with autoencoder.
• Figure 4: Autoencoder results using parameters (2,2) with 8 clients.
• Figure 5: 2-bit modulation of CIFAR-10 results with IID and Non-IID local datasets and full clients (8 clients) participation.