## 4 Structure of the Network and Training Scheme
### Multi-Head Attention Layer  
In Transformer architecture, the multi-head attention layer in Figure 1 is the core layer. Figure 2 shows the internal data processing of the layer. The particular attention function of Transformer is called Scaled Dot-Product Attention [36]. For each CSI matrix \(\mathbf{H_{a}}\), it is multiplied by \(\mathbf{W^{Q}}\), \(\mathbf{W^{K}}\), \(\mathbf{W^{V}}\in\mathbb{R}^{N_{1}\times d_{\text{model}}}\) to compute the corresponding query, key, and value matrix, which are denoted as \(\mathbf{Q}\), \(\mathbf{K}\), \(\mathbf{V}\). \(\mathbf{W^{Q}}\), \(\mathbf{W^{K}}\) and \(\mathbf{W^{V}}\) are parameters that the model needs to learn and \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) will be used to calculate the attention score. We take the row vector of a CSI matrix as the unit. The score is used to measure the magnitude of  
Figure 1: The structure of the encoder and decoder of our network.  
attention, which expresses how much we should pay attention to the other row vectors when coding each row vector. Note that \(d_{\text{model}}\) is the dimension of the model. In order to get higher feature expression for TransNet+ in the attention layer, \(d_{\text{model}}\) is adjusted to 2 times compared with the original method. To obtain the attention score, first we let \(\mathbf{Q}\) get multiplied by \(\mathbf{K}^{\text{T}}\). After that, we divide the score by \(\sqrt{d_{\text{model}}}\), thus making it more stable to find the gradient during backpropagation when the network is trained. Then, the softmax function is indispensable for normalizing the score. Finally, we do a weighted summation for each row of \(\mathbf{V}\) with attention score. The entire process is described by the equation  
\[\mathbf{Z}=\text{softmax}\bigg{(}\frac{\mathbf{Q}\mathbf{K}^{\text{T}}}{ \sqrt{d_{\text{model}}}}\bigg{)}\mathbf{V}\in\mathbb{R}^{2N_{\text{a}}\times d _{\text{model}}} \tag{5}\]  
The resulting matrix \(\mathbf{Z}\) characterizes the attention of each row of \(\mathbf{H}_{\text{a}}\) with respect to the other rows, i.e., the magnitude of correlation. The larger the calculated score means that the correlation is likely to be stronger.  
The multi-head attention mechanism further refines the self-attention mechanism. In TransNet, the number of attention heads named \(n_{\text{head}}\) is 2. In fact, bigger \(n_{\text{head}}\) has the potential to jointly attend to information from different representation subspaces at different positions [37]. As is shown in Figure 2, we modify \(n_{\text{head}}\) to 4, which means \(\mathbf{H}_{\text{a}}\) requires to be multiplied by 4 different sets of \(\mathbf{W}^{\text{Q}}\), \(\mathbf{W}^{\text{K}}\) and \(\mathbf{W}^{\text{V}}\). Note that the size of each \(\mathbf{W}^{\text{Q}}\), \(\mathbf{W}^{\text{K}}\) and \(\mathbf{W}^{\text{V}}\) should be modified to \(N_{\text{t}}\times(d_{\text{model}}/n_{\text{head}})\), so the number of columns of the final \(\mathbf{Z}\) can be \(d_{\text{model}}\). Since the feed-forward neural network receives only one matrix, we concatenate the results of the four attention heads and let it multiply with a weight matrix \(\mathbf{W}^{\text{O}}\in\mathbb{R}^{d_{\text{model}}\times d_{\text{model}}}\) following  
\[\mathbf{Z}=[\mathbf{Z}_{1}\ \mathbf{Z}_{2}\ \mathbf{Z}_{3}\ \mathbf{Z}_{4}] \mathbf{W}^{\text{O}}\in\mathbb{R}^{2N_{\text{a}}\times d_{\text{model}}} \tag{6}\]  
The result of the formula is the output of a single multi-head attention layer and will be fed into the add and norm layer.  
Figure 2: The internal structure of a single multi-head attention layer with \(n_{\text{head}}=4\). The second to fourth sets of \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) have also been multiplied by three different sets of \(\mathbf{W}^{\text{Q}}\), \(\mathbf{W}^{\text{K}}\) and \(\mathbf{W}^{\text{V}}\), which is omitted in the figure. The dotted arrow indicates that \(\mathbf{V}\) is not processed by the softmax layer.