KANsformer for Scalable Beamforming

Deep learning (DL) has revolutionized a wide range of application fields and achieved unprecedented success in tasks such as image recognition and natural language processing. Its ability to automatically extract high-level features from raw data enables deep neural networks to outperform traditional machine learning methods in complex problem domains. Recently, the DL-enabled designs for wireless networks have emerged as a hot research topic[1]. Some researchers attempted to apply multi-layer perceptrons (MLP) [2], convolutional neural networks (CNN) [3] and graph neural networks (GNN) [4] to deal with the power allocation and signal processing problems in wireless networks. Overall, the DL models can be trained to achieve close performance to traditional convex optimization (CVXopt)-based approaches but with a much faster inference speed. How to further improve the learning performance remains an open issue for DL-enabled wireless optimization.

Typically, task-oriented DL requires dedicated models for wireless networks. One promising way is to follow the “encoder-decoder” paradigm, where the encoder extracts features over the wireless networks while the decoder maps the extracted features to desired transmit design. Recent works have paid a great attention to the design of encoder. Particularly, the GNN shows good scalability and generalization performance by exploiting the graph topology of wireless networks [5]. In [6, 7, 8], the GNN was adopted as the encoder to develop the solution approaches for energy efficiency (EE) maximization, sum-rate maximization and max-min rate, respectively, for multi-user multi-input-single-output (MISO) networks, all of which were scalable to the number of users. In [9], a GNN based model was trained via unsupervised learning to solve the outage-constrained EE maximization problem. The GNNs in [6, 7, 8, 9] all leveraged the multi-head attention mechanism also known as the graph attention networks (GAT) to enhance the feature extraction, especially for inter-user interference. Similarly, the transformer is also built upon the self-attention mechanism and is popular for its highly predictive performance [10]. In [11], a transformer and a weighted A^∗ based algorithm were proposed to plan the unmanned aerial vehicle trajectory for age-of-information minimization, which outperformed traditional algorithms numerically. However, the above works on wireless networks all adopted MLP as decoder. Recently, Kolmogorov-Arnold network (KAN) has been proposed as a promising alternative to MLP with superior performance in terms of accuracy and interpretability [12].

To the best of our knowledge, the integration of transformer and KAN has not been applied to beamforming design. In this paper, we formulate the classic EE maximization problem for MISO networks [13]. We then propose an approach integrating transformer and KAN termed KANsformer, which utilizes the multi-head self-attention mechanism to extract hidden features among interference links and KAN to map the extracted features to the desired beamforming design. The KANsformer is trained via unsupervised learning and a scale function guarantees feasible solution. Via parameter sharing, the KANsformer is scalable to the number of users. Numerical results indicate that the KANsformer outperforms existing DL models and approaches the CVXopt-based solution accuracy with millisecond-level inference time. The major performance gain is contributed by the KAN via ablation experiment. Besides, we validate the scalability of KANsformer, which can be further enhanced by transfer learning at the expense of little training cost.

The rest of this paper is organized as follows. Section II gives the system model and problem formulation. Section III presents the structure of KANsformer. Section IV provides numerical results. Finally, Section V concludes the paper with future research directions.

Consider a downlink MISO network, where one $N_{\rm T}$-antenna transmitter intends to serve $K$ single-antenna mobile users (MUs) over a common spectral band. We use $\mathcal{K}\triangleq\{1,2,...,K\}$ to denote the index set of the MUs.

Denote the symbol for the $k$-th MU and the corresponding beamforming vector as $s_{k}$ and ${{{\bf{w}}_{k}}}\in\mathbb{C}^{N_{\rm T}}$, respectively. The received signal at the $k$-th MU is given by

where ${{{\bf{h}}_{k}}}\in{\mathbb{C}^{{N_{\rm{T}}}}}$ denotes the channel state information (CSI) of the $k$-th transmitter-MU link, and $n_{k}\sim\mathcal{CN}({0,{\sigma_{k}^{2}}})$ denotes the additive white Gaussian noise (AWGN) at the $k$-th MU. Without loss of generality, it is assumed that ${{\mathbb{E}}}\{{{{|{s_{k}}|}^{2}}}\}=1$ ($\forall k\in\mathcal{K}$). Then, the achievable rate at the $k$-th MU is expressed as

where $\{{\bf w}_{i}\}$ denotes the set of all admissible beamforming vectors. The weighted EE for the considered system is expressed as

where $\alpha_{k}$ is a preassigned weight for the $k$-th MU and ${P_{\rm C}}$ denotes the constant power consumption introduced by circuit modules.

Our goal is to maximize the EE of the considered network, which is mathematically formulated as an optimization problem:

where ${P_{\rm max}}$ denotes the power budget of the transmitter.

The problem (4) can be efficiently solved by existing CVX techniques but without close-form solution. By treating the CVXopt-based algorithm as a “black box”, it maps CSI to beamforming vectors via iterative computations. Such a mapping can also be regarded as a “function” represented by $\Pi(\cdot):{\mathbb{C}}^{N_{\rm T}\times K}\rightarrow{\mathbb{C}}^{N_{\rm T}\times K}$. Following the universal approximation theorem, we intend to utilize neural networks to solve the problem (4).

The main idea of the proposed KANsformer is to realize the mapping, i.e., $\Pi(\cdot)$, from $\{{\bf h}_{i}\}$ to $\{{\bf w}_{i}\}$ such that ${\rm EE}({\{{{{\bf{w}}_{i}}}\}})$ is close to ${\rm EE}({\{{{{\bf{w}}^{\star}_{i}}}\}})$. Particularly, we utilize unsupervised learning to train the KANsformer to alleviate the burden on collecting labelled training set. Denote $\bm{\theta}$ as learnable parameters of the KANsformer, the loss function to update the learnable parameters is given by

Note that via off-line training based on historical statistics, the KANsformer can derive the solution instantaneously at a low computational complexity instead of complex iterative calculation by the mathematical optimization approaches.

The structure of the KANsformer is illustrated in Fig. 1, which includes four modules: pre-processing module, transformer encoder module, KAN decoder module and post-processing module. The detailed processes in each module are described as follows.

This section provides numerical results to evaluate the proposed KANsformer in terms of generalization performance, transfer learning and ablation experiment under the following settings.

We have presented a DL model (i.e., the KANsformer shown in Fig. 1) with the transformer and KAN used in the encoder-decoder structure, respectively, for solving the beamforming design problem (cf. (4)).

Numerical results showed that the KANsformer outperforms the existing DL models in terms of both the performance accuracy and the inference time consumed. Furthermore, we would like to emphasize that, in response to the given input CSI $\{{\bf h}_{k}\}$, the KANsformer can yield the beamforming vector output $\{{\bf w}_{k}\}$, with the elapsed inference time almost fixed (thus insensitive to the problem size) and tremendously lower than the problem-size dependent running time required by the CVXopt-based approach; the performance accuracy of the former is quite close to that the latter (treated as the optimum). These results also motivate further development of more powerful encoders and decoders dedicated to wireless communication systems.