Imbalanced Learning-Enhanced Beam Codebooks towards Imbalanced User Distribution in Millimeter Wave and Terahertz Massive MIMO Systems

Abstract

Millimeter wave (mmWave) and terahertz (THz) massive MIMO architectures are pivotal in the advancement of mobile communications. These systems conventionally utilize codebooks to facilitate initial connection and to manage information transmission tasks. Traditional codebooks, however, are typically composed of numerous single-lobe beams, thus incurring substantial beam training overhead. While neural network-based approaches have been proposed to mitigate the beam training load, they sometimes fail to adequately consider the minority users dispersed across various regions. The fairness of the codebook coverage relies on addressing this problem. Therefore, we propose an imbalanced learning (IL) methodology for beam codebook construction, explicitly designed for scenarios characterized by an imbalanced user distribution. Our method begins with a pre-clustering phase, where user channels are divided into subsets based on their power response to combining vectors across distinct subareas. Then, each subset is refined by a dedicated sub-model, which contributes to the global model within each IL iteration. To facilitate the information exchange among sub-models during global updates, we introduce the focal loss mechanism. Our simulation results substantiate the efficacy of our IL framework in enhancing the performance of mmWave and THz massive MIMO systems under the conditions of imperfect channel state information and imbalanced user distribution.

1. Introduction

Millimeter wave (mmWave) and terahertz (THz) multiple-input multiple-output (MIMO) communication systems are crucial in driving the progress of the fifth generation (5G) mobile communication and future Beyond 5G. These systems incorporate large-scale antenna arrays, providing notable beamforming gains and adequate signal reception capability. Nonetheless, the implementation of fully digital transceiver architectures, necessitating an RF chain for each antenna, faces limitations due to the cost and power consumption issues linked with high-frequency mixed-signal circuits [1]. As a result, mmWave and THz systems frequently adopt either fully analog or hybrid analog/digital architectures for their combining operations [2,3]. In these configurations, transceivers utilize the networks of phase shifters. Due to the challenges associated with channel estimation and feedback, they commonly employ predefined single-lobe beamforming codebooks, such as DFT codebooks [4,5]. These classical pre-defined beamforming/beamsteering codebooks typically encompass a multitude of single-lobe beams, with each beam capable of steering the signal in a specific direction. These classical beamforming codebooks encounter two primary issues. Firstly, they involve extra beam training costs by covering all possible directions, which includes the redundant beams for beam training. Secondly, they slightly accommodate specific scenarios such as non-line-of-sight (NLOS) conditions or imbalanced user distribution.

Recent developments have focused on enhancing the adaptability of beam codebooks to user distributions in both LOS and NLOS conditions using neural network (NN)-based approaches [6,7]. Moreover, reinforcement learning has also been utilized to adapt the beam patterns based on the surrounding environment [8]. Nevertheless, the challenges posed by imbalanced or non-uniform user distribution, crucial for 5G deployment [9,10], remain largely unexplored. Typically, in imbalanced user distribution, majority users significantly outnumber the minority users. This is related to the imbalanced long-tail data challenge in machine learning, which hampers the neural network’s ability to generalize for minority classes [11]. In this case, imbalanced user distribution can lead to either overlooking or underperforming combined gain for these minority users, especially when applying NN-based methods in codebook design. Ensuring equitable codebook coverage necessitates an effective approach to address the data imbalance-related problem. Outside the realm of wireless communication, imbalanced learning (IL)-based solutions typically involve resampling the different classes in [12,13] and re-weighting featured by adjusting the loss calculated from the majority and minority classes in [14,15,16]. Furthermore, ensemble learning, a technique that aims to enhance the performance of a single classifier by integrating multiple complementary classifiers, has demonstrated its effectiveness in adapting to and generalizing from imbalanced data [17,18]. However, directly applying these strategies in mmWave and THz communications overlooks the richness of the geometric and physical attributes of the training data, i.e., the channels. Consequently, it is natural to propose an IL-based codebook specifically designed for imbalanced user distributions in mmWave and THz communication systems, which meets the following criteria. Firstly, the IL method aims to fully exploit the physical features of the channels. This is achieved by implementing pre-clustering, which is similar to a reverse process of DFT codebook-based combining. Secondly, the method inherits the superior feature extraction capability of NN-based approaches. This is accomplished through the utilization of multiple sub-models to deal with different classes of channels. Each of these sub-models contributes to the global model, which is similar to the essentials of ensemble learning.

In this paper, our work centers around the advancement of codebook design enhanced by the IL framework, particularly in scenarios where user distribution is notably imbalanced. The proposed approach commences with a pre-clustering of user channels based on the power responses of distinct combining vectors across various subregions, which enables a distinction between channels associated with majority and minority users. The classification creates multiple subsets. Moving forward, we introduced an innovative IL architecture to process these identified subsets. Each subset is then associated with a phase shift sub-model, which contributes to the global model after undergoing sub-forward and sub-backpropagation passes. To bolster information exchange across sub-models, we employ the focal loss function during global updating. Ultimately, iterative IL updating rounds for both global and sub-models allow the IL-based codebook to adapt to the imbalanced user distribution under imperfect channel conditions, which outperforms the NN-based baseline in the achievable rate metric.

Throughout this paper, vectors and matrices are represented by lower-case and upper-case boldface letters, respectively. The sets ${\mathbb{R}}^{M\times K}$ and ${\mathbb{C}}^{M\times K}$ denote the real and complex spaces of dimensions $M\times K$. The Hermitian transpose and the transpose of a matrix $\mathbf{A}$ are indicated by ${\mathbf{A}}^{†}$ and ${\mathbf{A}}^{\mathrm{T}}$, respectively. We denote an $M\times M$ identity matrix by ${\mathbf{I}}_M$ and an $M\times N$ zero matrix by ${\mathbf{0}}_{M\times N}$. A complex Gaussian random vector $\mathbf{x}$ is characterized as $\mathbf{x}\sim {\mathcal{N}}_{\mathbb{C}}(\overline{\mathbf{x}},\Sigma )$, where $\overline{\mathbf{x}}$ and $\Sigma$ represent the mean vector and the covariance matrix, respectively. The notation ${\parallel ·\parallel }_n$ denotes the n-norm of a vector.

The rest of this paper is organized as follows: Section 2 details the system model. Section 3 provides the problem description and defines the optimization objective. In Section 4, we introduce the proposed IL scheme. Section 5 describes the experimental setup and discusses the simulation results. Finally, Section 6 concludes the paper.

2. System Model

2.1. Channel Model

This study investigates a mmWave and THz massive MIMO communication system that accommodates U single-antenna users dispersed over F subregions, which utilizes a base station (BS) equipped with a uniform linear array (ULA) comprising M antennas. The single-cluster geometric channel vector between the single antenna uth user and a M-antenna BS can be written as [19]

$$\begin{array}{c}\hfill {\mathbf{h}}_u=\sqrt{\frac{M}{L}}\sum _{l=1}^L{\alpha }_{u,l}\mathbf{a}\left({\psi }_{u,l}\right),\end{array}$$

(1)

where $\sqrt{\frac{M}{L}}$ is the normalization factor, ${\alpha }_{u,l}$ denotes the complex gain of the lth path, and there are L paths in total. Also, the term ${\psi }_{u,l}$ represents the angle of arrival (AOA) of the lth path in the angle domain $\Psi \triangleq [0,\pi )$. The array response vector $\mathbf{a}(·)$ is given by the ULA geometric structure of the M-antennas as

$$\begin{array}{c}\hfill \mathbf{a}\left({\psi }_{u,l}\right)=\frac{1}{\sqrt{M}}{[1,\cdots ,e^{-j\frac{2\pi }{\lambda }d(m-1)cos\left({\psi }_{u,l}\right)},\cdots ,e^{-j\frac{2\pi }{\lambda }d(M-1)cos\left({\psi }_{u,l}\right)}]}^{\mathrm{T}}\in {\mathbb{C}}^{M\times 1},\end{array}$$

(2)

where $\lambda$ is the wavelength and d signifies the spacing between antennas. To comprehensively characterize the channel states of all U users, we introduce the concept of a channel set denoted as $\mathcal{H}$. It is formally defined as $\mathcal{H}\triangleq {\displaystyle \bigcup _{u=1}^U}{\mathbf{h}}_u$ and hence $|\mathcal{H}|=U$.

2.2. User Distribution

To streamline the representation of user distribution, we divide the channel set $\mathcal{H}$ into F subsets. This can be expressed as: $\mathcal{H}={\displaystyle \bigcup _{f=1}^F}{\mathcal{H}}_f$. Here, each user channel group ${\mathcal{H}}_f$ encompasses the channels of designated users belonging to class f ($f=1,\dots ,F$). In more detail, by employing the mean AOA of the user’s channel as the feature of a cluster [19], the channel subset ${\mathcal{H}}_f$ comprises channels with the same mean AOA values that fall within the interval ${\Psi }_f\triangleq \left[(f-1)\frac{\pi }{F},f\frac{\pi }{F}\right)$. Here, $\Psi$ is partitioned into F non-overlapping subregions. Subsequently, an imbalanced user distribution is manifested by a substantial difference in the cardinality between one subset and the others, i.e., $\exists n\in \{1,\dots ,F\}s.t.\forall f\in \{1,\dots ,F\}\mathrm{and}f\ne n,|{\mathcal{H}}_n|\gg |{\mathcal{H}}_f|$. Consequently, we can interpret ${\mathcal{H}}_n$ as ${\mathcal{H}}_{\mathrm{mjr}}$, representing the channel set of the majority of users. The channels not encompassed by ${\mathcal{H}}_{\mathrm{mjr}}$ are collectively categorized under ${\mathcal{H}}_{\mathrm{mnr}}$, which denotes the channel subsets corresponding to the minority user groups.

2.3. Uplink Transmission

In the uplink data transmission stage, the uth user sends the uplink symbol $x_u\in \mathbb{C}$ to the BS. After combining in BS, the processed signal $z_u\in \mathbb{C}$ can be represented as

$$\begin{array}{c}\hfill z_u={\mathbf{c}}_{{ϵ}_u}^{†}{\mathbf{h}}_u{x}_u+{\mathbf{c}}_{{ϵ}_u}^{†}{\mathbf{n}}_u,\end{array}$$

(3)

where the symbol $x_u$ has a mean power limitation $\mathbb{E}[|x_u{|}^2]=P_{x_u}$. Also, the noise vector ${\mathbf{n}}_u$ at the BS follows ${\mathcal{N}}_{\mathbb{C}}({\mathbf{0}}_{M\times 1},{\sigma }^2{\mathbf{I}}_M)$. Next, the combining vector ${\mathbf{c}}_{{ϵ}_u}$ is the ${ϵ}_u$th beam of the N-beam codebook $\mathcal{C}$, i.e., $\mathcal{C}\triangleq {\displaystyle \bigcup _{n=1}^N}\{{\mathbf{c}}_n\}$, $|\mathcal{C}|=N$, and ${ϵ}_u\in \{1,\cdots ,n,\cdots ,N\}$. In particular, ${\mathbf{c}}_{{ϵ}_u}$ is given as

$$\begin{array}{c}\hfill {\mathbf{c}}_{{ϵ}_u}=\frac{1}{\sqrt{M}}{\left[e^{j{\gamma }_{{ϵ}_u,1}},\dots ,e^{j{\gamma }_{{ϵ}_u,m}},\dots ,e^{j{\gamma }_{{ϵ}_u,M}}\right]}^{\mathrm{T}},\end{array}$$

(4)

where ${\gamma }_{{ϵ}_u,m}$ represents the phase shift applied at the mth BS antenna. Importantly, ${\mathbf{c}}_{{ϵ}_u}$ is only responsible for modifying the phase of the signal and does not affect the processed signal power. To this end, it is possible to optimize the configuration of ${\mathbf{c}}_{{ϵ}_u}$.

3. Problem Definition

In this paper, our primary goal is to devise a design for an N-beam codebook, referred to as $\mathcal{C}$, which can adeptly accommodate the imbalanced user distribution outlined in Section 2.2. Given the substantial generalization and adaptability attributes inherent in NN architectures, we selected this approach as the backbone network to create an adaptive codebook. Hence, we aim to reframe the construction of codebook $\mathcal{C}$ into an optimization problem based on NNs in this section.

First of all, our aim is to formulate a beamforming codebook $\mathcal{C}$ with the target of optimizing the combining gain. This gain is defined as $g_{u,{ϵ}_u}\triangleq {\left|{\mathbf{c}}_{{ϵ}_u}^{†}{\mathbf{h}}_u\right|}^2$ [6]. To this end, the optimal ${ϵ}_u$ is determined by exhaustively searching through ${\mathbf{c}}_n$ in the codebook, i.e.,

$$\begin{array}{cc}\hfill {ϵ}_u^{*}=arg\underset{n}{max}& g_{u,n}=arg\underset{n}{max}{\left|{\mathbf{c}}_n^{†}{\mathbf{h}}_u\right|}^2,\hfill \\ \hfill & \mathrm{s}.\mathrm{t}.n=1,\cdots ,N,\hfill \end{array}$$

(5)

where ${ϵ}_u^{*}$ signifies the selection of the most suitable codeword selected from the codebook $\mathcal{C}$. Following this, the optimization problem is elucidated. Specifically, the task is to maximize the combining gain across the entire dataset of user channels with the designed codebook. The expression of the superlative codebook ${\mathcal{C}}_{\mathrm{opt}}$ can be rewritten as

$$\begin{array}{cc}\hfill {\mathcal{C}}_{\mathrm{opt}}=& arg\underset{\mathcal{C}}{\max }\sum _{u=1}^U{\left|{\mathbf{c}}_{{ϵ}_u^{*}}^{†}{\mathbf{h}}_u\right|}^2,\hfill \\ & \mathrm{s}.\mathrm{t}.\left|{[{\mathbf{c}}_{{ϵ}_u^{*}}]}_m\right|=\frac{1}{\sqrt{M}},m=1,\dots ,M.\hfill \end{array}$$

(6)

The above-mentioned restriction is introduced to ensure adherence to power limitations imposed by the combining vectors.

Furthermore, a prominent challenge inherent in optimizing this problem lies in the imbalanced distribution of users, notably the substantial imbalance between the majority and minority subsets. This disproportionality leads to a faint combining gain improvement for minority user channels when employing conventional NN-based methods. This phenomenon can be attributed to the fact that the loss values computed from the minority user channels contribute insignificantly to the overall loss during batch updates in traditional NN-based methods. Consequently, the associated combining gains for the minority users from these updates remain almost the same. In the subsequent section, we will conduct an in-depth analysis of the underlying causes for the limited contribution of the minority user losses and propose an innovative IL-based solution to effectively mitigate this issue.

4. Imbalanced Learning

In this section, we employ IL to construct a more adaptive codebook in scenarios characterized by imbalanced user distribution, as compared to the NN-based approach. This adaptability arises from the independent updates of sub-models and the fair aggregation process applied to the global model. To achieve this, we introduce an initial pre-clustering method that capitalizes on the channel properties to effectively differentiate between ${\mathcal{H}}_{\mathrm{mnr}}$ and ${\mathcal{H}}_{\mathrm{mjr}}$, subsequently assigning them to separate subsets. Following this, we propose an IL framework to attain the desired beam patterns for these diverse subsets. Finally, we present a global model structure ensuring the equitable integration of the sub-models.

4.1. Model Architecture

4.1.1. Pre-Clustering Process

According to the user distribution elucidated in Section 2.2, our primary objective is to categorize users and their associated channels into $F^{\prime }$ distinct clusters. This process results in distinguishing between ${\mathcal{H}}_{\mathrm{mnr}}$ and the corresponding ${\mathcal{H}}_{\mathrm{mjr}}$. Additionally, segmenting $\mathcal{H}$ into $F^{\prime }$ clusters can be likened to partitioning $\Psi$ into $F^{\prime }$ discrete regions. This segmentation effectively characterizes each cluster by the median angle, denoted by ${\theta }_{f^{\prime }}$, of subregion $f^{\prime }$ (${\theta }_{f^{\prime }}\triangleq \frac{(2f^{\prime }-1)\pi }{2F^{\prime }})$ [20]. Consequently, by defining the median angle of each cluster as a distinctive feature, we can express the pre-clustered index $f^{\prime }\in \{1,\dots ,F^{\prime }\}$ as follows

$$\begin{array}{cc}\hfill f^{\prime }=arg\underset{i}{\max }& {\left[{\mathbf{P}}_{\mathrm{PRE}}\right]}_i=arg\underset{i}{\max }{\left[{\left|{\mathbf{C}}_{\mathrm{PRE}}^{†}{\mathbf{h}}_u\right|}^2\right]}_i,\hfill \end{array}$$

(7)

where the pre-clustering power vector ${\mathbf{P}}_{\mathrm{PRE}}={[{\mathcal{P}}_1,\dots ,{\mathcal{P}}_{f^{\prime }},\dots ,{\mathcal{P}}_{F^{\prime }}]}^{\mathrm{T}}\in {\mathbb{R}}^{F^{\prime }\times 1}$, the power response to each combining vector is defined as ${\mathcal{P}}_{f^{\prime }}\triangleq {\left|\mathbf{a}{\left({\theta }_{f^{\prime }}\right)}^{†}{\mathbf{h}}_u\right|}^2$ and the pre-clustering codebook matrix ${\mathbf{C}}_{\mathrm{PRE}}=\left[\mathbf{a}\left({\theta }_1\right),\dots ,\mathbf{a}\left({\theta }_{f^{\prime }}\right),\dots ,\mathbf{a}\left({\theta }_{F^{\prime }}\right)\right]\in {\mathbb{C}}^{M\times F^{\prime }}$. Also, the function $\mathbf{a}(·)$ signifies the array response for a given AOA. This further signifies that the channels in ${\mathcal{H}}_{f^{\prime }}$ share analogous power responses to the respective featured combining vector within the $f^{\prime }$th subregion. To implement the above equation in the form of NN, the complex matrix multiplication, power computation, and argmax can be mapped into the basic NN components. These components will be introduced in the next subsection. Then, the pre-clustered channels, which belong to the subset ${\mathcal{H}}_{f^{\prime }}$, will be stored in the $f^{\prime }$th subset register waiting to be processed by the corresponding sub-model. After partitioning $\mathcal{H}$ into $F^{\prime }$ distinct types, the next step involves deploying a model to acquire the features of ${\mathcal{H}}_{f^{\prime }}$. Specifically, this paper adopts a focal loss and cross entropy loss combined IL model improved from the original NN-based self-supervised method outlined in [7]. The details of the IL forward pass architecture will be elaborated in the subsequent parts.

Remark 1.

The pre-clustering process draws inspiration from the DFT codebook ${\mathbf{C}}_{\mathrm{DFT}}$, denoted by ${\mathbf{c}}_n^{\mathrm{DFT}}=\frac{1}{\sqrt{M}}{\left[1,\dots ,e^{j\frac{2\pi n}{M}},\dots ,e^{j\frac{2\pi n(M-1)}{M}}\right]}^{\mathrm{T}},n=0,1,\dots ,N-1$. Given that the angle domain Ψ is defined as the range $[0,\pi )$, it is divided into M subareas, each with an equal interval of $\frac{2\pi }{M}$. Consequently, when a channel’s mean AOA is congruent with a combining vector for a specific interval within Ψ, a superior power response is realized. In essence, the DFT codebook covers the entire angular space, which ensures that there is a combining vector catering to the specific channel, and thus the highest power response is attained among the pre-defined combining vectors. However, the best power response within the pre-defined combining vectors may not be optimal and requires the more adaptive NN-based method to optimize.

Informed by the exhaustive angular coverage of the DFT codebook, the initial step of our method involves calculating the power response for each channel across the subdivided angular sectors. Through this assessment, we identify the superior combining vector that results in the highest power ${\mathcal{P}}_{f^{\prime }}$ among the pre-defined combining vectors. This process facilitates the determination of a channel’s pre-clustering class $f^{\prime }$ where the channels share similar power responses.

4.1.2. Complex-Valued Fully Connected and Power Computation Layer

The layer is designed to perform complex-valued multiplicative and additive operations in the neural network backbone. In this context, the signal post-combining is represented by the resultant inner product, i.e.,

$$\begin{array}{c}\hfill \mathbf{y}={\mathbf{C}}_{f^{\prime }}^{†}\mathbf{h},\end{array}$$

(8)

where $\mathbf{y}={\left[y_1,\dots ,y_n,\dots ,y_{N_{f^{\prime }}}\right]}^{\mathrm{T}}\in {\mathbb{C}}^{N_{f^{\prime }}\times 1}$ signifies the received signal vector after combining and $\mathbf{h}\in {\mathbb{C}}^{M\times 1}$ denotes a channel from the subset ${\mathcal{H}}_{f^{\prime }}$. Also, the term ${\mathbf{C}}_{f^{\prime }}\in {\mathbb{C}}^{M\times N_{f^{\prime }}}$ denotes the $f^{\prime }$th sub-model codebook matrix. Moreover, unlike the conventional NN approaches in computer vision, the elements of the codebook matrix do not function as parameters of the fully connected layer. Instead, the codebook matrices are generated based on the real NN weights [6]. This operation is conducted through the phase-to-complex conversion, scaled by $\frac{1}{\sqrt{M}}$, which is given as

$$\begin{array}{c}\hfill {\mathbf{C}}_{f^{\prime }}=\frac{1}{\sqrt{M}}\left(cos\left({\mathbf{\Gamma }}_{f^{\prime }}\right)+j*sin\left({\mathbf{\Gamma }}_{f^{\prime }}\right)\right),\end{array}$$

(9)

where ${\mathbf{\Gamma }}_{f^{\prime }}=\left[{\gamma }_1,\cdots ,{\gamma }_n,\cdots ,{\gamma }_{N_{f^{\prime }}}\right]\in {\mathbb{R}}^{M\times N_{f^{\prime }}}$ denotes the phase shift matrix and ${\gamma }_n=[{\gamma }_{1,n},...,$${\gamma }_{M,n}],$$n\in \left\{1,\dots ,N_{f^{\prime }}\right\}$ represents a phase shift vector.

Following this, the power received from each combining vector is determined by evaluating the square modulus of every complex value in the received signal vector. This leads to the formulation of the received signal power $\mathbf{p}\in {\mathbb{R}}^{N_{f^{\prime }}\times 1}$, which is defined as

$$\begin{array}{c}\hfill \mathbf{p}={\left[p_1,\cdots ,p_n,\cdots ,p_{N_{f^{\prime }}}\right]}^{\mathrm{T}}={\left[|y_1{|}^2,\cdots ,|y_n{|}^2,\cdots ,{|y_{N_{f^{\prime }}}|}^2\right]}^{\mathrm{T}}.\end{array}$$

(10)

4.1.3. Softmax and Argmax Layer

This layer commences with a softmax operation to estimate the alignment ’probability’ between a combining vector and the current channel $\mathbf{h}$, which is indicated by the received power $\mathbf{p}$. Accordingly, the softmax vector, represented by $\mathbf{s}={[s_1,\dots ,s_n,\dots ,s_{N_{f^{\prime }}}]}^T\in {\mathbb{R}}^{N_{f^{\prime }}\times 1}$, is formulated as

$$\begin{array}{c}\hfill s_n=\frac{e^{p_n}}{{\displaystyle \sum _{i=1}^{N_{f^{\prime }}}}{e}^{p_i}}.\end{array}$$

(11)

On the other hand, the argmax layer is tasked with conducting the one-hot vector, denoted by $\mathbf{o}={[o_1,\cdots ,o_n,\cdots ,o_{N_{f^{\prime }}}]}^{\mathrm{T}}\in {\mathbb{R}}^{N_{f^{\prime }}\times 1}$, and it is given as

$$\begin{array}{c}\hfill o_n=\left\{\begin{array}{cc}\hfill & 1,n=arg\underset{i}{\max }{[\mathbf{s}]}_i,i\in \{1,\dots ,N_{f^{\prime }}\},\hfill \\ \hfill & 0,\mathrm{otherwise},\hfill \end{array}\right.,\end{array}$$

(12)

where the $argmax$ operation is to identify the position of the maximum value within the softmax vector. Additionally, the self-generated label vector $\mathbf{o}$ plays a crucial role in adjusting and updating the phase shifts in the codebook. This adaptation is achieved through the utilization of the cross-entropy loss during sub-updating and the focal loss during global updating. Also, it is worth noting that a similar $argmax$ operation is also required in the pre-clustering process to determine the class index $f^{\prime }$ associated with the optimal power response.

Transitioning to the process of learning the IL-based codebook, the subsequent subsection will delve into the details of sub-backpropagation and the comprehensive IL solution.

4.2. Learning the IL-Based Codebooks

After establishing the architecture for the forward pass of the sub-model, the subsequent phase entails the application of sub-backpropagation to enhance the performance of the sub-model ${\Gamma }_{f^{\prime }}$. Following the sub-backpropagation, the sub-models are aggregated to assess the overall performance of the global model, which then performs the global updating through the focal loss. Ultimately, the global model is partitioned into sub-models to initiate the next round of IL.

4.2.1. Sub-Backpropagation

The corresponding sub-model undergoes a series of forward-backward updating iterations based on ${\mathcal{H}}_{f^{\prime }}$. Within each iteration, as delineated in the preceding section, the forward processes obtain the softmax vectors $\mathbf{s}$ and one-hot labels $\mathbf{o}$ for a batch of channel vectors. In each of these iterations, the efficacy of the codebook is evaluated by quantifying the divergence between the obtained softmax vector $\mathbf{s}$ and the target one-hot vector $\mathbf{o}$, utilizing the cross-entropy loss function, i.e.,

$$\begin{array}{c}\hfill {\mathcal{L}}_B^{\mathrm{CE}}=\sum _{b=1}^B{\mathcal{L}}_b^{\mathrm{CE}}=-\sum _{b=1}^B\sum _{n=1}^{N_{f^{\prime }}}{o}_{nb}log{s}_{nb},\end{array}$$

(13)

where $s_{nb}$ and $o_{nb}$ represent the nth element of the bth tuple ${\mathbf{s}}_b$, ${\mathbf{o}}_b$ in a batch with the size of B. Moreover, channels for a single batch are indiscriminately chosen from both ${\mathcal{H}}_{\mathrm{mjr}}$ and ${\mathcal{H}}_{\mathrm{mnr}}$, which constitute the comprehensive dataset $\mathcal{H}$. In contrast, in the IL-based batch, the batch exclusively comprises channels from either ${\mathcal{H}}_{\mathrm{mjr}}$ or ${\mathcal{H}}_{\mathrm{mnr}}$.

Particularly, the aforementioned loss function considers the one-hot encoded label $\mathbf{o}$ as the target distribution for the model. Its objective is to minimize this value through the adjustment of phase shift vectors ${\gamma }_n\in {\mathbf{\Gamma }}_{f^{\prime }}$ with the aim of minimizing the discrepancy between $\mathbf{s}$ and $\mathbf{o}$ to the utmost degree. Subsequently, the backpropagation gradient of phase shift vectors $\mathbf{g}\left({\gamma }_n\right)$ is computed through the chain rule, i.e., [7]

$$\begin{array}{c}\hfill \mathbf{g}\left({\gamma }_n\right)={\left(\frac{\partial {\mathcal{L}}_B}{\partial {\gamma }_n}\right)}^{\mathrm{T}}=\sum _{b=1}^B{\left(\frac{\partial {\mathcal{L}}_b}{\partial {\mathbf{s}}_b}\right)}^{\mathrm{T}}·\frac{\partial {\mathbf{s}}_b}{\partial {\mathbf{p}}_b}·\frac{\partial {\mathbf{p}}_b}{\partial {\mathbf{y}}_b}·\frac{\partial {\mathbf{y}}_b}{\partial {\gamma }_n}.\end{array}$$

(14)

Hence, the stochastic gradient descent (SGD) based method can be employed to update the sub-model phase matrix ${\mathbf{\Gamma }}_{f^{\prime }}$, which is given as

$$\begin{array}{c}\hfill {{\gamma }_n}_{\mathrm{new}}={{\gamma }_n}_{\mathrm{current}}-\beta ·\mathbf{g}{\left({\gamma }_n\right)}^{\mathrm{T}},\end{array}$$

(15)

where $\beta$ denotes the learning rate of the SGD-based process. In conclusion, the pre-clustering process and sub-model updating are demonstrated in Figure 1. After the sub-model updates, the next stage will introduce the process of global aggregation and updating.

4.2.2. Global Updating

After the sub-model updating epochs, the phase shift matrices ${\mathbf{\Gamma }}_{f^{\prime }}$ are fused into the global model. As previously explained, ${\mathbf{\Gamma }}_{f^{\prime }}$ consists of $N_{f^{\prime }}$ distinct phase shift vectors, each corresponding to the number of combining vectors within the correlated codebook. Our objective is to evenly integrate the model across the various $N_{f^{\prime }}$ distributions, thereby addressing the challenge of imbalanced combining gain elevation in the NN-based approach. Driven by the specified objective, the global model is constructed by concatenating the phase shift matrices of varying dimensions $N_{f^{\prime }}$ to create a unified model, which is expressed by

$$\begin{array}{c}\hfill \mathbf{\Gamma }=\left[{\mathbf{\Gamma }}_1,\dots ,{\mathbf{\Gamma }}_{f^{\prime }},\dots ,{\mathbf{\Gamma }}_{F^{\prime }}\right],\end{array}$$

(16)

where $\mathbf{\Gamma }\in {\mathbb{R}}^{M\times N}$ ($N={\displaystyle \sum _{a=1}^{F^{\prime }}}{N}_a)$ represents the global phase shift matrix model and $\mathbf{C}$ denotes the corresponding global codebook. The above equation ensures the specialization of different sub-models.

Nonetheless, the process of concatenation introduces the invalid exchange of information among different sub-models, potentially sub-optimal in nature. The information exchange implies that a small proportion of the channels from a subset ${\mathcal{H}}_{f^{\prime }}$ tend to have better power responses on combining vectors in the adjacent sub-model after global updating, instead of its own sub-model ${\mathbf{\Gamma }}_{f^{\prime }}$. Consequently, to adapt the global user distribution within $\mathcal{H}$, the global model should undergo updates akin to the sub-model. However, the cross-entropy loss function is supposed to be replaced with a more adaptive focal loss function in [21]. The adaptive method is supposed to automatically rectify the imbalance in loss calculations, stemming from the imbalance of quantity between two types of channels: those seeking to change their assigned sub-models and those already well adapted to their sub-models. Specifically, the focal loss mechanism notes the lower received power (reflected on the softmax probabilities) of the exchanged channels, compared with the received power of a sub-model’s originally well-adapted channels. To this end, the mechanism will introduce a more balanced loss function by improving the proportion of the losses of the low received power channels. Therefore, the focal loss function is utilized and given as [21]

$$\begin{array}{c}\hfill {\mathcal{L}}_B^{\mathrm{FL}}=-\sum _{b=1}^B\sum _{n=1}^N{(1-s_{nb})}^{\zeta }{o}_{nb}log{s}_{nb},\end{array}$$

(17)

where the modulating factor ${(1-s_{nb})}^{\zeta }$ is added to the cross-entropy loss, with a tunable focusing parameter $\zeta ⩾0$. In scenarios where the received power is relatively high (i.e., $s_{nb}\to 1$), the modulating factor approaches zero. Conversely, for channels with lower received power, the modulating factor tends towards unity. As a result, when contrasted with the standard cross-entropy loss, the focal loss maintains its effectiveness for channels with lower received power while reducing the loss for those with higher received power. In essence, this adjustment amplifies the contribution of channels with lower receiver power that seeks to alter the serving sub-model within the loss function.

To initiate the subsequent IL iteration, $\mathbf{\Gamma }$ is separated at the columns corresponding to the prior concatenation, reverting to their pre-fusion dimensions. This step can be considered the reverse operation of Equation (16). To date, Algorithm 1 has summarized the overall IL process.

Algorithm 1 Imbalanced Learning.

1: Pre-cluster $\mathcal{H}$ and return $F^{\prime }$ different subsets 2: Initialize $F^{\prime }$ sub-models 3: for each IL round do 4:     for each subset ${\mathcal{H}}_{f^{\prime }}$ and the selected sub-model ${\mathbf{\Gamma }}_{f^{\prime }}$ do 5:       for each sub-model updating epoch do 6:         for each batch with size of B in ${\mathcal{H}}_{f^{\prime }}$ do 7:           for each channel in a batch indexed by $b\in \{1,\dots ,B\}$ do 8:              Execute the sub-forward pass and obtain ${\mathbf{s}}_b$, ${\mathbf{o}}_b$ 9:              Execute the sub-backpropagation pass:               ${\mathcal{L}}_B^{\mathrm{CE}}\leftarrow {\mathcal{L}}_B^{\mathrm{CE}}-{\displaystyle \sum _{n=1}^{N_{f^{\prime }}}}{o}_{nb}log{s}_{nb}$ 10:           end for 11:           for each $n\in \{1,\dots ,N_{f^{\prime }}\}$ do 12:              Update the selected sub-model ${\mathbf{\Gamma }}_{f^{\prime }}$ in the sub-backpropagation pass:               ${{\gamma }_n}_{\mathrm{new}}\leftarrow {{\gamma }_n}_{\mathrm{current}}-\beta ·\frac{\partial {\mathcal{L}}_B^{\mathrm{CE}}}{\partial {\gamma }_n}$ 13:           end for 14:         end for 15:      end for 16:    end for 17:    Obtain the global model $\mathbf{\Gamma }=\left[{\mathbf{\Gamma }}_1,\dots ,{\mathbf{\Gamma }}_{f^{\prime }},\dots ,{\mathbf{\Gamma }}_{F^{\prime }}\right]$ 18     for each global updating epoch do 19:      for each batch with size of B in $\mathcal{H}$ do 20:         for each channel in a batch indexed by $b\in \{1,\dots ,B\}$ do 21:           Execute the global forward pass and obtain ${\mathbf{s}}_b$, ${\mathbf{o}}_b$ 22:           Execute the global backpropagation pass:           ${\mathcal{L}}_B^{\mathrm{FL}}\leftarrow {\mathcal{L}}_B^{\mathrm{FL}}-{\displaystyle \sum _{n=1}^N}{(1-s_{nb})}^{\zeta }{o}_{nb}log{s}_{nb}$, $N={\displaystyle \sum _{a=1}^{F^{\prime }}}{N}_a$ 23:         end for 24:         for each $n\in \{1,\dots ,N\}$ do 25:           Update the global model $\mathbf{\Gamma }$ in the global backpropagation pass:            ${{\gamma }_n}_{\mathrm{new}}\leftarrow {{\gamma }_n}_{\mathrm{current}}-\beta ·\frac{\partial {\mathcal{L}}_B^{\mathrm{FL}}}{\partial {\gamma }_n}$ 26:         end for 27:      end for 28:    end for 29:    Execute the reverse operation of $\mathbf{\Gamma }=\left[{\mathbf{\Gamma }}_1,\dots ,{\mathbf{\Gamma }}_{f^{\prime }},\dots ,{\mathbf{\Gamma }}_{F^{\prime }}\right]$ and get the global-updated sub-models 30: end for

Remark 2.

While the loss function in the NN-based approach effectively quantifies the disparity between the current model and the desired response for users equally distributed, its primary limitation becomes apparent when applied in a batch-based context. Even though only the best-performing combining vectors associated with the current channel will be updated based on the one hot label, it is quite common for the same combining vector to serve both the majority and minority during a single batch. For instance, it arises when the codebook size is relatively small, and the previous batch update only contains the majority user channels and the current batch update contains minority user channels. To this end, in cases where there is an imbalance in the number of users across different regions, the loss associated with channels from ${\mathcal{H}}_{\mathrm{mnr}}$ might be unfairly averaged and consequently downplayed.

To illustrate this, it is reasonable to consider a NN-based updating process for a phase shift matrix $\mathbf{\Gamma }\in {\mathbb{R}}^{M\times N}$ based on the global dataset $\mathcal{H}$. Let the number of channels from ${\mathcal{H}}_{\mathrm{mjr}}$ in a batch be denoted by $B_{\mathrm{mjr}}$ and the number of channels from ${\mathcal{H}}_{\mathrm{mnr}}$ as $B_{\mathrm{mnr}}$, with $B_{\mathrm{mjr}}>B_{\mathrm{mnr}}$ and $B_{\mathrm{mjr}}+B_{\mathrm{mnr}}=B$. In such scenarios, the loss function for NN with a batch size of B can be represented as follows

$$\begin{array}{c}\hfill {\mathcal{L}}_B=-\left(\sum _{b=1}^{B_{\mathrm{mnr}}}\sum _{n=1}^N{o}_{nb}log{s}_{nb}+\sum _{b=1}^{B_{\mathrm{mjr}}}\sum _{n=1}^N{o}_{nb}log{s}_{nb}\right).\end{array}$$

(18)

If $B_{\mathrm{mjr}}\gg B_{\mathrm{mnr}}$, the equation above is approximated as

$$\begin{array}{c}\hfill {\mathcal{L}}_B\approx {\mathcal{L}}_{B_{\mathrm{mjr}}}=-\sum _{b=1}^{B_{\mathrm{mjr}}}\sum _{n=1}^N{o}_{nb}log{s}_{nb}.\end{array}$$

(19)

This insight reveals a tendency of the NN-based method to underestimate the calculated loss in specific regions when the size of subsets in those regions is significantly smaller than the region with the predominant users. In the IL-based approaches, the shortcomings of the loss function are effectively mitigated. During the update phase of the phase shift matrix in NN, the enhancement in combining gain for minority users can be measured by the contribution of ${\mathcal{H}}_{\mathrm{minor}}$ to the overall loss function, expressed as

$$\begin{array}{c}\hfill {\mathcal{Q}}_{\mathrm{NN}}\triangleq \frac{{\mathcal{L}}_{B_{\mathrm{mnr}}}}{{\mathcal{L}}_B}=\frac{{\displaystyle \sum _{b=1}^{B_{\mathrm{mnr}}}}{\displaystyle \sum _{n=1}^N}{o}_{nb}log{s}_{nb}}{{\displaystyle \sum _{b=1}^B}{\displaystyle \sum _{n=1}^N}{o}_{nb}log{s}_{nb}},\end{array}$$

(20)

where $\mathcal{Q}$ denotes the fraction of the loss derived from ${\mathcal{H}}_{\mathrm{mnr}}$ in the cumulative loss. For a comparative perspective, by assuming the update processes of all sub-models for ${\mathcal{H}}_{\mathrm{mnr}}$ are considered, the proportional contribution ${\mathcal{Q}}_{\mathrm{IL}}$ in the IL framework can be given as

$$\begin{array}{c}\hfill {\mathcal{Q}}_{\mathrm{IL}}=\frac{{\mathcal{L}}_{B_{\mathrm{mnr}}}}{{\mathcal{L}}_{B_{\mathrm{mnr}}}}=\frac{{\displaystyle \sum _{b=1}^{B_{\mathrm{mnr}}}}{\displaystyle \sum _{n=1}^N}{o}_{nb}log{s}_{nb}}{{\displaystyle \sum _{b=1}^{B_{\mathrm{mnr}}}}{\displaystyle \sum _{n=1}^N}{o}_{nb}log{s}_{nb}},\end{array}$$

(21)

where ${\mathcal{Q}}_{\mathrm{IL}}$ exceeds ${\mathcal{Q}}_{\mathrm{NN}}$ is evident in the stipulated scenario. Consequently, from Equations (20) and (21), it is inferred that minority users contribute more significantly to the loss function within the proposed framework. This ensures a more pronounced combining gain enhancement for minority users in comparison to the NN approach.

5. Simulation Result

In this section, we begin by providing an in-depth exposition of the dataset generation, model training, and testing configurations for the proposed algorithms. The simulation was conducted in a setting with imperfect channel conditions and a notably imbalanced user distribution.

5.1. Scenarios, Datasets, and Training Parameters

To assess the validity of the IL method in addressing the challenges posed by the imbalance within the channel set $\mathcal{H}$, we create a scenario featuring 10 non-overlapping and equally spaced subregions. The numbers of pre-clustering subsets $F^{\prime }$ in pre-clustering, BS antennas $M$ and paths $L$ are set as 5, 64 and 5, respectively. The AOA within each subregion follows a Laplacian distribution, all sharing the same mean angle and an angle spread of 10 degrees. More precisely, the mean AOAs in subregion n fall within the interval $[\frac{\pi }{10}(n-1),\dots ,\frac{\pi }{10}n)$. We generate a total of 10,000 realizations of $\mathbf{h}$ in the 10th subregion, representing the majority of users. In contrast, the minority group is represented by 250 realizations of $\mathbf{h}$, with subregions 1,3,4,5 containing 100, 50, 50, and 50 realizations, respectively. The average achievable rate for all users is denoted by $\frac{1}{U}{\displaystyle \sum _{u=1}^U}{log}_2\left(1+\frac{P_{x_u}}{{\sigma }^2}{\left|{\mathbf{c}}_{{ϵ}_u^{*}}^{†}{\mathbf{h}}_u\right|}^2\right)$ and the term $P_{x_u}/{\sigma }^2$ is the received SNR. For assessing the performance under imperfect channel conditions, we introduce synthetic additive white Gaussian noise to each channel’s data and define ${\mathrm{SNR}}_{\mathrm{tr}}=20{log}_{10}\frac{{[|{\mathbf{h}}_u]}_i{|}^2}{{{\sigma }^{\prime }}^2}$ [20]. For simulations, ${\mathrm{SNR}}_{\mathrm{tr}}$ is set as 10 dB. To assess the efficacy of our proposed method, we perform a comparative analysis. Our method is compared with the N-beam DFT codebook ${\mathbf{c}}_n^{\mathrm{DFT}}=\frac{1}{\sqrt{M}}{[1,\dots ,e^{j\frac{2\pi n}{M}},\dots ,e^{j\frac{2\pi n(M-1)}{M}}]}^{\mathrm{T}},n=0,1,\dots ,N-1$ and with the self-supervised NN solution outlined in prior work [7]. In addition, we define the upper bound equal-gain combining vector as ${\mathbf{c}}_{u,\mathrm{egc}}=\frac{1}{\sqrt{M}}{\left[e^{j\angle h_{u,1}},e^{j\angle h_{u,2}},\cdots ,e^{j\angle h_{u,M}}\right]}^{\mathrm{T}}$. This upper bound vector is defined based on the corresponding received power for each user, given as $p={\left|{\mathbf{c}}_{u,\mathrm{egc}}^{†}{\mathbf{h}}_u\right|}^2=\frac{1}{M}||{\mathbf{h}}_u{||}_1^2$ [22]. To simplify the IL model training, we maintain uniform codebook sizes across all sub-models. In a given training instance, the batch sizes during the sub-model and global updates are configured to 10 and 1000, respectively. Similarly, the epochs are set to 20 for sub-model updating and 5 for global updating. Moreover, the learning rate, the focal loss focusing parameter $\zeta$, and the validation rates are set as 0.001, 2, and 0.1, respectively.

5.2. Performance Analysis

5.2.1. Pre-Clustering Result

As delineated in Section 5.1, the constructed dataset comprises 10,000 realizations for region 10, while all other regions consist of 250 realizations. After pre-clustering, the channel sets ${\mathcal{H}}_1$, ${\mathcal{H}}_2$, ${\mathcal{H}}_3$, and ${\mathcal{H}}_5$ encompass 125, 49, 95, and 9981 channels, respectively. These outcomes suggest that certain non-ideal conditions, such as constrained subregion segmentation, imperfect channel state information, and angle spacing, may impede precise clustering. However, the distinction between ${\mathcal{H}}_{\mathrm{mjr}}$ and ${\mathcal{H}}_{\mathrm{mnr}}$ remains generally discernible.

5.2.2. Achievable Rate and Model Convergence

Figure 2a depicts the average achievable rate against the codebook size at received SNRs of 0 dB and 5 dB. The equal-gain receiver serves as the performance upper bound. Observations from the results indicate that increased codebook size leads to augmented gains and achievable rates, which implies enhanced feature extraction during updates. Notably, saturation in the achievable rate is observed at a 64-beam codebook size, which suggests marginal gains when transitioning from 32 to 64 beams. In terms of the proposed IL method’s efficacy, it surpasses the NN-based approach in the achievable rate for all cases except the 64-beam codebook. Specifically, with smaller codebooks (sizes of 8, 12, or 16 beams), the IL method demonstrates superior adaptability to ${\mathcal{H}}_{\mathrm{mnr}}$ and the imbalanced user distribution. For codebooks exceeding 16 beams, the IL method’s performance is commensurate with that of the 64-beam DFT codebook, whereas the NN-based codebooks lag behind. Moreover, Figure 2b depicts the validation losses for the 16-beam IL codebook over the IL rounds, each comprising sub-model and global updating cycles. It is close to saturation in the 8th round.

5.2.3. Beam Pattern

Figure 3 showcases the beam patterns emanating from six learned codebooks of sizes 20, 32, and 64, which demonstrates the learning outcomes for both the proposed and NN-based methods. Notably, unlike the DFT codebook, these patterns reveal that beams do not span the entire azimuth plane but are instead oriented toward user locations. However, the NN codebook tends to overlook minority channel groups in region 1, and the magnitude of beams for user channels in regions 3, 4, and 5 are not significantly enhanced after updating. This is particularly evident in the 64-beam NN codebook, where most beams cater to the majority user channels. Consequently, it fails to capture the desired beam pattern, which corroborates the analysis presented in Remark 2. In contrast, the IL-based method flexibly adapt to the minority user channels in multipath scenarios using multi-lobe beams or precisely aligns with the minority users’ orientation with a magnitude exceeding 0.7.

6. Conclusions

In this paper, our focus was on the design of codebooks aided by the IL framework in the mmWave and THz massive MIMO communication systems with imbalanced user distribution. The primary contributions of our work to the field of machine learning-enhanced mmWave and THz massive MIMO communication lie in revealing the negative impact of imbalanced datasets in practical communication on the overall performance of NN-based codebook design and proposing an IL-based method to handle the imbalanced samples. We began by pre-clustering $\mathcal{H}$ based on the power responses of distinct combining vectors across various subregions, which allows us to differentiate between channels associated with majority and minority users. This classification resulted in multiple subsets. Subsequently, we introduced an innovative IL architecture to process these $F^{\prime }$ subsets. Each subset was associated with a phase shift sub-model ${\mathbf{\Gamma }}_{f^{\prime }}$, which contributed to the global model after undergoing sub-forward and sub-backpropagation passes. To enhance the information exchange across different sub-models, we utilized the focal loss function, which improved the balance during global updating. Compared to the NN-based method in previous work, the IL process ensured combining gain improvement during updates by augmenting the minority user channel’s contribution to the overall loss function in each batch. Through iterative IL updating rounds for both global and sub-models, our IL-based approach effectively adapted to the imbalanced user distribution under imperfect channel conditions, surpassing the performance of the NN-based baseline in terms of achievable rate. Moreover, the depicted beam patterns further intuitively confirmed the method’s effectiveness in accommodating minority user channels. In future work, we plan to broaden the scope of our methodology to encompass hybrid beamforming techniques. This expansion will involve integrating additional modules in the backbone network for designing baseband precoders. Additionally, to test the validation of the proposed IL framework in a more practical scenario, we will consider the potential hardware impairments on the antenna array geometric that are prevalent in the real-world mmWave and THz massive MIMO systems.