---
license: cc-by-nc-4.0
library_name: pytorch
tags:
- pde
- neural-operator
- tokenizer
- scientific-machine-learning
- fluid-dynamics
- compressible-euler
- masked-autoencoder
- safetensors
datasets:
- llingsch/phaedra-tokens
- camlab-ethz/CE-KH
- camlab-ethz/CE-CRP
- camlab-ethz/CE-RPUI
---

# Phaedra: tokenizers, neural operators and masked autoencoders

Weights accompanying **Phaedra: Learning High-Fidelity Discrete Tokenization for the Physical Sciences** (Levi Lingsch, Georgios Kissas, Johannes Jakubik, Siddhartha Mishra, NeurIPS 2026) — [paper](https://arxiv.org/abs/2602.03915) · [project page](https://camlab-ethz.github.io/Phaedra/) · [code](https://github.com/camlab-ethz/Phaedra) · [token datasets](https://huggingface.co/datasets/llingsch/phaedra-tokens).

Phaedra tokenizes each 4×4 patch of a physical field into **two discrete tokens: an amplitude token and a
morphology token**. This repository contains every model evaluated in the paper: the Phaedra tokenizer and
the three tokenizer baselines, the 21 neural operators of the main comparison (7 architectures × 3 datasets),
and the masked autoencoders. All weights are `safetensors`; the folder layout mirrors the output directory
the [code](https://github.com/camlab-ethz/Phaedra) expects, so downloaded weights are used without any renaming.

## Contents

| folder | model | parameters | codes / input |
|---|---|---|---|
| `tokenizers/phaedra_4x4` | Phaedra (`Phaedra_AE_FSQ_4x4`) | 97.84M | amplitude 1024 x morphology 8640 (FSQ [5,4,4,3,3,3,2,2]) on 32x32 |
| `tokenizers/fsq` | FSQ (`AE_FSQ`) | 97.83M | 8640 (FSQ [5,4,4,3,3,3,2,2]) on 32x32 |
| `tokenizers/vqvae2` | VQ-VAE-2 (`AE_VQVAE2`) | 112.19M | top 4096 on 16x16 + bottom 16384 on 32x32 |
| `tokenizers/continuous` | continuous AE (`AE_Continuous`) | 97.83M | 8 float channels on 32x32 (no quantization) |
| `operators/phaedra_38m_{kh,rc,rkh}` | Phaedra transformer | 38.06M | Phaedra tokens (amplitude + morphology) |
| `operators/fsq_38m_{kh,rc,rkh}` | FSQ transformer | 37.30M | FSQ tokens |
| `operators/vqvae2_38m_{kh,rc,rkh}` | VQ-VAE-2 transformer | 44.30M | VQ-VAE-2 tokens (top + bottom) |
| `operators/continuous_38m_{kh,rc,rkh}` | Continuous-latent transformer | 38.24M | continuous AE latents |
| `operators/fno_38m_{kh,rc,rkh}` | FNO | 37.83M | physical fields |
| `operators/cno_38m_{kh,rc,rkh}` | CNO | 38.78M | physical fields |
| `operators/vit_38m_{kh,rc,rkh}` | ViT | 37.92M | physical fields |
| `mae/mae_{phaedra,fsq}_3pde` | token MAE, pre-trained on KH + RKH + RC | 48.00M / 47.48M | Phaedra / FSQ tokens |
| `mae/mae_{phaedra,fsq}_finetune_{kh,rc,rkh}` | the 3-PDE MAE fine-tuned on one dataset | 48.00M / 47.48M | Phaedra / FSQ tokens |

Every folder holds `model.safetensors`, the training `config.yaml` (the exact recipe; paths are relative to
`$PHAEDRA_DATA_ROOT` / `$PHAEDRA_OUTPUT_ROOT`) and `metadata.json` (provenance, sha256, results).

## Quick start

```bash
git clone https://github.com/camlab-ethz/Phaedra && cd Phaedra
pip install -r requirements.txt
cp env.example.sh env.sh && source env.sh      # sets PHAEDRA_DATA_ROOT / PHAEDRA_OUTPUT_ROOT

# everything needed to evaluate one model (operator + its tokenizer + its test tokens)
python scripts/download_pretrained.py eval phaedra_38m_kh
# ground-truth fields from the public Poseidon datasets (numeric chunk order, hash-verified)
python scripts/prepare_poseidon_data.py --datasets CE-KH --download
python -m evaluation.eval_downstream --models phaedra_38m_kh --mode rollout662 --timesteps 6 12 14 --deterministic
```

Use a tokenizer on its own:

```python
from tokenizer.pretrained import load_tokenizer
tok = load_tokenizer("Phaedra_AE_FSQ_4x4")     # downloads tokenizers/phaedra_4x4 from this repository
amp, morph = tok.encode(x)                       # x: [B, 1, 128, 128], normalized per variable (table below)
x_rec = tok.decode((amp, morph))
```

## Results

Relative $L_1$ error (%) at the final time $t = 0.7$ on the 240 test trajectories of each dataset; each
model under its best prediction strategy (direct $0\to t$, 2-step autoregressive, or 6-6-2 rollout); ± is
the 95 % bootstrap confidence interval of the mean. Reproduced from these files with
`python -m evaluation.eval_downstream --deterministic` (bit-identical to the paper's evaluation on an RTX 4090).

| model | data | ρ | u | v | p | average | strategy |
|---|---|---|---|---|---|---|---|
| Phaedra transformer | KH (CE-KH) | 4.78 | 8.43 | 24.29 | 0.51 | **9.50** ± 0.34 | direct |
| Phaedra transformer | RC (CE-CRP) | 20.66 | 40.73 | 40.42 | 7.02 | **27.21** ± 0.71 | 6-6-2 |
| Phaedra transformer | RKH (CE-RPUI) | 6.52 | 16.01 | 15.64 | 2.74 | **10.23** ± 0.80 | 6-6-2 |
| FSQ transformer | KH (CE-KH) | 5.30 | 10.91 | 28.27 | 0.62 | **11.28** ± 0.34 | 6-6-2 |
| FSQ transformer | RC (CE-CRP) | 28.12 | 57.08 | 56.77 | 10.21 | **38.05** ± 1.05 | 6-6-2 |
| FSQ transformer | RKH (CE-RPUI) | 10.37 | 27.73 | 28.78 | 5.08 | **17.99** ± 1.37 | 6-6-2 |
| VQ-VAE-2 transformer | KH (CE-KH) | 7.95 | 13.90 | 38.04 | 0.68 | **15.14** ± 0.48 | 6-6-2 |
| VQ-VAE-2 transformer | RC (CE-CRP) | 29.54 | 63.17 | 64.39 | 12.48 | **42.40** ± 0.63 | direct |
| VQ-VAE-2 transformer | RKH (CE-RPUI) | 14.73 | 38.09 | 37.29 | 7.97 | **24.52** ± 1.15 | 6-6-2 |
| Continuous-latent transformer | KH (CE-KH) | 4.64 | 8.34 | 23.53 | 0.61 | **9.28** ± 0.30 | 6-6-2 |
| Continuous-latent transformer | RC (CE-CRP) | 77.95 | 197.31 | 156.93 | 44.00 | **119.05** ± 2.96 | direct |
| Continuous-latent transformer | RKH (CE-RPUI) | 7.09 | 18.09 | 17.87 | 3.18 | **11.56** ± 0.70 | 6-6-2 |
| FNO | KH (CE-KH) | 6.24 | 11.18 | 32.23 | 0.54 | **12.55** ± 0.36 | 6-6-2 |
| FNO | RC (CE-CRP) | 26.31 | 53.62 | 53.42 | 9.16 | **35.63** ± 0.69 | 6-6-2 |
| FNO | RKH (CE-RPUI) | 10.08 | 26.36 | 25.67 | 4.54 | **16.66** ± 0.91 | 6-6-2 |
| CNO | KH (CE-KH) | 5.06 | 9.10 | 25.93 | 0.51 | **10.15** ± 0.32 | 2-step |
| CNO | RC (CE-CRP) | 22.65 | 44.44 | 44.66 | 7.94 | **29.92** ± 0.74 | 2-step |
| CNO | RKH (CE-RPUI) | 6.79 | 18.83 | 18.50 | 3.17 | **11.82** ± 0.82 | 6-6-2 |
| ViT | KH (CE-KH) | 4.64 | 8.25 | 23.33 | 0.51 | **9.18** ± 0.27 | 6-6-2 |
| ViT | RC (CE-CRP) | 28.76 | 60.17 | 59.51 | 10.28 | **39.68** ± 0.58 | direct |
| ViT | RKH (CE-RPUI) | 9.09 | 20.86 | 19.66 | 3.18 | **13.20** ± 0.79 | 6-6-2 |

Tokenizer reconstruction floor (relative $L_1$ %, $t = 0.7$, mean of ρ, u, v, p): error of decoding the
ground-truth tokens, i.e. the best any operator on these tokens can do.

| tokenizer | KH | RC | RKH |
|---|---|---|---|
| Phaedra | 0.54 | 3.30 | 1.49 |
| FSQ | 1.01 | 5.16 | 2.36 |
| VQ-VAE-2 | 1.37 | 8.39 | 3.45 |
| continuous AE | 0.27 | 1.48 | 0.65 |

Masked autoencoders (evaluation of the paper; relative $L_1$ % of the reconstructed fields, mean of ρ, u, v, p).
The masks are random: these numbers come from one unseeded run each; `python -m mae.test --seed 0` with the
released weights reproduces them to within 0.5–2 % (Phaedra) and ≤ 5 % (FSQ fine-tunes, scored on a single
test trajectory) relative.

| model | tokens | trained on | parameters | evaluated on | test trajectories | rel. $L_1$ |
|---|---|---|---|---|---|---|
| `mae/mae_fsq_3pde` | fsq | KelvinHelmholtz, RiemannKelvinHelmholtz, RiemannCurved | 47.48M | RKH | 10 | 16.91 |
| `mae/mae_fsq_finetune_kh` | fsq | KelvinHelmholtz | 47.48M | KH | 1 | 6.29 |
| `mae/mae_fsq_finetune_rc` | fsq | RiemannCurved | 47.48M | RC | 1 | 28.20 |
| `mae/mae_fsq_finetune_rkh` | fsq | RiemannKelvinHelmholtz | 47.48M | RKH | 1 | 7.07 |
| `mae/mae_phaedra_3pde` | phaedra | KelvinHelmholtz, RiemannKelvinHelmholtz, RiemannCurved | 48.00M | RKH | 10 | 16.43 |
| `mae/mae_phaedra_finetune_kh` | phaedra | KelvinHelmholtz | 48.00M | KH | 10 | 4.76 |
| `mae/mae_phaedra_finetune_rc` | phaedra | RiemannCurved | 48.00M | RC | 10 | 22.17 |
| `mae/mae_phaedra_finetune_rkh` | phaedra | RiemannKelvinHelmholtz | 48.00M | RKH | 10 | 7.27 |

## Checkpoint notes

* **Weights evaluated in the paper.** Tokenizers: EMA weights. Operators: the final training step, except
  `phaedra_38m_rkh` (step 56000 of 58000) and `continuous_38m_rkh` (epoch 59 of 100 — that run diverged
  after epoch ~63; validation latent-$L_1$ 3.24 → 11.3).
* `continuous_38m_rc` did not train successfully (predictions collapse to the mean; ~119 % error). It is
  released unchanged because it is the model the paper reports.
* `vqvae2_38m_*` were trained for 5 epochs only (compute-bound); the VQ-VAE-2 tokenizer uses ~135 of its
  4096 top codes and ~380 of its 16384 bottom codes on these datasets.
* `mae_phaedra_3pde` also ships `model_raw.safetensors` (non-EMA): the fine-tuning runs were initialized
  from these, and `mae.train` picks them up automatically for `warm_start_from`.
* The ViT operators for RC/RKH were trained with the KH normalization statistics (as in their configs); the
  evaluation code handles this.

## Normalization

Fields are normalized per variable, $x' = (x - \mu) / \sigma$, before tokenization (μ / σ):

| data | ρ | u | v | p |
|---|---|---|---|---|
| KH (CE-KH) | 0.75 / 0.22776 | -0.016416 / 0.112602 | 5e-05 / 0.044217 | 1 / 0.0083488 |
| RC (CE-CRP) | 0.548238 / 0.322666 | 0.00042 / 0.275455 | 0.0026327 / 0.275455 | 0.552022 / 0.169886 |
| RKH (CE-RPUI) | 0.543686 / 0.361854 | -0.00332036 / 0.210691 | 0.00215155 / 0.215305 | 0.548822 / 0.199823 |

## Verification and integrity

`SHA256SUMS` lists every file. Before publication every operator was re-evaluated from these files with the released code (all 240 test trajectories x 4 variables x all timesteps, all three prediction modes; 342,720 errors in total): all are bit-identical to the paper's evaluation. The MAE weights were A/B-tested against the original checkpoints under a fixed seed: 8/8 give identical metrics.

## License

Weights: CC-BY-NC-4.0 (the models are trained on the CC BY-NC 4.0 Poseidon datasets). Code: MIT.

## Citation

```bibtex
@inproceedings{lingsch2026phaedra,
  title         = {Phaedra: Learning High-Fidelity Discrete Tokenization for the Physical Sciences},
  author        = {Lingsch, Levi and Kissas, Georgios and Jakubik, Johannes and Mishra, Siddhartha},
  booktitle     = {Advances in Neural Information Processing Systems},
  year          = {2026},
  eprint        = {2602.03915},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2602.03915}
}
```