README.md

---
license: cc-by-nc-4.0
library_name: pytorch
pipeline_tag: feature-extraction
tags:
  - wireless
  - channel-state-information
  - channel-foundation-model
  - masked-autoencoder
  - vision-transformer
---

# 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*](https://arxiv.org/abs/2601.03789).

- Code: [GREAT-ISAC/CSI-MAE](https://github.com/GREAT-ISAC/CSI-MAE)
- Reference data generation: [GREAT-ISAC/Channel-Simulation-Data](https://github.com/GREAT-ISAC/Channel-Simulation-Data)
- Paper: [arXiv:2601.03789](https://arxiv.org/abs/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.

```bash
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:

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

Expected output includes:

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

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

```bash
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](https://github.com/GREAT-ISAC/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](https://creativecommons.org/licenses/by-nc/4.0/)
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

```bibtex
@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](https://github.com/facebookresearch/mae)
implementation. The original attribution notices are retained in the source
repository.

## Contact

For questions, contact Jun Jiang at
[Jun.Jiang25@student.xjtlu.edu.cn](mailto:Jun.Jiang25@student.xjtlu.edu.cn).