📡 CSI-MAE: A Masked Autoencoder-based Channel Foundation Model

This is the official code release for CSI-MAE: A Masked Autoencoder-based Channel Foundation Model (arXiv:2601.03789). CSI-MAE is a Channel Foundation Model (CFM) that adapts masked autoencoding to channel state information (CSI), learning transferable wireless-channel representations for both communication and sensing tasks.

CSI-MAE is released by GREAT Wireless AI, an open research organization developing channel foundation models and AI-native systems for wireless communications, sensing, and radio access networks.

✨ Contributions

CSI-MAE is designed as a reusable CFM for wireless channels:

• Unified perception and generation: supports CSI generation tasks such as channel feedback and extrapolation, as well as CSI perception tasks such as user positioning.
• Masked channel modeling: treats complex CSI as a two-channel real/imaginary representation and learns robust spatial-frequency features through masked reconstruction.
• Efficient adaptation: supports full fine-tuning, supervised training from scratch, and lightweight decoder/head fine-tuning with a frozen encoder.
• Cross-scenario transfer: targets reusable CSI representations with strong cross-scenario and zero-shot transfer capability.

🧩 Architecture

CSI-MAE follows the asymmetric MAE encoder-decoder design and adapts it to CSI matrices. The encoder receives only visible CSI patches, while the decoder reconstructs the full channel using latent features and mask tokens. Downstream tasks attach task-specific decoders or heads to the learned CSI encoder.

📁 Repository Contents

• main_pretrain.py: self-supervised CSI-MAE pre-training.
• main_generate.py: channel feedback and channel extrapolation training.
• main_finetune.py: positioning fine-tuning.
• inference.py: checkpoint evaluation for feedback/extrapolation.
• models_mae.py, models_generate.py, models_vit.py: model definitions.
• dataset.py: CSI dataset loaders.
• pretrain.sh, feedback.sh, extrapolation.sh, positioning.sh: path-anonymized run templates.
• load_pretrained.py: strict checkpoint loading and reconstruction smoke test.

Generated data, checkpoints, TensorBoard logs, and experiment outputs are not stored in this Git repository. Official weights are distributed separately on Hugging Face. Reproducible simulator and model-input preparation code is maintained in Channel Simulation Data, including the shared <scenario>/cfr.npy contract used by CSI-MAE and CSI-CLIP. The committed configurations are model-compatible reference examples and do not reconstruct the complete checkpoint training data.

⚙️ Environment

Required packages include:

• Python 3.8+
• PyTorch
• torchvision
• timm==0.3.2
• numpy
• tensorboard
• safetensors

Example installation:

pip install -r requirements.txt

⚖️ Pre-trained Weights

The official release provides both ViT-Base/16 and ViT-Large/16 checkpoints pretrained on the simulated Sionna/3GPP channel corpus with a 75% masking ratio. These are the epoch-300 checkpoints used by this release, not the separate DeepMIMO experimental checkpoints.

Variant	Recommended weight	Compatible PyTorch weight
CSI-MAE Base	csi-mae-base.safetensors	csi-mae-base.pth
CSI-MAE Large	csi-mae-large.safetensors	csi-mae-large.pth

The safetensors files are recommended for standalone loading. Model-only .pth files support the existing fine-tuning scripts. Optimizer state, AMP scaler state, local paths, and other training-resume metadata are removed.

After downloading the matching Base or Large .safetensors file, verify strict loading and a masked reconstruction forward pass with:

python load_pretrained.py \
  --model base \
  --checkpoint /path/to/CSI-MAE/csi-mae-base.safetensors

Expected output includes reconstruction shape (1, 256, 512). Real complex CSI must first use the two-channel conversion and normalization implemented in dataset.py. The example performs this preprocessing when an input file is provided:

python load_pretrained.py \
  --model large \
  --checkpoint /path/to/CSI-MAE/csi-mae-large.safetensors \
  --input /path/to/scenario/cfr.npy \
  --sample-index 0

Pre-training and feature extraction standardize the real and imaginary channels using the fixed statistics recorded in dataset.py. This input normalization is part of the published weight contract and must not be replaced with per-sample min-max normalization.

Download the official release from GREAT-Wireless-AI/CSI-MAE.

File	SHA-256
csi-mae-base.safetensors	db7ad83987ae7f8a5f0b42da9ba1749f792a27aa37356c2f5ba94aaa09f44a46
csi-mae-base.pth	c1ac1b2cddd4cd41b7edc0970a34e34668d48dec26c73be4feec68ccde846694
csi-mae-large.safetensors	d3ea94ead704488b410b1e86ceaa23f8c6446dfad93af24827715d453c915f25
csi-mae-large.pth	c65dba88c94c584d42f7e3cdb3975840d8af273992f1a4f8352f7a4931b59f38

🗂️ Data Layout

Prepare data locally and pass paths through environment variables or command-line arguments.

Pre-training data:

DATA_ROOT/
  scenario_a/
    cfr.npy
  scenario_b/
    cfr.npy

Feedback/extrapolation data:

DATA_ROOT/
  scenario_a/
    train_data.npy
    val_data.npy
  scenario_b/
    train_data.npy
    val_data.npy

Positioning data:

DATA_ROOT/
  scenario_a/
    train_csi.npy
    val_csi.npy
    train_pos.npy
    val_pos.npy

🚀 How to Run

🧠 Pre-train CSI-MAE

DATA_ROOT=/path/to/pretrain_data \
OUTPUT_ROOT=./outputs \
MODEL_VARIANT=base \
bash pretrain.sh

Train the Large variant with the recorded Large-model defaults:

DATA_ROOT=/path/to/pretrain_data \
OUTPUT_ROOT=./outputs \
MODEL_VARIANT=large \
bash pretrain.sh

🧪 Train Without Pre-training

Use MODE=base to train the downstream model from scratch as a supervised baseline.

Channel feedback:

DATA_ROOT=/path/to/sionna_data \
SCENARIO=scenario_a \
MODE=base \
bash feedback.sh

Channel extrapolation:

DATA_ROOT=/path/to/sionna_data \
BASE_SCENARIO=scenario_a \
PRED_SCENARIO=scenario_b \
MODE=base \
bash extrapolation.sh

Positioning:

DATA_ROOT=/path/to/positioning_data \
SCENARIO=scenario_a \
MODE=base \
bash positioning.sh

🔧 Fine-tune With a Pre-trained Checkpoint

Use MODE=finetune for full-parameter fine-tuning.

DATA_ROOT=/path/to/sionna_data \
SCENARIO=scenario_a \
CKPT_PATH=/path/to/pretrained_checkpoint.pth \
MODE=finetune \
bash feedback.sh

DATA_ROOT=/path/to/sionna_data \
BASE_SCENARIO=scenario_a \
PRED_SCENARIO=scenario_b \
CKPT_PATH=/path/to/pretrained_checkpoint.pth \
MODE=finetune \
bash extrapolation.sh

DATA_ROOT=/path/to/positioning_data \
SCENARIO=scenario_a \
CKPT_PATH=/path/to/pretrained_checkpoint.pth \
MODE=finetune \
bash positioning.sh

❄️ Fine-tune a Lightweight Decoder With Frozen Encoder

Use MODE=freeze for feedback/extrapolation to load a pre-trained CSI-MAE encoder and train only the task-side generation module.

DATA_ROOT=/path/to/sionna_data \
SCENARIO=scenario_a \
CKPT_PATH=/path/to/pretrained_checkpoint.pth \
MODE=freeze \
bash feedback.sh

DATA_ROOT=/path/to/sionna_data \
BASE_SCENARIO=scenario_a \
PRED_SCENARIO=scenario_b \
CKPT_PATH=/path/to/pretrained_checkpoint.pth \
MODE=freeze \
bash extrapolation.sh

Positioning with a frozen encoder:

DATA_ROOT=/path/to/positioning_data \
SCENARIO=scenario_a \
CKPT_PATH=/path/to/pretrained_checkpoint.pth \
MODE=freeze \
bash positioning.sh

📊 Evaluate a Feedback/Extrapolation Checkpoint

python inference.py \
  --base_dir /path/to/base_scenario \
  --pred_dir /path/to/target_scenario \
  --checkpoint /path/to/checkpoint.pth

Evaluate a positioning checkpoint:

DATA_ROOT=/path/to/positioning_data \
SCENARIO=scenario_a \
CKPT_PATH=/path/to/best_checkpoint.pth \
MODE=eval \
bash positioning.sh

📧 Contact

If you have any questions, please feel free to contact Jun Jiang at Jun.Jiang25@student.xjtlu.edu.cn.

🙏 Acknowledgement

This codebase is adapted from the excellent MAE repository. We thank the MAE authors for releasing their PyTorch implementation. Original MAE attribution notices are retained in source files where applicable.

📜 License

This project is released for research and other non-commercial use only under the Creative Commons Attribution-NonCommercial 4.0 International license. Commercial use is prohibited unless prior written authorization is obtained from the authors. See LICENSE for details.

📝 Citation

If you find this work helpful, please consider citing:

@article{jiang2026csimae,
  title={CSI-MAE: A Masked Autoencoder-based Channel Foundation Model},
  author={Jiang, Jun and Ruan, Xiaolong and Xu, Shugong},
  journal={arXiv preprint arXiv:2601.03789},
  year={2026}
}

Some other related papers and resources:

• A MIMO Wireless Channel Foundation Model via CIR-CSI Consistency
  Paper: IEEE Xplore | Code: GREAT-ISAC/CSI-CLIP
• Towards Channel Foundation Models (CFMs): Motivations, Methodologies and Opportunities
  Paper: arXiv:2507.13637 | GitHub: GREAT-ISAC/Awesome-Channel-Foundation-Models

@inproceedings{jiang2025csi_clip,
  title={A MIMO Wireless Channel Foundation Model via CIR-CSI Consistency},
  author={Jiang, Jun and Yu, Wenjun and Li, Yunfan and Gao, Yuan and Xu, Shugong},
  booktitle={2025 IEEE International Conference on Machine Learning for Communication and Networking (ICMLCN)},
  pages={1--6},
  year={2025},
  doi={10.1109/ICMLCN64995.2025.11140262}
}

@article{jiang2025cfmsurvey,
  title={Towards Channel Foundation Models (CFMs): Motivations, Methodologies and Opportunities},
  author={Jiang, Jun and Gao, Yuan and Wu, Xinyi and Xu, Shugong},
  journal={arXiv preprint arXiv:2507.13637},
  year={2025}
}