Two-Stage Heterogeneous Graph Neural Network for RIS-Aided Physical-Layer Security

1. Problem Statement

The paper addresses the challenge of optimizing Physical-Layer Security (PLS) in a Reconfigurable Intelligent Surface (RIS)-assisted Multi-Input Single-Output (MISO) system.

• System Setup: A Base Station (BS) with multiple antennas serves multiple Legitimate Users (LUs) while facing multiple Eavesdroppers (Eves). The system utilizes an RIS with $L$ reflecting elements to enhance channel quality and security.
• Objective: Maximize the Secrecy Energy Efficiency (SEE), defined as the ratio of the sum secrecy rate to the total power consumption (including transmit power and circuit power).
• Optimization Variables: The problem requires the joint optimization of:
  1. Beamforming vectors ($w_k$) for information transmission to LUs.
  2. Artificial Noise (AN) vectors ($z_m$) to disrupt Eves.
  3. RIS Phase Shift Matrix ($\Phi$).
• Challenges:
  • Deep Coupling: The variables are non-convex and deeply coupled, making traditional convex optimization (CVX) computationally expensive (high latency).
  • Scalability: Existing Deep Learning (DL) models (e.g., MLPs, CNNs) often fail to generalize when the number of RIS elements, users, or eavesdroppers changes dynamically due to fixed input/output dimensions.
  • Heterogeneity: The system involves distinct entities (BS, RIS elements, LUs, Eves) with different roles and interactions, which standard homogeneous Graph Neural Networks (GNNs) struggle to model effectively.

2. Methodology: Two-Stage Heterogeneous GNN (HGNN)

The authors propose a Two-Stage HGNN trained in an unsupervised manner (using the SEE objective as a loss function) to map Channel State Information (CSI) directly to the optimization variables. The architecture is designed to be permutation-equivariant, ensuring scalability.

Stage 1: RIS Phase Shift Optimization

• Graph Representation: The system is modeled as a Bipartite Graph ($G^{(1)}$) with three node types:
  1. BS-RIS Link Nodes: Representing individual RIS reflecting elements.
  2. BS-LU Link Nodes: Representing legitimate users.
  3. BS-Eve Link Nodes: Representing eavesdroppers.
• Mechanism:
  • Uses Edge-based and Edge-free Graph Attention operators to capture interactions between RIS elements and users/eavesdroppers.
  • Incorporates Residual Connections to mitigate gradient vanishing in deep layers.
  • Output: Generates the RIS phase shift matrix ($\Phi$). The output of this stage is fed into Stage 2 and the loss function.

Stage 2: Beamforming and AN Optimization

• Graph Representation: The system is modeled as a Fully-Connected Heterogeneous Graph ($G^{(2)}$) with two node types:
  1. BS-RIS-LU Nodes: Combining effective CSI (after applying $\Phi$) and augmented features from Stage 1.
  2. BS-RIS-Eve Nodes: Combining effective CSI and augmented features from Stage 1.
• Mechanism:
  • Uses Edge-free Graph Attention and a Semantic Attention mechanism to learn the complex relationships between users and eavesdroppers in the presence of the optimized RIS.
  • Output Strategies: Two methods are proposed for generating beamforming ($w_k$) and AN ($z_m$) vectors:
    1. Beam-Direct: Directly maps node embeddings to vectors via MLPs.
    2. Model-Based: Integrates theoretical knowledge (Hybrid Zero-Forcing and Maximum Ratio Transmission) to predict direction and power parameters, ensuring physical constraints are met more robustly.

Key Architectural Features

• Unsupervised Training: The loss function is the negative of the SEE, eliminating the need for labeled optimal solutions (which are hard to obtain).
• Scalability: By treating RIS elements, LUs, and Eves as nodes in a graph, the model's parameters are independent of the number of nodes. New nodes can be added/removed without retraining.

3. Key Contributions

1. Novel Framework: First work to integrate a multi-stage framework with Heterogeneous Graph Neural Networks (HGNNs) specifically for RIS-aided PLS. This addresses both the algorithmic structure (decomposing the problem) and the system heterogeneity.
2. Scalable Architecture: The proposed HGNN is permutation-equivariant, meaning it can handle dynamic changes in the number of RIS elements, users, and eavesdroppers without retraining, overcoming the limitations of MLPs and CNNs.
3. Dual-Stage Design:
   • Stage 1 explicitly models the RIS structure (bipartite) to optimize phase shifts.
   • Stage 2 models the user-eavesdropper interaction (fully connected) to optimize beamforming and AN.
4. Model-Based Output: Introduction of a "Model-Based" output layer that hybridizes DL with classical signal processing (ZF/MRT), significantly improving performance and stability compared to purely data-driven outputs.

4. Numerical Results

The model was evaluated against Convex Optimization (CVXopt), MLPs, and other GNN baselines (CO-GNN, BHGNN).

• Performance vs. CVXopt:
  • The proposed model achieves >96% of the optimal SEE obtained by CVXopt.
  • Inference Time: It reduces running time by three orders of magnitude (e.g., from ~82 seconds for CVXopt to ~11.6 ms for the proposed model).
• Performance vs. Baselines:
  • Outperforms state-of-the-art GNNs (CO-GNN, BHGNN) and MLPs significantly.
  • In scalability tests (varying $L$, $K$, $M$), the proposed model maintains high performance (>95% of optimal), whereas baselines degrade significantly (e.g., MLP drops to ~20-30%).
• Ablation Study:
  • The Two-Stage Framework alone provided a 19.07% gain.
  • The Model-Based Output provided the largest gain (43.56%), highlighting the importance of integrating domain knowledge.
  • Residual Connections contributed a 1.54% gain.

5. Significance

• Real-Time Security: The drastic reduction in computational time (milliseconds vs. minutes) makes the approach viable for real-time PLS in dynamic 6G networks where channel conditions change rapidly.
• Dynamic Adaptability: The ability to scale to varying numbers of RIS elements and users without retraining is crucial for practical deployment in heterogeneous networks.
• Generalization: The study demonstrates that aligning DL architectures with the physical topology of the wireless system (via HGNNs) and decomposing the optimization problem (via multi-stage frameworks) yields superior results compared to "black-box" deep learning approaches.
• Energy Efficiency: By maximizing SEE, the method aligns with green communication goals, balancing security performance with energy consumption.

In conclusion, this paper presents a robust, scalable, and efficient solution for RIS-aided physical-layer security, bridging the gap between theoretical optimization and practical, real-time implementation.