Spatial attention and quantization-based contrastive learning framework for mmWave massive MIMO beam training

Abstract

Deep learning (DL)-based beam training schemes have been exploited to improve spectral efficiency with fast optimal beam selection for millimeter-wave (mmWave) massive multiple-input multiple-output (MIMO) systems. To achieve high prediction accuracy, these DL models rely on training with a tremendous amount of labeled environmental measurements, such as mmWave channel state information (CSI). However, demanding a large volume of ground truth labels for beam training is inefficient and infeasible due to the high labeling cost and the requirement for expertise in practical mmWave massive MIMO systems. Meanwhile, a complex environment incurs critical performance degradation in the continuous output of beam training. In this paper, we propose a novel contrastive learning framework, named self-enhanced quantized phase-based transformer network (SE-QPTNet), for reliable beam training with only a small fraction of the labeled CSI dataset. We first develop a quantized phase-based transformer network (QPTNet) with a hierarchical structure to explore the essential features from frequency and spatial views and quantize the environmental components with a latent beam codebook to achieve robust representation. Next, we design the SE-QPTNet including self-enhanced pre-training and supervised beam training. SE-QPTNet pre-trains by the contrastive information of the target user and others with the unlabeled CSI, and then, it is utilized as the initialization to fine-tune with a reduced volume of labeled CSI. Finally, the experimental results show that the proposed framework improves beam prediction accuracy and data rates with 5% labeled data compared to existing solutions. Our proposed framework further enhances flexibility and breaks the limitation of the quantity of label information for practical beam training.