The training procedure of DSA-HBF is more challenging than the two previous structures since the matrix \(\mathbf {S_{\sf {DSA}}}\)  becomes an additional variable to be jointly optimized with the \(\mathbf {A}_q\)  and \(\mathbf {W}\)  matrices.
According to the subarray constraint defined in (REF ), each antenna must be connected to only one RF chain at any given time. Note that each RF chain can still be connected to multiple antennas.
Therefore, \(\mathbf {S_{\sf {DSA}}}\)  can be seen as a concatenation of \(N_{\sf {T}}\)  “one-hot” vectors, each of length \(N_{\sf {RF}}\) , where a “one-hot” vector is defined as a binary vector with a Hamming weight of 1.
In classification tasks, the need to choose one of many categories is typically handled by applying the “softmax” function during training, which is replaced by a (hard) maximum during the test phase. However, we found that this approach does not lead to good results for unsupervised learning, because the sum rate measured during training can then be very different from the actual test-time sum rate.
To solve this issue, we propose a differentiable approximation during training inspired by the “Gumbel-Softmax” estimator [1]}.
Gumbel-Softmax is a technique that allows sampling from a categorical distribution during the forward pass of a neural network, by combining the reparameterization trick and smooth relaxation.
Defining \(p_{i,j}\)  the probability that antenna \(i\)  is connected to RF chains \(j\) , we can form a \(1 \times N_{\sf {RF}}\)  vector \(\mathbf {p}_i = [p_{i,1},..., p_{i,N_{\sf {RF}}}]\)  that corresponds to the probability states between antenna \(i\)  and all RF chains. Note that \(\sum _j p_{i,j} = 1 \; \forall i \in [1,N_{\sf {T}}]\)  due to constraint (REF ). The Gumbel-Softmax function applied to that vector, \(G(\mathbf {p}_i)\) , can then be defined as
\(G(\mathbf {p}_i) = \frac{\exp ((\log (\mathbf {p}_i) + g_i)/\tau )}{\sum _{j=1}^{N_{\sf {RF}}} \exp ((\log (p_{i,j}) + g_{i,j})/\tau )},\) 
