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 · project page · code · token datasets.

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

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:

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

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