Positional Attention-based Graph Neural Network for Learning Permutation Non-equivariant Wireless Policies

Baichuan Zhao, Chenyang Yang, Jianyu Zhao, Di Zhang

Graph neural networks (GNNs) have emerged as a promising approach to learning wireless policies efficiently by leveraging topology prior and incorporating relational inductive biases. However, when the optimal policy is not permutation equivariant (PE), conventional GNNs suffer from mismatched inductive biases, leading to degraded performance or poor generalizability. This issue arises in wireless tasks with expected objectives, such as channel estimation and end-to-end (E2E) precoding, where the PE property of the optimal policy depends on the underlying channel distribution. In this paper, we propose a novel positional attention-based GNN to learn permutation nonequivariant policies efficiently. The core idea is to incorporate relative positions of vertices into the attention mechanism via an embedding function, enabling the GNNs to capture asymmetric relationships. Consequently, the proposed GNN can represent permutation non-equivariant functions, while retaining high learning efficiency and size generalizability through parameter sharing. We consider channel estimation and E2E precoding as case studies, and prove that their policies are PE to users but not to antennas under spatially correlated channels. We employ the proposed GNN to learn the policies, where the embedding function is designed based on the channel covariance matrix. Simulation results demonstrate that the proposed GNN outperforms existing channel estimation and E2E precoding methods, requires fewer samples for training, and can be generalized to systems with different numbers of antennas and users.