HeterCSI: Channel-Adaptive Heterogeneous CSI Pretraining Framework for Generalized Wireless Foundation Models

TL;DR

HeterCSI tackles the generalization gap in wireless foundation models caused by dual CSI heterogeneity—scale and scenario—by introducing a channel-adaptive pretraining framework. The method couples a scale-aware batching strategy with a double masking MAE-based self-supervised objective, effectively mitigating destructive gradient interference from scale differences while preserving constructive scenario diversity. Empirical results on QuaDRiGa-generated data across 12 unseen scenarios show substantial NMSE improvements over zero-shot baselines (averages of 7.19 dB, 4.08 dB, and 5.27 dB for CSI reconstruction, time-domain, and frequency-domain predictions, respectively) and a 53% reduction in training time compared with global shuffling. The work demonstrates scalable, generalizable wireless foundation models without scenario-specific finetuning, offering practical impact for 6G CSI processing and beyond, with future directions toward physics-informed encodings and architecture optimizations.

Abstract

Wireless foundation models promise transformative capabilities for channel state information (CSI) processing across diverse 6G network applications, yet face fundamental challenges due to the inherent dual heterogeneity of CSI across both scale and scenario dimensions. However, current pretraining approaches either constrain inputs to fixed dimensions or isolate training by scale, limiting the generalization and scalability of wireless foundation models. In this paper, we propose HeterCSI, a channel-adaptive pretraining framework that reconciles training efficiency with robust cross-scenario generalization via a new understanding of gradient dynamics in heterogeneous CSI pretraining. Our key insight reveals that CSI scale heterogeneity primarily causes destructive gradient interference, while scenario diversity actually promotes constructive gradient alignment when properly managed. Specifically, we formulate heterogeneous CSI batch construction as a partitioning optimization problem that minimizes zero-padding overhead while preserving scenario diversity. To solve this, we develop a scale-aware adaptive batching strategy that aligns CSI samples of similar scales, and design a double-masking mechanism to isolate valid signals from padding artifacts. Extensive experiments on 12 datasets demonstrate that HeterCSI establishes a generalized foundation model without scenario-specific finetuning, achieving superior average performance over full-shot baselines. Compared to the state-of-the-art zero-shot benchmark WiFo, it reduces NMSE by 7.19 dB, 4.08 dB, and 5.27 dB for CSI reconstruction, time-domain, and frequency-domain prediction, respectively. The proposed HeterCSI framework also reduces training latency by 53% compared to existing approaches while improving generalization performance by 1.53 dB on average.

Figures (4)

• Figure 1: (a) CSI varies significantly across scales ($T, K, U$) and propagation scenarios. (b) Jointly leveraging heterogeneous CSI aggregates complementary gradients ($g_1, g_2$) to guide optimization toward a robust global optimum, while conflicting gradients ($g_1, g_3$) lead to destructive interference. (c) Similarity statistics confirm that scale discrepancy is the primary source of gradient conflict (left: scale heterogeneity), while scenario diversity maintains positive alignment when scales are compatible (middle and right).
• Figure 2: Overview of the HeterCSI framework. The framework introduces scale-aware adaptive batching strategy, which synergizes intra-bucket gradient alignment with inter-bucket stochastic scheduling. Furthermore, a double masking mechanism is designed to rigorously isolate valid signals from padding artifacts via attention masking, while simultaneously driving self-supervised representation learning through MAE random masking. The processed tokens are finally reconstructed via a standard Transformer-based encoder-decoder architecture.
• Figure 3: Zero-shot NMSE performance versus the number of training datasets for CSI reconstruction (left), time-domain prediction (middle), and frequency-domain prediction (right). The proposed method consistently outperforms baselines as data volume increases.
• Figure 4: Impact of bucket granularity on zero-shot NMSE performance and training time under different heterogeneity levels.

Theorems & Definitions (2)

• Theorem 1
• proof