EPFL-IVRL/wsr

Weight Space Representation Learning via Neural Field Adaptation — model zoo

Pretrained checkpoints for the CVPR 2026 paper Weight Space Representation Learning via Neural Field Adaptation (Zhuoqian Yang, Mathieu Salzmann, Sabine Süsstrunk — EPFL).

Paper | Code | Project page | ShapeNetSDF dataset

The folder layout mirrors the code repo, so downloading into the repo root restores every path the configs and REPRODUCE.md expect:

hf download EPFL-IVRL/wsr --local-dir .

The pipeline has three stages (Algorithm 1 in the paper). This repo holds each stage's outputs for mLoRA-Asym, the paper's main representation, so any stage can be skipped.

Stage 1 — base neural fields (3 checkpoints, ~200 MB)

Modulated base fields trained as variational autodecoders. Needed only for the LoRA / mLoRA representations. torch.load(...)['g_ema'] is the EMA model used everywhere.

Path	Dataset
neural_field/train_outputs_2d/nf-res-2d-o32-ffhq-128-5k/checkpoints/checkpoint_500000.pt	FFHQ 128 (5,000 images)
neural_field/train_outputs_3d/nf-res-3d-airplane/checkpoints/checkpoint_350000.pt	ShapeNet airplane
neural_field/train_outputs_3d/nf-res-3d-all5k10c-o1/checkpoints/checkpoint_350000.pt	ShapeNet multi-category (5k10c)

Stage 2 — fitted-weight datasets (3 folders, ~3 GB)

One folder per dataset (mLoRA-Asym representation); each holds one PyTorch state_dict per instance (weights/image_000000.pt ... in 2D, weights/shape_000000.pt ... in 3D). These are the training data for Stage 3 and the inputs for the reconstruction / discriminative evaluations (paper Tables 1 and 4).

Folder	Dataset	Instances
neural_field/overfit_outputs_2d/2d_nf-res-2d-o32-ffhq-128-5k_mlora_r4_asym_0.0	FFHQ 128	5,000
neural_field/overfit_outputs_3d/3d_nf-res-3d-airplane_mlora_r3_asym_0.0	airplane	4,045
neural_field/overfit_outputs_3d/3d_nf-res-3d-all5k10c-o1_mlora_r3_asym_0.0	multi (10k10c superset)	10,000

Note on the multi-category splits: 5k10c (4,999 shape ids) is a subset of 10k10c (9,999 ids) but not a line-order prefix. The 10,000-instance folder is indexed in 10k10c line order; extracting the 5k10c subset requires the id-to-index mapping from the split files ShapeNetSDF/meta/all/{5k10c,10k10c}.txt in the ShapeNetSDF dataset.

Stage 3 — weight-space diffusion (3 checkpoints, ~16 GB)

One Diffusion Transformer per dataset, trained on the mLoRA-Asym weights (fp32 PyTorch Lightning checkpoints, weights only, no optimizer state — ~5 GB each). Directory names equal the Hydra config names in weight_space_diffusion/configs/diffusion_configs/; each holds last.ckpt.

Config (weight_space_diffusion/checkpoints/<CONFIG>/last.ckpt)	Dataset	Paper
train_ffhq_128_nfres_mlora_r4_asym	FFHQ 128	Table 2
train_airplane_nfres3d_mlora_r3_asym	ShapeNet airplane	Table 3
train_all_5k10c_nfres3d_mlora_r3_asym	ShapeNet multi (5k10c)	Table 3

Sampling example (from the code repo, see REPRODUCE.md for details):

cd weight_space_diffusion
python apps/inference.py --config-name=<CONFIG> \
    inference.model_path=checkpoints/<CONFIG>/last.ckpt

License

CC BY-NC-SA 3.0. The Stage-3 diffusion code derives from HyperDiffusion (CC BY-NC-SA 3.0); the neural_field code is MIT with NVIDIA-licensed CUDA ops. See the code repo's LICENSE files and ACKNOWLEDGEMENT.md.

Citation

@inproceedings{yang2026wsr,
  title     = {Weight Space Representation Learning via Neural Field Adaptation},
  author    = {Yang, Zhuoqian and Salzmann, Mathieu and S{\"u}sstrunk, Sabine},
  booktitle = {Proceedings of the IEEE/CVF Conference on
               Computer Vision and Pattern Recognition (CVPR)},
  year      = {2026}
}