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.
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.
1git clone https://github.com/GREAT-ISAC/CSI-MAE.git
2cd CSI-MAE
3pip install -r requirements.txt
After downloading a weight, run the strict loading and forward smoke test. The
--model value must match the selected checkpoint:
1python load_pretrained.py \
2 --model base \
3 --checkpoint /path/to/csi-mae-base.safetensors
Expected output includes:
1Reconstruction shape: (1, 256, 512)
2Mask shape: (1, 256)
Run the same check on a complex cfr.npy sample:
1python load_pretrained.py \
2 --model large \
3 --checkpoint /path/to/csi-mae-large.safetensors \
4 --input /path/to/scenario/cfr.npy \
5 --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
1@article{jiang2026csimae,
2 title={CSI-MAE: A Masked Autoencoder-based Channel Foundation Model},
3 author={Jiang, Jun and Ruan, Xiaolong and Xu, Shugong},
4 journal={arXiv preprint arXiv:2601.03789},
5 year={2026}
6}
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.
