learning-bayesian-statistics/baygent-skills/amortized-workflow

Amortized Bayesian Workflow

Workflow overview

Every amortized Bayesian analysis follows this sequence. Do not skip steps — especially simulator validation and model criticism.

1. Formulate — Define the generative story. What latent variables or parameters generated the observations?
2. Specify the simulator regime — The first iteration always uses offline training for fast turnaround, regardless of simulator speed. The simulator regime only determines the simulation budget for the pilot run:

• Fast simulator (< 0.05 s per draw): pre-simulate 20 000 datasets, train for 100 epochs
• Slow simulator (> 1 s – minutes per draw): pre-simulate 3 000–5 000 datasets, train for 100 epochs
• No simulator / pre-existing bank: use whatever is available; switch to disk training if it does not fit in memory

Online training is a refinement step — use it only after the first offline pass shows healthy diagnostics and you want to squeeze out more performance.

1. Define prior + observation model or simulation bank

• Implement prior and observation model and wrap them in a simulator
• Pre-simulate the pilot budget into a dict (using workflow.simulate(N)) for offline training
• If the simulator is external or proprietary, ensure simulations are already generated from the intended prior and data-generating process

1. Choose the architecture — this step is critical; getting it wrong ruins inference. See references/conditioning.md for the full conditioning logic and decision table.

• "Simple vector" means the observation is a single fixed-length feature vector whose element order is meaningful (e.g., 5 named sensor readings, a pre-computed summary statistic). Only then: route through inference_conditions with no summary network.
• Set-based / exchangeable data — If the simulator produces N observations that are exchangeable, the data is a set, not a vector. This includes: N i.i.d. draws, regression datasets with (x, y) pairs, repeated measurements, trial-level data, cross-sectional samples. Route through summary_variables with a SetTransformer. Never put this in `inference_conditions`.
• Time series — ordered sequences: route through summary_variables with TimeSeriesTransformer or TimeSeriesNetwork.
• Images as conditions / observations for parameter inference — route through summary_variables with ConvolutionalNetwork.
• Images as inferential targets — conditional image generation, spatial field generation, denoising, and other image-valued outputs require an image-capable diffusion inference network. Use bf.networks.DiffusionModel(subnet=...) with UNet, UViT, or ResidualUViT; see references/image-generation.md.
• A workflow can use both slots simultaneously. Fixed-length metadata (e.g., sample size N, scalar design variables) can go in inference_conditions while structured observations go in summary_variables.
• When in doubt, use a summary network. It is always safer to include one than to omit one; a summary network will always be needed if the data has more than one axis.

1. Build the workflow — Prefer bf.BasicWorkflow(...)

• Decide on which variables to auto-standardize. Prefer standardize="all" unless you have verfied that the simulator outputs are already in a good range for the networks.

1. Run simulation sanity checks — Before training, verify that simulated data look plausible and span the relevant range of real observations. Again, pay attention to what needs to be standardized.
2. Train the amortizer — First iteration always uses offline training for fast feedback:

• workflow.fit_offline(...) with the pre-simulated pilot budget (default first pass)
• workflow.fit_online(...) only as a refinement step after offline diagnostics look healthy, or when the user explicitly requests it
• workflow.fit_disk(...) if streaming simulations from disk

Always offer to run training in the terminal so the user can monitor progress interactively.

1. Diagnose in silico — Use held-out simulations with known ground truth using the workflow's built-in diagnostics: workflow.compute_default_diagnostics(...) for numerical results and workflow.plot_default_diagnostics(...) for visual diagnostics.
2. Amortized inference on real data — Use workflow.sample(...)
3. Posterior predictive checks (PPCs) — Re-simulate data from posterior samples and compare to the real data using model-specific test quantities
4. Write a report — Use references/reporting.md to generate a structured report outlining results and next steps.

Hard rules — MUST and NEVER

These rules are non-negotiable. Violating any of them will silently produce wrong results.

• MUST use `bf.Adapter()` for data routing. Build an explicit adapter chain with .as_set(), .constrain(), .concatenate(), etc. and pass adapter= to BasicWorkflow, as described in references/adapter.md. Do NOT do manual preprocessing — the adapter handles training and inference identically. The naming shorthand (inference_variables=, summary_variables= as kwargs to BasicWorkflow) is ONLY acceptable when the simulator output already has the exact shapes/dtypes the networks expect AND no parameter has bounded support. When in doubt, use an explicit adapter.
• MUST start with the Base network configuration from references/model-sizes.md. Scale up to Large or XL ONLY if diagnostics show poor recovery or calibration after sufficient training. Oversized networks waste compute and can hurt calibration on simple problems.
• MUST use `workflow.simulate(N)` to generate train/test data — not a Python for-loop over simulator(). The simulator returned by bf.make_simulator is a batched object; workflow.simulate(N) calls it efficiently and returns data in the format the workflow expects.
• MUST use `workflow.compute_default_diagnostics(test_data=...)` and `workflow.plot_default_diagnostics(test_data=...)` for in-silico diagnostics. NEVER hand-roll coverage, bias, or calibration computations — the built-in methods are correct, complete, and consistent with the house thresholds.
• For image-valued inference targets, follow `references/image-generation.md`. MUST use bf.networks.DiffusionModel(subnet=...) with UNet, UViT, or ResidualUViT — not the default low-dimensional setup. Conditions must be spatially concatenable with the image target (broadcast (B, D) to (B, H, W, D)). The standard diagnostic report does not apply; use visual sample grids instead.
• `workflow.sample()` returns original parameter names, NOT `"inference_variables"`. The adapter's reverse transform restores the original keys from the simulator (e.g., "alpha", "beta", "sigma"). Each parameter has shape (batch, num_samples) for scalars or (batch, num_samples, d) for vectors. NEVER index into "inference_variables" — that key does not exist in the output.
• MUST reuse the existing simulator functions for PPCs. NEVER re-implement the generative model by hand for posterior predictive checks. Loop over a subset of posterior draws (50 is a good default), indexing over the num_samples axis, and pass each draw through the simulator's forward model.
• MUST save `history.history` as JSON (not CSV, not a DataFrame — it is a plain dict). Then run scripts/inspect_training.py or call inspect_history() in-process.
• MUST pass `validation_data=` to all `fit_*` calls. For offline training, hold out ~300 simulations as a separate validation dict. For online training (refinement step only), pass an integer (e.g., validation_data=300) to auto-simulate.
• NEVER mix an explicit `adapter=` with the naming shorthand (inference_variables=, summary_variables=, inference_conditions= as kwargs). They are mutually exclusive. Passing both causes silent conflicts.
• NEVER flatten structured data into `inference_conditions`. Sets, time series, and images MUST go through summary_variables with an appropriate summary network.
• `workflow.plot_default_diagnostics()` ALWAYS returns a `dict[str, Figure]`. Iterate directly over .items() to save figures. Do not type-check or branch on the return type.
• NEVER skip in-silico diagnostics. Good training loss does not imply good inference.
• MUST generate `report.md` after every training + diagnostics run. Store all artifacts in <slug>/ (see references/reporting.md for naming and structure). Save all diagnostic figures with their standard names, save metrics.csv, and produce a self-contained markdown diagnostic report. If real data is available, include the optional real-data sections in the same report. For image-valued targets, skip the standard report and use visual sample grids instead.
• NEVER use `fit_online` as the first training pass unless the user explicitly requests it. The first iteration MUST use fit_offline with a pre-simulated pilot budget (10k sims for fast simulators, 3k–5k for slow ones) to maximize iteration speed. Online training is a refinement step for subsequent iterations.
• MUST offer to run training in the terminal so the user can monitor progress. Training scripts should be runnable standalone; do not silently execute long training runs without giving the user access to the live output.

Installation

Install BayesFlow and a backend. Prefer JAX unless the user has a strong reason to use PyTorch or TensorFlow.

pip install "bayesflow"

BayesFlow workflow template

import bayesflow as bf
import numpy as np
# --------------------------------------------------
# 1. Define prior + observation model
# --------------------------------------------------
def my_prior():
    theta = ...
    return {"parameters": theta}
def my_observation_model(parameters):
    x = ...
    return {"observables": x}
# bf.make_simulator returns a BATCHED simulator object.
# Use simulator.sample(batch_size) or workflow.simulate(N) — NEVER loop with simulator() in Python.
simulator = bf.make_simulator([my_prior, my_observation_model])
# --------------------------------------------------
# 2. Choose architecture
# --------------------------------------------------
# See references/conditioning.md for the full conditioning logic and decision table.
# See references/model-sizes.md for different configurations — always start with Base.
# Example summary network for set-based data (exchangeable observations):
summary_net = bf.networks.SetTransformer(...) 
inference_net = bf.networks.FlowMatching(...)
# alternatives:
# bf.networks.StableConsistencyModel() # faster sampler, less performant
# bf.networks.DiffusionModel() # slower sampler, good for image generation
# bf.networks.CouplingFlow(depth=4, transform="spline") # good-old normalizing flow
# --------------------------------------------------
# 3. Build the adapter (MUST use bf.Adapter)
# --------------------------------------------------
# See references/adapter.md for the full API and a step-by-step example.
# The adapter routes simulator output to the correct network slots and handles
# parameter constraints, set assembly, dtype conversion, and concatenation.
adapter = (
    bf.Adapter()
    .as_set(["observables"])              # (N,) -> (N, 1) for SetTransformer
    .constrain("parameters", lower=0)     # if parameters have bounded support
    .convert_dtype("float64", "float32")
    .concatenate(["observables"], into="summary_variables")
    .concatenate(["parameters"], into="inference_variables")
)
# --------------------------------------------------
# 4. Create results folder and workflow
# --------------------------------------------------
import os
results_dir = "<slug>"  # e.g., "churn-model" or "churn-model-v2" for iterations
os.makedirs(results_dir, exist_ok=True)