A Multi-Task Foundation Model for Wireless Channel Representation Using Contrastive and Masked Autoencoder Learning

TL;DR

ContraWiMAE tackles the challenge of learning transferable wireless channel representations by combining masked reconstruction with a wireless-specific masked contrastive objective in an asymmetric encoder-decoder transformer. The approach yields data-efficient, robust embeddings that transfer across unseen scenarios and frequencies, enabling strong performance on cross-frequency beam selection, LoS detection, and channel estimation without heavy supervision. Key contributions include the wireless-inspired contrastive mechanism, an efficient asymmetric architecture, and comprehensive analysis of complexity, robustness, and design choices, supported by large-scale DeepMIMO-based pretraining and unseen-scene evaluation. The work demonstrates practical impact by achieving competitive results with lightweight downstream models and releasing model weights and training pipelines for reproducibility and future research.

Abstract

Current applications of self-supervised learning to wireless channel representation often borrow paradigms developed for text and image processing, without fully addressing the unique characteristics and constraints of wireless communications. To bridge this gap, we introduce ContraWiMAE, Wireless Contrastive Masked Autoencoder, a transformer-based foundation model that unifies masked reconstruction and masked contrastive learning for wireless channel representation. Our key innovation is a new wireless-inspired contrastive objective that exploits the inherent characteristics of wireless environment, including noise, fading, and partial observability, as natural augmentation. Through extensive evaluation on unseen scenarios and conditions, we demonstrate our method's effectiveness in multiple downstream tasks, including cross-frequency beam selection, line-of-sight detection, and channel estimation. ContraWiMAE exhibits superior linear separability and adaptability in diverse wireless environments, demonstrating exceptional data efficiency and competitive performance compared with supervised baselines under challenging conditions. Comparative evaluations against a state-of-the-art wireless channel foundation model confirm the superior performance and data efficiency of our approach, highlighting its potential as a powerful baseline for future research in self-supervised wireless channel representation learning. To foster further work in this direction, we release the model weights and training pipeline for ContraWiMAE.

Figures (7)

• Figure 1: Illustration of the ContraWiMAE blocks, detailing preprocessing steps, encoder-decoder processing, and the contrastive head.
• Figure 2: Channel estimation on unseen scenarios
• Figure 3: Top-1 beam selection accuracy for a codebook size of 32 achieved by ContraWiMAE vs. baselines followed by (a) linear probing, (b) ResNet-Slim, and (c) ResNet-Wide.
• Figure 4: Magnitude (left) and phase (right) reconstruction for ContraWiMAE, obtained with patch size $(N_{\text{p,s}}, N_{\text{p,f}})=\textcolor{black}{(16,1)}$. Each row is a channel sample. Columns correspond to original, masked, and reconstructed channels.
• Figure 5: Beam selection performance under varying SNR for in-distribution, same-frequency OOD, and cross-frequency OOD settings with beam sizes 16 (top row) and 32 (bottom row)