Cooperative Beamforming and RISs Association for Multi-RISs Aided
  Multi-Users MmWave MIMO Systems through Graph Neural Networks
Reconfigurable intelligent surface (RIS) is considered as a promising
solution for next-generation wireless communication networks due to a variety
of merits, e.g., customizing the communication environment. Therefore,
deploying multiple RISs helps overcome severe signal blocking between the base
station (BS) and users, which is also a practical and effective solution to
achieve better service coverage. However, reaping the full benefits of a
multi-RISs aided communication system requires solving a non-convex,
infinite-dimensional optimization problem, which motivates the use of
learning-based methods to configure the optimal policy. This paper adopts a
novel heterogeneous graph neural network (GNN) to effectively exploit the graph
topology in the wireless communication optimization problem. First, we
characterize all communication link features and interference relations in our
system with a heterogeneous graph structure. Then, we endeavor to maximize the
weighted sum rate (WSR) of all users by jointly optimizing the active
beamforming at the BS, the passive beamforming vector of the RIS elements, as
well as the RISs association strategy. Unlike most existing work, we consider a
more general scenario where the cascaded link for each user is not fixed but
dynamically selected by maximizing the WSR. Simulation results show that our
proposed heterogeneous GNNs perform about 10 times better than other
benchmarks, and a suitable RISs association strategy is also validated to be
effective in improving the quality services of users by 30%.
Mengbing Liu, Chongwen Huang, Marco Di Renzo, Merouane Debbah, Chau Yuen
Cooperative Beamforming and RISs Association for Multi-RISs Aided Multi-Users MmWave MIMO Systems through Graph Neural Networks

###### Abstract

Reconfigurable intelligent surface (RIS) is considered as a promising solution for next-generation wireless communication networks due to a variety of merits, e.g., customizing the communication environment. Therefore, deploying multiple RISs helps overcome severe signal blocking between the base station (BS) and users, which is also a practical and effective solution to achieve better service coverage. However, reaping the full benefits of a multi-RISs aided communication system requires solving a non-convex, infinite-dimensional optimization problem, which motivates the use of learning-based methods to configure the optimal policy. This paper adopts a novel heterogeneous graph neural network (GNN) to effectively exploit the graph topology in the wireless communication optimization problem. First, we characterize all communication link features and interference relations in our system with a heterogeneous graph structure. Then, we endeavor to maximize the weighted sum rate (WSR) of all users by jointly optimizing the active beamforming at the BS, the passive beamforming vector of the RIS elements, as well as the RISs association strategy. Unlike most existing work, we consider a more general scenario where the cascaded link for each user is not fixed but dynamically selected by maximizing the WSR. Simulation results show that our proposed heterogeneous GNNs perform about 10 times better than other benchmarks, and a suitable RISs association strategy is also validated to be effective in improving the quality services of users by 30\(\%\).

 Beamforming, graph neural network, reconfigurable intelligent surfaces, RISs association strategy.

## I Introduction

Multiple-input multiple-output (MIMO) technology continues to play an indispensable role in providing high spectral efficiency in the current and next-generation wireless systems [1]. Although MIMO technology can bring significant antenna gains, the signal transmitted at millimeter wave (mmWave) bands experiences serious propagation attenuation. Moreover, due to the existence of obstacles in the environment, MIMO technology cannot provide seamless coverage between transceivers [2]. As a remedy, reconfigurable intelligent surface (RIS) is considered as a promising technology to improve the links' quality of service as well as recover communications in the dead zone. Compared with other technologies, such as ultra dense networks, RIS is regarded as a promising paradigm to assist mmWave MIMO networks via a large number of effective and low-cost reflective elements [3, 4, 5].

Due to the benefits above, RIS has been introduced to various communication systems to balance the network capacity and energy consumption. To fully reap the benefits, the cooperative beamforming design on base stations (BSs) with MIMO and RISs is a crucial challenge since many beamforming designs for the distributed RISs system are non-convex and heavily computational. Thus, a great deal of efforts have been devoted to developing effective algorithms to tackle this challenge in RIS-aided systems [6, 7]. Recently, the cooperative beamforming design in multi-RISs aided system has been investigated to compensate for the coverage hole and the serious propagation attenuation in high-frequency bands [8]. However, these problem-specific algorithms are laborious processes and require much problem-specific knowledge, especially in multi-RISs aided systems.

Inspired by the recent successes of deep learning technology in many application domains, e.g., the computer vision [9], deep neural network (DNN) models have been utilized for RIS-aided communication systems in recent years [10, 11]. However, such DNN models may not be very suitable for the problems in wireless communications. This is because DNNs fail to capture the underlying topology of wireless communication networks, resulting in dramatic performance degradation when the network size becomes large. Thus, in
order to improve the scalability and generalization of the network, a long-standing idea is to incorporate the structures of the target task into the neural network architecture [12].

From this perspective, the graph neural network (GNN) model has emerged as a new direction to exploit the graph topology of wireless communication networks [13]. There are some works applying the GNN for resource allocation in device-to-device (D2D) communication networks [14, 15]. In these works, the service links and interference relations are modeled as nodes and edges, respectively. Based on these graph embedding features, the optimization problem on the weighted sum rate (WSR) can be solved efficiently.

In this paper, we leverage the advantage of the GNN to solve the beamforming design in a multi-RISs aided communication system. Unfortunately, the modeling method for above D2D systems is not suitable for our system. In D2D systems, the transceiver pair (including one transmitter and one receiver) is denoted as one node without interference inside, and the interference in the system only exists between different transceiver pairs. However, we consider the communication links between one transmitter (BS) and multiple receivers (users) in this paper, where the inter-user interference is included. Thus, a new modeling approach is needed to better map the multi-RISs aided system into the graph.

Aside from the beamforming design, we consider a more practical and general problem in the multi-RISs scenario, assuming that each user only selects one of the RISs to enhance the links' service quality. In fact, selecting all RISs is not an efficient way to assist each user, because RISs are generally deployed in different positions. The distances between the user and some RISs may be so far that the signals reflected by the RISs are severely attenuated. For these reasons, the RISs association is introduced into this problem which needs to be jointly optimized with the cooperative beamforming at both the BS and RISs to obtain better performance. Thus, our main contributions in this work are concluded as

1) We first propose a new GNN structure to map the optimization problem in a multi-RISs aided system and validate its effectiveness through extensive experiments.

2) A RISs association strategy is considered in our optimization problem. By comparing the results obtained with fully connected RISs, we find that the strategy can efficiently improve the quality services of users by \(30\%\).

3)In order to exhibit the benefits of applying structure information, DNNs with different network sizes are adopted as benchmarks. It is clear that our proposed heterogeneous GNNs perform about 10 times better than DNNs.

## II Multi-RISs Aided Massive MIMO Communication Networks

In this section, we introduce a multi-RISs aided MIMO system model with a RISs association strategy.

### _System Model_

We consider a multi-RISs aided system with a \(N_{t}\)-antenna BS, a set of single-antenna users \(\mathcal{K}=\{1,2,\cdots,K\}\) and a set of RISs \(\mathcal{R}=\{1,2,\cdots,R\}\). Each RIS is equipped with \(M=M_{x}\times M_{y}\) reflective elements, where \(M_{x}\) and \(M_{y}\) represent the the number of horizontal and vertical elements, respectively. Assume that the BS transmits \(K\) data streams to users through the beamforming architecture. Thus, the transmitted signal \(\mathbf{x}\) from the BS can be expressed as

\[\mathbf{x}\triangleq\mathbf{F}\mathbf{s}=\sum_{k=1}^{K}\mathbf{f}_{k}s_{k}, \tag{1}\]

where \(\mathbf{s}=[s_{1},\cdots,s_{K}]^{T}\in\mathbb{C}^{K\times 1}\) indicates the symbol vector and \(\mathbb{E}(\mathbf{s}\mathbf{s}^{H})=\mathbf{I}_{K}\). \(\mathbf{F}=[\mathbf{f}_{1},\cdots,\mathbf{f}_{K}]\in\mathbb{C}^{N_{t}\times K}\) denotes the beamforming matrix at the BS and the power constraint is denoted by \(\left\|\mathbf{F}\right\|_{\mathcal{F}}^{2}\leq P_{max}\), where \(P_{max}\) is the maximum transmit power by the BS. It is crucial to investigate the joint design of RISs association, BS beamforming, and the phase configuration of RISs. In this paper, we choose one of all RISs for each user to maximize the WSR of the multi-RISs aided massive MIMO system. To simply express the association between RISs and users, a matrix \(\mathbf{U}=[\mathbf{u}_{1},\cdots,\mathbf{u}_{K}]^{T}\in\mathbb{C}^{K\times R}\) consisting of binary variables \(u_{k,i}\) is defined as

\[u_{k,i}=\left\{\begin{array}{l}1,\text{user}_{k}\text{ is assisted by RIS}_{i}.\\ 0,\text{ otherwise}.\end{array}\right. \tag{2}\]

Thus, based on the RISs association strategy, the equivalent channel \(\mathbf{h}_{k}\) from the BS to user\({}_{k}\) can be expressed as

\[\mathbf{h}_{k}=\sum_{i\in\mathcal{R}}u_{k,i}\mathbf{h}_{ik}\mathbf{\Phi}_{i} \mathbf{G}_{i}, \tag{3}\]

where \(\mathbf{G}_{i}\in\mathbb{C}^{M\times N_{t}}\) and \(\mathbf{h}_{ik}\in\mathbb{C}^{1\times M}\) denote the BS-RIS\({}_{i}\) and RIS\({}_{i}\)-user\({}_{k}\), respectively. \(\mathbf{\Phi}_{i}=\mathrm{diag}\left(\theta_{i,1},\cdots,\theta_{i,M}\right)\in \mathbb{C}^{M\times M}\) with the diagonal element \(\theta_{i,m}=e^{j\phi_{i,m}}\) denoting the phase shift matrix of RIS\({}_{i}\).

According to above assumption, the signal to interference plus noise ratio (SINR) at user\({}_{k}\) can be given as

\[\mathrm{SINR}_{k}=\frac{\left|\mathbf{h}_{k}\mathbf{f}_{k}\right|^{2}}{\sum_{ j\in\mathcal{K},j\neq k}^{K}\left|\mathbf{h}_{k}\mathbf{f}_{j}\right|^{2}+ \sigma^{2}}, \tag{4}\]

where \(\sigma^{2}\) denotes the effective noise variance.

Fig. 1: Illustration of the multi-RISs aided massive MIMO system model.
### _Channel Model_

In this paper, we consider the mmWave channel and adopt the typical Saleh-Valenzuela channel model [16] with limited scattering paths to characterize the sparse nature. Supposing a uniform linear array (ULA) is equipped at the BS and a uniform planar array (UPA) is equipped at each RIS. Hence, we can get the array response vectors by the following equations

\[\mathbf{a}_{\mathrm{ULA}}(\phi) =\frac{1}{\sqrt{N}}\left[e^{-j\frac{2\pi d}{\lambda}\phi_{i}} \right]_{i\in\mathcal{I}(N)}, \tag{5}\] \[\mathbf{a}_{\mathrm{UPA}}(\phi^{1},\phi^{2}) =\mathbf{a}_{\mathrm{ULA}}^{\mathrm{x}}(\phi^{1})\otimes\mathbf{ a}_{\mathrm{ULA}}^{\mathrm{y}}(\phi^{2}), \tag{6}\]

where \(\lambda\) and \(d\) represent the signal wavelength and the antenna spacing, respectively. As for a ULA with \(N\) antennas, \(\phi\) denotes the angle of arrival (AoA) or the angle of departure (AoD). For a UPA with \(M=M_{y}\times M_{x}\) elements, \(\phi^{1}\) and \(\phi^{2}\) denote the azimuth and elevation angles, respectively. \(\mathbf{a}_{\mathrm{ULA}}^{\mathrm{x}}(\phi^{1})\) and \(\mathbf{a}_{\mathrm{ULA}}^{\mathrm{y}}(\phi^{2})\) are defined in the same manner as \(\mathbf{a}_{\mathrm{ULA}}(\phi)\) with \(\mathcal{I}(N)=\{n-(N-1)/2,n=0,1,\cdots,N-1\}\), and \(\otimes\) denotes the Kronecker product. The array element spacings of ULA and UPA are assumed to be \(\lambda/2\). Hence, the effective channels can be represented as

\[\mathbf{h}_{ik} =\sum_{l_{h}=1}^{L_{h}}\beta_{l_{h}}\mathbf{a}_{t}^{H}\left( \phi_{l_{h}}^{\mathrm{r},1},\phi_{l_{h}}^{\mathrm{r},2}\right), \tag{7}\] \[\mathbf{G}_{i} =\sum_{l_{g}=1}^{L_{g}}\beta_{l_{g}}\mathbf{a}_{t}\left(\phi_{l_{ g}}^{\mathrm{r},1},\phi_{l_{g}}^{\mathrm{r},2}\right)\mathbf{a}_{b}^{H}\left( \phi_{l_{g}}^{\mathrm{h}}\right), \tag{8}\]

where \(L_{s}\) is the number of paths that contains one LoS path and \((L_{s}-1)\)\(\mathrm{NLoS}\) paths for link, and \(\beta_{l_{s}}\) denotes the complex path gain of \(l_{s}\) path with \(s\in\{h,g\}\). For simply notation, the subscript \(h,g\) represent \(\mathbf{h}_{ik}\) and \(\mathbf{G}_{i}\), respectively.

### _Problem Formulation_

Our main objective is to maximize the WSR by designing the beamforming matrix \(\mathbf{F}\), the RISs association matrix \(\mathbf{U}\) as well as the passive beamforming matrices \(\mathbf{\Phi}_{i},\forall i\in\mathcal{R}\) in multi-RISs aided communication systems. Therefore, the maximization problem can be formulated as

\[\text{P1:}\ \max_{\mathbf{F},\mathbf{U},\mathbf{\Phi}_{i}, \forall i\in\mathcal{R}} \sum_{k=1}^{K}w_{k}\mathrm{log}_{2}(1+\mathrm{SINR}_{k}) \tag{9a}\] \[\text{s.t.} \|\mathbf{F}\|_{\mathcal{F}}^{2}\leq P_{max}\] (9b) \[|\mathbf{\Phi}_{i,m}|=1,\forall m=1,2,\cdots,M\] (9c) \[\sum_{i\in\mathcal{R}}u_{k,i}=1,\forall k\in\mathcal{K}\] (9d) \[u_{k,i}=\left\{0,1\right\},\forall k\in\mathcal{K},\forall i\in \mathcal{R}, \tag{9e}\]

where \(w_{k}\) denotes the weight for user\({}_{k}\) with \(\sum_{k=1}^{K}w_{k}=1\).

## III Sum-Rate Maximization Via GNN

In this paper,we introduce the heterogeneous GNN which contains abundant information with structural relations among multi-typed nodes as well as unstructured content associated with each node [17]. As for a heterogeneous graph \(\mathcal{G}=(\mathcal{V},\mathcal{E},\mathcal{O}_{\mathcal{V}},\mathcal{R}_{ \mathcal{E}})\), there are two additional type sets, including the predefined node types \(\mathcal{O}_{\mathcal{V}}\) and edge types \(\mathcal{R}_{\mathcal{E}}\) with \(|\mathcal{O}_{\mathcal{V}}|+|\mathcal{R}_{\mathcal{E}}|>2\)[18]. The node types \(\mathcal{O}_{\mathcal{V}}\) includes BS node \(\mathbf{v}_{b}\), RIS\({}_{i}\) node \(\mathbf{v}_{r_{i}}\) and user\({}_{k}\) node \(\mathbf{v}_{u_{k}}\). According to the wireless communication scenario, the attributes of \(\mathbf{v}_{u_{k}}\) include the weight \(w_{k}\), noise variance \(\sigma^{2}\) as well as \(\mathbf{U}\). Similarly, the attributes \(\mathbf{v}_{b}\) and \(\mathbf{v}_{r_{i}}\) include \(\mathbf{F}\) and \(\mathbf{\Phi}_{i}\), respectively. Given the heterogeneous GNN structure, the goal of interest is to find a policy \(p_{\Theta}(\cdot)\) with parameters \(\Theta\), mapping the graph \(\mathcal{G}\) to estimations of the optimal \(\mathbf{F},\mathbf{\Phi}_{i}\) and \(\mathbf{U}\).

### _The Initialization Layer of Proposed GNN_

The initial layers take the CSI information as the input features. In order to better learn the phase configuration of each RIS, we extract the phase shift matrix \(\mathbf{\theta}_{i}\). Thus, the equivalent channel \(\mathbf{h}_{k}\) can be expressed as

\[\mathbf{h}_{k} =\sum_{i\in\mathcal{R}}u_{k,i}\mathbf{\theta}_{i}\mathrm{diag}( \mathbf{h}_{ik})\mathbf{G}_{i} \tag{10}\] \[=\sum_{i\in\mathcal{R}}u_{k,i}\mathbf{\theta}_{i}\mathbf{H}_{ik}^{ \mathrm{cas}},\]

where \(\mathbf{\theta}_{i}=[\theta_{i,1},\cdots,\theta_{i,M}]\) and \(\mathbf{H}_{ik}^{\mathrm{cas}}=\mathrm{diag}(\mathbf{h}_{ik})\mathbf{G}_{i}\) denotes the cascaded channel links, which is considered as the input feature for each user node, i.e., \(\mathbf{v}_{u_{k}}^{(\mathrm{in})}=[\mathrm{Re}(\mathbf{H}_{ik}^{\mathrm{ cas}}),\mathrm{Im}(\mathbf{H}_{ik}^{\mathrm{cas}})]\).

### _Message Passing Paradigm_

Graph network block is adopted as the basic computation unit over graphs [19], as illustrated in Fig. 2. The update function for each type of node is illustrated in detail as follows.

#### Iii-B1 Update for RIS Nodes

The updated feature for the RIS nodes at step \(t\)\(\mathbf{v}_{r_{i}}^{[t]}\) is obtained by combining the features from the last state \(\mathbf{v}_{r_{i}}^{[t-1]}\) and all the user nodes \(\mathbf{v}_{u_{k}}^{[t-1]},\forall k\in\mathcal{K}\). Then, the messages from different type of nodes are given as

\[\mathbf{\xi}_{rr_{i}}=\phi_{r_{i}-r_{i}}^{(\mathrm{up})}(\mathbf{\xi}_{r_{i}},\mathbf{v }_{r_{i}}^{[t-1]}), \tag{11}\]

\[\mathbf{\xi}_{ur_{i}}=\phi_{u\to r_{i}}^{(\mathrm{up})}(\rho_{u\to r_{i}}(\mathbf{\xi}_{u_{1}}, \cdots,\mathbf{\xi}_{u_{K}}),\mathbf{v}_{r_{i}}^{[t-1]}), \tag{12}\]

where \(\mathbf{\xi}_{\mathrm{src}}=\phi_{\mathrm{src}\to r_{i}}(\mathbf{v}_{r_{i}}^{[t-1]})\) denotes message from the source node and \(\mathrm{src}\in\{r_{i},u_{1},\cdots,u_{K}\}\) for simple notation. Meanwhile, both \(\phi_{\mathrm{src}\to r_{i}}\) and \(\phi_{\mathrm{src}\to r_{i}}^{(\mathrm{up})}\) are two-layer MLPs. Besides, \(\rho_{u\to r_{i}}\) is the mean aggregation function for all the different user nodes. Then, in heterogeneous GNNs, the updated RIS nodes \(\mathbf{v}_{r_{i}}^{[t]}\) aggregate the feature information from different-type neighbors, which is given as

\[\mathbf{v}_{r_{i}}^{[t]}=\rho_{r_{i}}(\mathbf{\xi}_{rr_{i}},\mathbf{\xi}_{ur_{i}})+ \mathbf{v}_{r_{i}}^{[t-1]}, \tag{13}\]

where \(\rho_{r_{i}}\) denotes the operation of the mean aggregation. Besides, a residual operation is added in this step to avoid degradation and gradient explosion of our network [20].
#### Iii-B2 Update for User Nodes

The updated feature \(\mathbf{v}_{u_{k}}^{[t]}\) is obtained by combining the features from all the user nodes \(\mathbf{v}_{u_{k}}^{[t-1]}\) and RIS nodes \(\mathbf{v}_{r_{i}}^{[t-1]}\). In the same manner, the messages from different type of nodes are given as

\[\boldsymbol{\xi}_{ru_{k}}=\phi_{r\to u_{k}}^{(\mathrm{up})}(\rho_{r\to u_{k}}( \boldsymbol{\xi}_{r_{1}},\cdots,\boldsymbol{\xi}_{r_{R}}),\mathbf{v}_{u_{k}}^{ [t-1]}), \tag{14}\]

\[\boldsymbol{\xi}_{uu_{k}}=\rho_{u_{k}}^{(u)}(\phi_{u\to u_{k}}^{(\mathrm{up})}( \boldsymbol{\xi}_{u}),\phi_{u_{k}\to u_{k}}^{(\mathrm{up})}(\boldsymbol{\xi}_ {u_{k}},\mathbf{v}_{u_{k}}^{[t-1]}))), \tag{15}\]

where \(\boldsymbol{\xi}_{u}=\rho_{u\to u_{k}}(\boldsymbol{\xi}_{u_{1}},\cdots, \boldsymbol{\xi}_{u_{K}})\) and \(\boldsymbol{\xi}_{\mathrm{src}}=\phi_{\mathrm{src}\to r_{i}}(\mathbf{v}_{ \mathrm{src}}^{[t-1]})\) denotes message from the source node and \(\mathrm{src}\in\{r_{1},\cdots,r_{R},u_{1},\cdots,u_{K}\}\) for simple notation. Meanwhile, both \(\phi_{\mathrm{src}\to u_{k}}\) and \(\phi_{\mathrm{src}\to u_{k}}^{(\mathrm{up})}\) are two-layer MLPs. Besides, \(\rho_{u\to u_{k}}\) and \(\rho_{r\to u_{k}}\) are the max and mean aggregation functions, respectively. Then, in heterogeneous GNNs, the updated RIS nodes \(\mathbf{v}_{u_{k}}^{[t]}\) aggregate the feature information from different-type neighbors, which can be obtained by

\[\mathbf{v}_{u_{k}}^{[t]}=\rho_{u_{k}}(\boldsymbol{\xi}_{ru_{k}},\boldsymbol{ \xi}_{uu_{k}})+\mathbf{v}_{u_{k}}^{[t-1]}, \tag{16}\]

where \(\rho_{u_{k}}\) denotes the operation of the max aggregation.

#### Iii-B3 Output Layers for Node Features

After updating all the nodes, each node has multiple feature embeddings from different-type nodes. In order to obtain the final feature of each node \(\mathbf{v}_{b}^{[T]}\), \(\mathbf{v}_{r_{i}}^{[T]}\) and \(\mathbf{v}_{u_{k}}^{[T]}\), a decoder layer is used to output all the optimized features, which is given as

\[\mathbf{v}_{\mathrm{dst}}^{[T]}=\rho_{\mathrm{dst}}^{(\mathrm{out})}( \boldsymbol{\xi}_{\mathrm{udst}},\boldsymbol{\xi}_{\mathrm{rdst}}), \tag{17}\]

where the destination node is denoted as \(\mathrm{dst}\in\{b,r_{1},\cdots,r_{R},u_{1},\cdots,u_{K}\}\) and \(\phi_{\mathrm{src}\to\mathrm{dst}}^{(\mathrm{out})}\) is utilized to extract the features from different-type nodes, as shown by the black line in Fig. 2. Finally, the new state \(\mathbf{v}_{\mathrm{dst}}^{[T]}\) will be obtained by a mean aggregation layer \(\rho_{\mathrm{dst}}^{(\mathrm{out})}\).

Suppose that \(T\) iterations are executed during the training process, and we defined each update as \(\text{HGN}_{t},\forall t\in T\). An encoder-process-decoder architecture is utilized to facilitate training [19], including one encoder block \(\text{HGN}_{0}\), one decoder block \(\text{HGN}_{T}\), and one core block (concatenate by \(\text{HGN}_{1},\cdots\), \(\text{HGN}_{T-1}\)). After the encoder-process-decoder architecture, we can obtain the final result for the optimal \(\mathbf{F}\), \(\mathbf{U}\) and \(\boldsymbol{\Phi}_{i},\forall i\in\mathcal{R}\).

### _Model Architecture of Proposed Heterogeneous GNN_

In addition to the massage passing paradigm, the proposed heterogeneous GNN needs some additional layers to complete the whole learning process as follows

#### Iii-C1 The Normalization Layers

Note that the neural network is based on the real numbers, we learn the vector consisting of the real and imaginary parts of the beamforming matrix \(\mathbf{F}\) and the phase matrices \(\boldsymbol{\Phi}_{i},\forall i\in\mathcal{R}\). Therefore, before calculating the WSR loss function, additional layers are necessary to merge the real-valued results to complex-valued results.

* To satisfy the constraint on transmit power in (9b), a normalization layer is employed at the output feature \(\mathrm{Re}\{\mathbf{F}\}\)and \(\mathrm{Im}\{\mathbf{F}\}\), which can be expressed as \[\hat{\mathbf{F}}=\sqrt{P_{max}}\frac{\mathrm{Re}\{\mathbf{F}\}+j\cdot\mathrm{ Im}\{\mathbf{F}\}}{\sqrt{\left\|\mathrm{Re}\{\mathbf{F}\}\right\|_{\mathcal{F}}^{2}+ \left\|\mathrm{Im}\{\mathbf{F}\}\right\|_{\mathcal{F}}^{2}}}.\] (18)
* Similarly, the normalization layer for the reflection coefficients can be represented as \[\theta_{i,m}=\frac{\mathrm{Re}\{\theta_{i,m}\}+j\cdot\mathrm{Im}\{\theta_{i,m} \}}{\sqrt{|\,\mathrm{Re}\{\theta_{i,m}\}|^{2}+|\mathrm{Im}\{\theta_{i,m}\}|^{2 }}}.\] (19)

Fig. 2: The message passing paradigm in the heterogeneous GNN network.
#### Iii-C2 The RIS Association Strategy

Each user chooses a suitable one of all RISs as the reflection link to enhance the link quality, as the constraint (9d). Thus, the matrix \(\mathbf{U}\) for calculating the WSR can be obtained by

\[c_{u_{k,i}}=\mathrm{SoftMax}\left(u_{k,i}\right)=\frac{e^{u_{k,i}}}{\sum\limits_ {i=1}^{R}e^{u_{k,i}}}. \tag{20}\]

After obtaining the SoftMax score \(c_{u_{k,i}}\), we sort the scores of each user, i.e., each column in \(\mathbf{U}\). We select the largest value and set it to 1, indicating that the user\({}_{k}\) selects this RIS to assist the communication links while setting the remaining positions of this column to 0 to get the final matrix.

#### Iii-C3 The Loss Function

Due to the connection matrix is composed of binary elements, which are discontinuous variants. Thus, the process to determine \(\mathbf{U}\) as well as the optimization of \(\mathbf{F}\) and \(\mathbf{\Phi}\) is very difficult to learn. Considering this, we add a regular term related to the distance between users and RISs to the training loss. We adopt distance-dependent labels to represent the association relationship between the users and their nearest RIS and the categorical cross-entropy loss function to achieve it. Thus, the loss function is composed of the negative WSR and the categorical cross-entropy loss function, which can be expressed as

\[\mathrm{Loss}= \tag{21}\] \[-\sum\limits_{k=1}^{K}\left(w_{k}\mathrm{log}_{2}\left(1\!+\! \frac{\left|\mathbf{h}_{k}\mathbf{f}_{k}\right|^{2}}{\sum\limits_{j\in\mathcal{ K},j\neq k}^{K}\left|\mathbf{h}_{k}\mathbf{f}_{j}\right|^{2}\!+\!\sigma^{2}}\right)\!-\! \eta\mathbf{u}_{k}^{(1)}\log\mathbf{u}_{k}\right),\]

where \(\mathbf{u}_{k}^{(1)}\) and \(\mathbf{u}_{k}\) denote the distance-dependent label and optimized value for the user\({}_{k}\), respectively. \(\eta\) is the coefficient of the penalty term, and its role is to strike the balance of the significance of these two parts in the loss function, which can be calculated by \(\eta=\mathrm{WSR}_{p}/\mathrm{WSR}_{p_{0}}\). We set \(\mathrm{WSR}_{p}\) as the results in the pre-training process where the loss function is negative WSR, and \(\mathrm{WSR}_{p_{0}}\) is the baseline case when \(P_{max}=30\)dBm. For the training details, Adam optimizer [21] is adopted with a total of 200 epochs on one NVIDIA RTX 2080Ti GPU. The initial learning rate, weight decay, and batch size are set to \(1\times 10^{-4}\), \(5\times 10^{-5}\) and 128.

## IV Simulation Results

In this section, numerical simulations are provided to verify the effectiveness of our proposed heterogeneous GNN. Throughout the simulation, the BS is equipped with \(N_{t}=8\) antennas, and the number of RIS elements \(M\) is 16 with \(M_{x}=M_{y}=4\). The BS is assumed to be located on the origin and two RISs are located in \([30,25]\) and \([30,-25]\). In addition, two users are randomly located in a rectangular region in \([40:50,-25:25]\). Without loss of generality, the weights \(w_{k}\) are set to be equal. According to [16], the channel gain is taken as \(\beta\sim\mathcal{CN}\left(0,10^{-0.1\mathrm{PL}(r)}\right)\) where \(\mathrm{PL}(r)=\varrho_{a}+10\varrho_{b}lg(r)+\xi\) with \(\xi\sim\mathcal{N}(0,\sigma_{\xi}^{2})\). In addition, the channel realizations are produced by setting \(\sigma_{n}^{2}=-85dBm\), \(\varrho_{\mathrm{a}}=61.4,\varrho_{\mathrm{b}}=2\), and \(\sigma_{\xi}=5.8\ \mathrm{dB}\). The total number of samples is 110,000 where 100,000 and 10,000 of the whole generated data are used for training and validation, respectively. To see the performance, WMMSE with a random phase configuration is used as the performance baseline. Also, in order to validate the effectiveness of the graph structure, we compare the performance in our proposed heterogeneous GNN with that in DNN.

Fig. 3 shows the sum rate with respect to the maximum transmit power of the BS. Our proposed GNN structure performs well under all levels of transmit power, which exhibits the generality and effectiveness of our structure. First, the sum rate obtained by random configuration and heterogeneous GNN becomes higher with the increase of maximum transmit power \(P_{max}\). However, in the absence of optimization, the sum rate boosting from the increasing power is limited. In contrast, our proposed GNN structure exhibits its superb ability to solve the optimization problem P1 with different number of reflective elements (M = 9, 16, 25) and different \(P_{max}\), from \(10\) dBm to \(30\) dBm with an interval of \(5\) dBm.

Fig. 4 illustrates the total number of antennas with respect to the WSR performance. We also exhibit the WSR values in two other cases, including the results by WMMSE with the

Fig. 4: The weighted sum rate versus the number of antennas at the BS

Fig. 3: The sum rate versus the maximum transmit power at the BS.
random phase configuration and the optimal value obtained by our proposed GNN in the case of one fixed RIS, i.e., all the users receive the signal reflected by the RIS located at \([30,25]\). It is obvious that the WSR increase with the growth of the number of antennas under all cases, and heterogeneous GNN can achieve a higher WSR. We can choose a suitable one of all RISs to replace the fix RISs to provide service links for each user, which can improve the performance about \(30\%\) under all the different number of antennas.

Fig. 5 illustrates the WSR of heterogeneous GNNs and two DNNs with respect to the same size of the training set under the networks with different sizes. We can easily find that the performance in heterogeneous GNNs performs significantly better than those of DNNs, especially when the number of training sample is \(2\times 10^{5}\), GNNs perform 10 times better than DNNs. Since the introduction of the RISs association strategy will exacerbate the influence of structural information on optimization results, DNNs perform very poorly due to their lack of the ability to capture the underlying model structure. In addition, comparing the results of two DNNs, it can be seen that no matter how to increase the number of layers and the number of channels cannot improve the performance of DNNs. In addition, GNNs can get better optimization results with fewer samples, even if it is less than 10,000, and get better results with increasing training samples.

## V Conclusion

We proposed a heterogeneous GNN structure to solve the NP-hard beamforming optimization problem in a multi-RISs aided MIMO system by mapping the structure of the wireless system into a graph neural network. Different from most beamforming problems, a RISs association strategy was introduced into our problem, which makes the beamforming design more complex, but a suitable RISs association strategy can enhance the users' quality of service. The effectiveness of our proposed GNN structure and the RISs association strategy were validated through extensive simulations. According to the simulation results, we found that our proposed GNNs performed about 10 times better than the conventional DNNs. Moreover, the benefits of the RISs association strategy were illustrated by comparing the optimal results with/without this strategy, which improved the performance about \(30\%\).
