Association-Aware GNN for Precoder Learning in Cell-Free Systems

1. Problem Statement

The paper addresses the challenge of optimizing precoders in Cell-Free (CF) multi-user multi-antenna systems. Unlike traditional cellular systems where a User Equipment (UE) connects to a single Access Point (AP), CF systems allow a UE to be served by a dynamic subset of APs.

• Core Challenge: The optimal precoding policy depends not only on Channel State Information (CSI) but also on the dynamic UE-AP association status.
• Limitations of Existing Methods:
  • Numerical Algorithms (e.g., WMMSE): While near-optimal, they suffer from high computational complexity, making real-time implementation difficult.
  • Existing Deep Learning (DL) Methods: Most existing DL approaches (FNNs, CNNs, or standard GNNs) assume a fixed association (e.g., UE connected to all APs) or fail to fully exploit the Permutation Equivariance (PE) properties inherent in CF systems.
  • Generalization Issues: Standard models often struggle to generalize to varying system sizes (different numbers of UEs, APs, or antennas) and fluctuating association patterns without retraining.

2. Methodology: Association-Aware GNN (AAGNN)

The authors propose a novel Association-Aware Graph Neural Network (AAGNN) designed to explicitly incorporate UE-AP association status into the precoding design while adhering to the system's symmetry properties.

A. System Modeling & Problem Formulation

• Association-Aware Channel Definition: The authors decompose the channel matrix $H$ based on the association matrix $D$. They define two variants:
  • $\tilde{H}_D$: Represents desired signal channels (where $d_{km}=1$).
  • $\hat{H}_{\bar{D}}$: Represents interference channels (where $d_{km}=0$).
• 3D Permutation Equivariance (3D-PE): The optimal precoding policy in CF systems exhibits 3D-PE with respect to UEs, APs, and AP antennas. This means if the indices of these entities are permuted in the input, the output precoder must permute correspondingly. The proposed model is designed to strictly satisfy this property.

B. Network Architecture Design

1. Parameter-Sharing Structure:

   • To guarantee 3D-PE, the weight matrices in the GNN layers are constrained to a specific parameter-sharing structure.
   • Instead of learning a full dense weight matrix, the model learns only 6 free parameters per layer (e.g., $\tilde{o}_1, \tilde{o}_2, \tilde{p}, \tilde{q}_1, \tilde{q}_2, \tilde{r}$ and their counterparts for the interference path).
   • This structure ensures that the network treats all UEs, APs, and antennas symmetrically, drastically reducing the number of learnable parameters.

2. Graph Construction:

   • The system is modeled as a heterogeneous graph with UE vertices and AP antenna vertices.
   • Edges connect UEs to AP antennas. Edge features represent channel coefficients.
   • The update rule aggregates information from:
     • The edge itself (self-loop).
     • Other antennas serving the same UE.
     • Other UEs served by the same AP.
     • Other APs serving the same UE.

3. Attention Mechanism:

   • To enhance generalization across varying numbers of UEs, a specialized attention mechanism is introduced.
   • Instead of simple summation, the model uses learnable attention coefficients ($\alpha, \beta$) weighted by the correlation between the channel vector and the aggregated features. This allows the network to dynamically prioritize relevant information during the update step.

3. Key Contributions

• Explicit Association Awareness: The first GNN architecture to explicitly model the dynamic UE-AP association status as an input, distinguishing between desired and interference channels.
• 3D-PE Compliant Design: Derivation of a rigorous parameter-sharing structure that satisfies the 3D-PE property of CF precoding, ensuring the model is permutation-equivariant with respect to UEs, APs, and antennas simultaneously.
• Enhanced Generalization: Integration of a tailored attention mechanism that significantly improves performance when the number of UEs changes, a common requirement in dynamic networks.
• Efficiency: The constrained parameter structure leads to a massive reduction in model complexity (parameters) compared to standard DNNs and unconstrained GNNs.

4. Simulation Results

The AAGNN was evaluated against benchmarks including WMMSE (numerical baseline), AAGNN without attention (AAGNN-woA), EGAT, SUNet (CNN), and other GNN variants.

• Learning Performance:
  • AAGNN achieves the highest normalized sum rate.
  • Sample Efficiency: It maintains over 80% performance of the optimal WMMSE with only 16 training samples, whereas other methods require thousands of samples to converge.
• Generalization Capabilities:
  • Varying UEs: AAGNN generalizes well to unseen numbers of UEs (tested $K \in [2, 16]$), whereas AAGNN-woA and EGAT degrade significantly.
  • Varying APs/Antennas: AAGNN maintains high performance when the number of APs or antennas changes, outperforming SUNet and other GNNs which fail to generalize to these dimensions.
• Complexity:
  • Training: Requires significantly fewer samples, less training time, and fewer parameters (space complexity) to reach the target performance (90% of WMMSE).
  • Inference: Inference time is comparable to or faster than baseline methods, making it suitable for real-time deployment.

5. Significance

This work bridges the gap between theoretical optimality and practical deployment in Cell-Free networks. By respecting the inherent symmetries (3D-PE) of the physical system and explicitly modeling the dynamic association logic, the AAGNN offers a low-complexity, high-generalization solution for precoding. It demonstrates that incorporating domain-specific structural constraints (like permutation equivariance) into deep learning architectures is crucial for achieving robust performance in wireless systems with variable topologies.