CSI-MAE ViT-Base/16 and ViT-Large/16

CSI-MAE is a masked-autoencoder channel foundation model that learns reusable CSI representations through masked channel reconstruction. This repository contains the official model-only ViT-Base/16 and ViT-Large/16 pre-training weights associated with CSI-MAE: A Masked Autoencoder-based Channel Foundation Model.

• Code: GREAT-ISAC/CSI-MAE
• Reference data generation: GREAT-ISAC/Channel-Simulation-Data
• Paper: arXiv:2601.03789

Weight files

The release contains epoch-300 Base and Large checkpoints pretrained on the simulated Sionna/3GPP channel data. These are not the separate DeepMIMO experimental checkpoints. None of the files contains an optimizer, AMP scaler, epoch/resume state, training data, or a downstream task head.

File	Type	Intended use
csi-mae-base.safetensors	Base model-only pre-trained weights	Recommended Base weight for safe standalone loading and feature extraction
csi-mae-base.pth	Base model-only PyTorch checkpoint with a model key	Base compatibility weight for existing scripts
csi-mae-large.safetensors	Large model-only pre-trained weights	Recommended Large weight for safe standalone loading and feature extraction
csi-mae-large.pth	Large model-only PyTorch checkpoint with a model key	Large compatibility weight for existing scripts

These are channel-foundation-model pre-training weights, not final checkpoints for positioning, channel feedback, or channel extrapolation. The corresponding downstream architecture must be initialized from these weights and then adapted or fine-tuned. Exact SHA-256 values are recorded in manifest.json.

Model and input contract

Property	Value
Architecture	MAE ViT-Base/16 or ViT-Large/16
Input shape	[batch, 2, 256, 256]
Channel order	Real, imaginary
Dtype	float32
Pre-training channel mean	[-5.5445e-10, -5.5445e-10]
Pre-training channel variance	[5.7916e-10, 5.7916e-10]
Normalization	Channel-wise standardization using the fixed statistics above
Patch size	16
Pre-training mask ratio	0.75
Normalized pixel loss	Enabled

The same preprocessing must be used during training, fine-tuning, and inference. The implementation in dataset.py is authoritative.

Because normalized-pixel loss was enabled during pre-training, the decoder's patch predictions are normalized reconstruction targets. The loading example is therefore a forward/reconstruction smoke test, not recovery of CSI in its original physical scale. Physical-scale reconstruction additionally requires unpatchifying and applying the appropriate inverse patch and sample normalization.

Usage

Install the code and its minimal dependencies. CSI-MAE requires the pinned timm==0.3.2; newer timm releases remove the qk_scale API used by the released implementation.

git clone https://github.com/GREAT-ISAC/CSI-MAE.git
cd CSI-MAE
pip install -r requirements.txt

After downloading a weight, run the strict loading and forward smoke test. The --model value must match the selected checkpoint:

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

Expected output includes:

Reconstruction shape: (1, 256, 512)
Mask shape: (1, 256)

Run the same check on a complex cfr.npy sample:

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

Training data

The Base and Large models were pre-trained on simulated Sionna/3GPP CSI. The separate DeepMIMO experiments are not the source of these published weights. Generated training arrays are not included in this model repository. A reproducible, model-compatible reference data-generation pipeline is available in Channel Simulation Data. Its committed Sionna configuration is a runnable reference example; it does not reconstruct the complete checkpoint training data.

Intended use

• Research on wireless/channel foundation models.
• Masked CSI representation learning and reconstruction studies.
• Initialization for positioning, channel feedback, and channel extrapolation.
• Non-commercial evaluation and reproducibility studies.

Limitations and out-of-scope use

• The model was trained on simulated data; performance on measured channels is not guaranteed.
• The public architecture assumes two-channel 256 x 256 inputs. Other antenna or subcarrier layouts require an explicitly validated adaptation.
• Decoder outputs are not physical-scale complex CSI without the documented inverse-processing steps.
• The released weights do not provide final downstream predictions without adaptation or fine-tuning.
• The model is not intended for safety-critical deployment or commercial use.

License

The original CSI-MAE code, these model weights, and the repository-owned data-generation scripts are released under the Creative Commons Attribution-NonCommercial 4.0 International license (CC BY-NC 4.0). Attribution is required and commercial use is not permitted without prior written authorization from the copyright holders. Third-party software, simulators, datasets, and scenario assets remain subject to their respective licenses.

Citation

@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}
}

Acknowledgement

CSI-MAE is adapted from the public Masked Autoencoders implementation. The original attribution notices are retained in the source repository.

Contact

For questions, contact Jun Jiang at Jun.Jiang25@student.xjtlu.edu.cn.