Learned Digital Codes for Over-the-Air Computation in Federated Edge Learning

TL;DR

This paper introduces a learned digital over-the-air aggregation method for federated edge learning that enhances robustness and accuracy in low SNR conditions by integrating neural decoding with digital coding schemes.

Contribution

It proposes a novel end-to-end learned digital OTA framework combining URA codes, vector quantization, and AMP-DA-Net, extending OTA aggregation to various symmetric functions and improving low SNR performance.

Findings

Extends reliable digital OTA operation by over 10 dB in low SNR regimes.

Achieves comparable or improved performance across all SNR levels.

Maintains effectiveness under message corruption and nonlinear aggregation.

Abstract

Federated edge learning (FEEL) enables wireless devices to collaboratively train a centralised model without sharing raw data, but repeated uplink transmission of model updates makes communication the dominant bottleneck. Over-the-air (OTA) aggregation alleviates this by exploiting the superposition property of the wireless channel, enabling simultaneous transmission and merging communication with computation. Digital OTA schemes extend this principle by incorporating the robustness of conventional digital communication, but current designs remain limited in low signal-to-noise ratio (SNR) regimes. This work proposes a learned digital OTA framework that improves recovery accuracy, convergence behaviour, and robustness to challenging SNR conditions while maintaining the same uplink overhead as state-of-the-art methods. The design integrates an unsourced random access (URA) codebook with vector quantisation and AMP-DA-Net, an unrolled approximate message passing (AMP)-style decoder trained end-to-end with the digital codebook and parameter server local training statistics. The proposed design extends OTA aggregation beyond averaging to a broad class of symmetric functions, including trimmed means and majority-based rules. Experiments on highly heterogeneous device datasets and varying numbers of active devices show that the proposed design extends reliable digital OTA operation by more than 10 dB into low SNR regimes while matching or improving performance across the full SNR range. The learned decoder remains effective under message corruption and nonlinear aggregation, highlighting the broader potential of end-to-end learned design for digital OTA communication in FEEL.