Optimized Amplify-and-Forward Relaying for Hierarchical Over-the-Air Computation

Over-the-air computation (AirComp) is an emerging wireless technique with wide applications (e.g., in distributed edge learning), which can swiftly compute functions of distributed data from different wireless devices (WDs) by exploiting the super-position property of wireless channels. Different from prior works focusing on the AirComp over one single cell in a small area, this paper considers a new hierarchical architecture to enable AirComp in a large area, in which a set of intermediate relays are exploited to help the fusion center to aggregate data from massive WDs for functional computation. In particular, we present a two-phase amplify-and-forward (AF) relaying design for hierarchical AirComp. In the first phase, the WDs simultaneously send their data to the relays, while in the second phase, the relays amplify the received signals and concurrently forward them to the fusion center for aggregation. Under this setup, we minimize the computation distortion measured by the mean squared error (MSE), by jointly optimizing the transmit coefficients at the WDs and relays and the de-noising factor at the fusion center, subject to their individual transmit power constraints. For the highly non-convex MSE minimization problem, we develop an alternating-optimization-based algorithm to obtain a high-quality solution. The optimized solution shows that for each WD, the phase of its transmit coefficient is opposite to that of the composite channel from the WD itself to the relays to the fusion center, such that they can be aligned at the fusion center, and its transmit power follows a regularized composite-channel-inversion structure to strike a balance between minimizing the signal misalignment error and the noise-induced error. Numerical results show that our proposed design achieves a significant MSE performance gain over benchmark schemes with full-power transmission at the WDs and/or relays.