Towards Achieving Integer and Load-balancing User Association in Wireless Networks with a Reparameterized Attention-based GNN

Machine learning (ML) is a promising method for user association in dense wireless networks, where each user (UE) must connect to a unique base station (BS) to balance the loads and maximize network capacity. Two key challenges that hinder direct ML use are producing integer‐valued associations while still maintaining gradients, and satisfying BS load constraints. Most existing approaches relax the integer association requirement by using \textit{softmax}, leading to suboptimal results. To address this problem, we propose an attention‐based Graph Neural Network with Gumbel‐Softmax reparameterization for near‐integer outputs, together with Sinkhorn‐Knopp normalization and loss regularization for load balancing satisfactions. Numerical results show that our method outperforms all existing ML and non-ML solutions, approaches optimal exhaustive search in small networks, and generalizes well to larger and more dynamic networks without retraining.