Hierarchical Transformer for Power Allocation in Multi-Cell Massive MIMO Systems

Abstract

Massive multiple-input multiple-output (MIMO) systems is one of the key scenarios for 5G and future 6G networks. Power allocation is a critical yet computationally expensive component in multi-cell MIMO downlink systems. While the weighted minimum mean-square error (WMMSE) algorithm achieves near-optimal sum-rate performance, its iterative nature incurs prohibitive latency for real-time deployment. Inspired by the great success of transformer architecture in computer vision, we propose a deep learning method with hierarchical transformer architecture for joint multi-cell power allocation in this paper. The proposed method employs a shared link encoder to extract per-link channel features, followed by a two-level transformer structure. This hierarchical architecture contains an intra-BS transformer that captures channel correlations among all users within each base station (BS) via self-attention, and an inter-BS transformer that models crosscell interference coordination through BS-level self-attention. By directly maximizing a differentiable sum-rate objective, the model is trained without requiring WMMSE labels, significantly reducing data generation overhead. Simulation results show that the proposed method achieves near optimal performance to the WMMSE method, while achieving a $125 \times$ inference speedup over WMMSE. The proposed method consistently outperforms the traditional learning-based method such as fully-connected network and convolutional neural network.