nanoACE — source checkpoints

Trained PyTorch checkpoints for nanoACE, a small, readable implementation of the Amortized Conditioning Engine (ACE) (Chang et al., AISTATS 2025). ACE treats data, interpretable latents, and runtime prior information all as tokens: condition on one token set, predict distributions over another.

These are the full-precision source checkpoints — the .pt files written by each example's --save-checkpoint, each {cfg, seed, state_dict} (the extensions also carry a config provenance record). They load straight back into the example scripts. The interactive playground uses a different artifact — fp16 browser blobs derived from these, hosted separately at acerbilab/nanoACE-playground-weights.

Checkpoints

File	Task	Architecture (d_model / layers / heads / MDN K)	Seed	Steps
gaussian_toy.pt	Gaussian ACEP — infer mu/log_sigma with runtime Beta priors	96 / 3 / 4 / 8 (~0.55M)	0	320k
gp1d.pt	GP-1D regression — kernel + hyperparameters as latents	128 / 4 / 4 / 8 (~1.24M)	0	200k
sbi_sir.pt	SIR simulation-based inference — epidemic rates beta/gamma	128 / 4 / 4 / 8 (~1.24M)	0	100k
bo1d.pt	BO-1D — optimum location/value as latents, robust prior injection	192 / 6 / 16 / 12 (~3.96M)	0	200k
gp1d_arbuffer.pt	arbuffer extension — causal AR buffer (Hassan et al., 2026)	base 128 / 4 / 4 / 8 + buffer stream	0	200k
gp1d_aline.pt	aline extension — joint inference + active acquisition (Huang et al., 2025)	base 128 / 4 / 4 / 8 + policy decoder	0	35k

The two extension checkpoints are warm-started from gp1d.pt (concat-read for arbuffer; the served 35k policy fine-tune for aline) and carry a full config run-provenance record, as does the gaussian_toy.pt 320k retrain; the other three core checkpoints carry cfg + seed only.

Usage

huggingface_hub is the only extra needed to fetch (pip install huggingface_hub); loading uses the nanoACE example scripts, so clone the repo and run from its root:

from huggingface_hub import hf_hub_download
import gp1d                                        # nanoACE example module

path = hf_hub_download("lacerbi/nanoACE", "gp1d.pt")
model = gp1d.load_checkpoint(path, "cpu")          # returns a ready ACE model

The same pattern works for gaussian_toy, sbi_sir, and bo1d. The extension checkpoints load through their own modules (extensions/arbuffer/gp1d_arbuffer.py, extensions/aline/gp1d_aline.py) — see each extension's README for the exact --load-checkpoint invocation.

Or just point the example's CLI at a downloaded file:

python gp1d.py --eval-only --load-checkpoint <path-to>/gp1d.pt

Regenerating

Each checkpoint is regenerable, not a mystery binary: it stores its cfg and seed, and the training data stream is a pure function of (seed, step). Re-running the matching example at the listed step count reproduces it (see the nanoACE README and DEVLOG.md). The checkpoints are example artifacts for inspection and reuse, not a packaged runtime product.

References

@inproceedings{chang2025amortized,
  title={Amortized Probabilistic Conditioning for Optimization, Simulation and Inference},
  author={Chang, Paul E and Loka, Nasrulloh and Huang, Daolang and Remes, Ulpu and Kaski, Samuel and Acerbi, Luigi},
  booktitle={The Twenty-eighth International Conference on Artificial Intelligence and Statistics (AISTATS 2025)},
  year={2025}
}

@inproceedings{hassan2026efficient,
  title={Efficient Autoregressive Inference for Transformer Probabilistic Models},
  author={Conor Hassan and Nasrulloh Ratu Bagus Satrio Loka and Cen-You Li and Daolang Huang and Paul Edmund Chang and Yang Yang and Francesco Silvestrin and Samuel Kaski and Luigi Acerbi},
  booktitle={The Fourteenth International Conference on Learning Representations (ICLR 2026)},
  year={2026}
}

@inproceedings{huang2025aline,
  title={ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition},
  author={Daolang Huang and Xinyi Wen and Ayush Bharti and Samuel Kaski and Luigi Acerbi},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems (NeurIPS 2025)},
  year={2025}
}

License

MIT — see the nanoACE repository. Developed by the Machine and Human Intelligence group at the University of Helsinki.