README.md

---
license: mit
tags:
  - amortized-inference
  - probabilistic-models
  - transformer
  - bayesian-optimization
  - simulation-based-inference
  - gaussian-process
---

# nanoACE — source checkpoints

Trained PyTorch checkpoints for [**nanoACE**](https://github.com/acerbilab/nanoACE),
a small, readable implementation of the
[Amortized Conditioning Engine (ACE)](https://acerbilab.github.io/amortized-conditioning-engine/)
(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](https://acerbilab.github.io/nanoACE/) uses a
*different* artifact — fp16 browser blobs derived from these, hosted separately at
[acerbilab/nanoACE-playground-weights](https://github.com/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](https://github.com/acerbilab/nanoACE/tree/main/extensions/arbuffer) extension — causal AR buffer (Hassan et al., 2026) | base 128 / 4 / 4 / 8 + buffer stream | 0 | 200k |
| `gp1d_aline.pt` | [aline](https://github.com/acerbilab/nanoACE/tree/main/extensions/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](https://github.com/acerbilab/nanoACE) and run from its root:

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

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

```bibtex
@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](https://github.com/acerbilab/nanoACE/blob/main/LICENSE).
Developed by the [Machine and Human Intelligence group](https://www.helsinki.fi/en/researchgroups/machine-and-human-intelligence)
at the University of Helsinki.