Integrating sbi-trained likelihoods into HSSM (NRE via ONNX)¶

This tutorial mirrors the structure of bayesflow_lre_integration.ipynb for the third major SBI library in the HSSM ecosystem: sbi (mackelab). It demonstrates how to:

1. Train a neural ratio estimator (NRE) on synthetic DDM simulations using sbi.
2. Export the trained estimator to ONNX via lanfactory.onnx.transform_sbi_to_onnx.
3. Load the ONNX file into HSSM exactly like any other LAN-style approximator and run MCMC inference.
4. Compare against HSSM's analytical DDM likelihood as the gold-standard posterior.

Why NRE only, not NLE/MNLE? Vanilla NLE with a MAF flow misbehaves on DDM data because rt is continuous but choice is discrete (∈ {−1, +1}). The flow treats choice as continuous, can't represent the support boundary rt > t_nd, and produces qualitatively wrong posteriors (we observed v ≈ 0.12 vs truth 0.5, with spurious bimodality on a). The correct sbi method is MNLE (Mixed Neural Likelihood Estimator), which splits x into discrete and continuous dims and models each properly. But MNLE's categorical lookup uses torch.searchsorted, which torch.onnx.export doesn't support and jaxonnxruntime lacks a handler for — a ~50-line upstream PR adding a SearchSorted handler to jaxonnxruntime would unlock both MNLE and NSF flows in one stroke.

Until then, NRE is the working integration path — it doesn't need to model density at all, just a classifier between joint and marginal pairs, which is robust to the discrete/continuous mixing.

Environment note: This tutorial requires both hssm and lanfactory[all] (which pulls sbi and nflows) in the same environment. JAX/flax/numpyro pins must be resolved jointly across the two packages.

Part 1 — Setup¶

In [1]:

# Enable x64 BEFORE any other JAX-touching import. sbi-exported ONNX graphs carry
# int64 shape/index tensors (typical for torch.onnx.export of normalizing flows);
# HSSM's ONNX-likelihood path relies on JAX x64 so those values are preserved
# exactly. x64 is also enabled automatically by `import hssm` (pytensor sets it
# from the default floatX="float64"); setting it explicitly here keeps the
# notebook self-contained and import-order-independent.
import jax

jax.config.update("jax_enable_x64", True)

from pathlib import Path

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

# lanfactory>=0.7.0 ships the sbi exporter; HSSM's `notebook` dependency group
# pins it (alongside sbi / nflows), so the import is direct.
from lanfactory.onnx import transform_sbi_to_onnx
from sbi.inference import NRE_A
from sbi.neural_nets import classifier_nn
from sbi.utils import BoxUniform
from ssms.basic_simulators.simulator import simulator
from torch import nn

import hssm

np.random.seed(0)
torch.manual_seed(0)

# Training budget for NRE_A — the last known-working configuration: 1M (theta, x)
# pairs from ssm-simulators (1 sample per theta) at the wider HSSM-default prior
# bounds, with a moderate-sized MLP classifier. Part 3 explains why this baseline
# is used rather than NRE_B / FCEmbedding / atomic contrastive.
N_TRAIN = 1_000_000
N_OBS = 500
NUM_EPOCHS = 300
STOP_AFTER_EPOCHS = 50
TRAINING_BATCH_SIZE = 500
HIDDEN_FEATURES = 100
# Smaller MCMC budget for verification runs. Together with target_accept=0.8
# and max_tree_depth=8 below, this keeps the MCMC step bounded.
MCMC_DRAWS = 500
MCMC_TUNE = 500
MCMC_CHAINS = 2

# Enable x64 BEFORE any other JAX-touching import. sbi-exported ONNX graphs carry # int64 shape/index tensors (typical for torch.onnx.export of normalizing flows); # HSSM's ONNX-likelihood path relies on JAX x64 so those values are preserved # exactly. x64 is also enabled automatically by `import hssm` (pytensor sets it # from the default floatX="float64"); setting it explicitly here keeps the # notebook self-contained and import-order-independent. import jax jax.config.update("jax_enable_x64", True) from pathlib import Path import arviz as az import matplotlib.pyplot as plt import numpy as np import pandas as pd import torch # lanfactory>=0.7.0 ships the sbi exporter; HSSM's `notebook` dependency group # pins it (alongside sbi / nflows), so the import is direct. from lanfactory.onnx import transform_sbi_to_onnx from sbi.inference import NRE_A from sbi.neural_nets import classifier_nn from sbi.utils import BoxUniform from ssms.basic_simulators.simulator import simulator from torch import nn import hssm np.random.seed(0) torch.manual_seed(0) # Training budget for NRE_A — the last known-working configuration: 1M (theta, x) # pairs from ssm-simulators (1 sample per theta) at the wider HSSM-default prior # bounds, with a moderate-sized MLP classifier. Part 3 explains why this baseline # is used rather than NRE_B / FCEmbedding / atomic contrastive. N_TRAIN = 1_000_000 N_OBS = 500 NUM_EPOCHS = 300 STOP_AFTER_EPOCHS = 50 TRAINING_BATCH_SIZE = 500 HIDDEN_FEATURES = 100 # Smaller MCMC budget for verification runs. Together with target_accept=0.8 # and max_tree_depth=8 below, this keeps the MCMC step bounded. MCMC_DRAWS = 500 MCMC_TUNE = 500 MCMC_CHAINS = 2

mlflow not available
mlflow not available

Part 2 — Simulate observed DDM data¶

We use the standard 4-parameter DDM (v, a, z, t) from ssm-simulators. The true parameters and parameter ranges below match the BayesFlow LRE tutorial so that posteriors from the two tutorials can be compared apples-to-apples.

In [2]:

DDM_PARAM_NAMES = ["v", "a", "z", "t"]
# Training prior matches HSSM's default bounds for model="ddm" with
# loglik_kind="approx_differentiable" — see hssm.defaults.default_model_config.
# This is important: if the training prior is narrower than HSSM's posterior
# can explore, MCMC will walk into parameter regions the surrogate never saw
# and extrapolate badly.
PRIOR_LOW = np.array([-3.0, 0.3, 0.0, 0.0], dtype=np.float32)
PRIOR_HIGH = np.array([3.0, 2.5, 1.0, 2.0], dtype=np.float32)
TRUE_THETA = np.array([0.5, 1.2, 0.5, 0.25], dtype=np.float32)

out = simulator(theta=TRUE_THETA[None, :], model="ddm", n_samples=N_OBS)
obs_data = pd.DataFrame(
    {
        "rt": out["rts"].squeeze().astype(np.float32),
        "response": out["choices"].squeeze().astype(np.float32),
    }
)
print(f"observed: {len(obs_data)} trials at true theta = {TRUE_THETA}")
obs_data.head()

DDM_PARAM_NAMES = ["v", "a", "z", "t"] # Training prior matches HSSM's default bounds for model="ddm" with # loglik_kind="approx_differentiable" — see hssm.defaults.default_model_config. # This is important: if the training prior is narrower than HSSM's posterior # can explore, MCMC will walk into parameter regions the surrogate never saw # and extrapolate badly. PRIOR_LOW = np.array([-3.0, 0.3, 0.0, 0.0], dtype=np.float32) PRIOR_HIGH = np.array([3.0, 2.5, 1.0, 2.0], dtype=np.float32) TRUE_THETA = np.array([0.5, 1.2, 0.5, 0.25], dtype=np.float32) out = simulator(theta=TRUE_THETA[None, :], model="ddm", n_samples=N_OBS) obs_data = pd.DataFrame( { "rt": out["rts"].squeeze().astype(np.float32), "response": out["choices"].squeeze().astype(np.float32), } ) print(f"observed: {len(obs_data)} trials at true theta = {TRUE_THETA}") obs_data.head()

observed: 500 trials at true theta = [0.5  1.2  0.5  0.25]

Out[2]:

	rt	response
0	2.528641	-1.0
1	2.497200	1.0
2	2.062349	1.0
3	6.837864	1.0
4	1.071579	1.0

In [3]:

fig, ax = plt.subplots(1, 1, figsize=(8, 4))
for resp, color in zip([-1, 1], ["C0", "C1"]):
    mask = obs_data["response"] == resp
    ax.hist(
        obs_data.loc[mask, "rt"],
        bins=40,
        alpha=0.6,
        label=f"choice={int(resp)}",
        color=color,
    )
ax.set_xlabel("RT (s)")
ax.set_ylabel("count")
ax.set_title("Observed RT histogram by choice")
ax.legend()
plt.tight_layout()
plt.show()

fig, ax = plt.subplots(1, 1, figsize=(8, 4)) for resp, color in zip([-1, 1], ["C0", "C1"]): mask = obs_data["response"] == resp ax.hist( obs_data.loc[mask, "rt"], bins=40, alpha=0.6, label=f"choice={int(resp)}", color=color, ) ax.set_xlabel("RT (s)") ax.set_ylabel("count") ax.set_title("Observed RT histogram by choice") ax.legend() plt.tight_layout() plt.show()

Part 3 — Train an sbi NRE_A classifier on DDM simulations¶

NRE_A (Hermans et al. 2020) learns a binary classifier that distinguishes joint (θ, x) pairs from marginal (θ', x) pairs (where θ' is drawn from the prior). The output logit equals log p(x | θ) − log p(x) up to a constant, so it serves directly as the HSSM log-likelihood for MCMC (the θ-independent constant drops out under MCMC's accept ratios).

This tutorial uses a validated NRE configuration:

• NRE_A (binary classifier, not contrastive)
• 1M (θ, x) training pairs (1 sample per θ)
• hidden_features = 100 (sbi default is 50)
• No embedding net on θ
• norm_layer = nn.Identity (jaxonnxruntime doesn't implement LayerNormalization, so the MLP norm layer is disabled)

We keep this baseline because a more ambitious configuration (NRE_B + atomic contrastive + multi-sample-per-θ + FCEmbedding + hidden_features=128) gave HSSM a near-constant log-likelihood at MCMC time — the chains explored the entire prior with no concentration. Scaling the classifier up from this robust starting point is left as follow-up work.

In [4]:

prior = BoxUniform(
    low=torch.from_numpy(PRIOR_LOW),
    high=torch.from_numpy(PRIOR_HIGH),
)
theta_train = prior.sample((N_TRAIN,))

# Batched ssm-simulators: theta of shape (N, 4) with n_samples=1 returns
# rts/choices of shape (N, 1). Much faster than a Python loop for large N.
sim = simulator(
    theta=theta_train.numpy().astype(np.float32),
    model="ddm",
    n_samples=1,
)
x_train = torch.from_numpy(
    np.stack([sim["rts"].squeeze(-1), sim["choices"].squeeze(-1)], axis=-1).astype(
        np.float32
    )
)
print(f"training set: theta={theta_train.shape}, x={x_train.shape}")

prior = BoxUniform( low=torch.from_numpy(PRIOR_LOW), high=torch.from_numpy(PRIOR_HIGH), ) theta_train = prior.sample((N_TRAIN,)) # Batched ssm-simulators: theta of shape (N, 4) with n_samples=1 returns # rts/choices of shape (N, 1). Much faster than a Python loop for large N. sim = simulator( theta=theta_train.numpy().astype(np.float32), model="ddm", n_samples=1, ) x_train = torch.from_numpy( np.stack([sim["rts"].squeeze(-1), sim["choices"].squeeze(-1)], axis=-1).astype( np.float32 ) ) print(f"training set: theta={theta_train.shape}, x={x_train.shape}")

training set: theta=torch.Size([1000000, 4]), x=torch.Size([1000000, 2])

In [5]:

from torch.utils.tensorboard import SummaryWriter

# User-configurable: where sbi writes tensorboard training logs. Default is
# ~/sbi_logs_tutorial — outside the HSSM repo so notebook re-runs don't leave
# artifacts in your working tree. Set to None below (and drop summary_writer
# from the NRE_A call) to use sbi's default of ./sbi-logs/ relative to cwd.
TUTORIAL_LOG_DIR = Path.home() / "sbi_logs_tutorial"
TUTORIAL_LOG_DIR.mkdir(parents=True, exist_ok=True)

# Build the classifier. LayerNorm is disabled because jaxonnxruntime
# doesn't implement LayerNormalization. No embedding net on theta in this
# baseline-revert iteration (see Part 3 markdown for the bisect context).
classifier_builder = classifier_nn(
    model="mlp",
    norm_layer=nn.Identity,
    hidden_features=HIDDEN_FEATURES,
)
inference_nre = NRE_A(
    prior=prior,
    classifier=classifier_builder,
    summary_writer=SummaryWriter(log_dir=str(TUTORIAL_LOG_DIR)),
)
classifier_nre = inference_nre.append_simulations(theta_train, x_train).train(
    training_batch_size=TRAINING_BATCH_SIZE,
    max_num_epochs=NUM_EPOCHS,
    stop_after_epochs=STOP_AFTER_EPOCHS,
)
classifier_nre.eval()
print("NRE_A training complete")

from torch.utils.tensorboard import SummaryWriter # User-configurable: where sbi writes tensorboard training logs. Default is # ~/sbi_logs_tutorial — outside the HSSM repo so notebook re-runs don't leave # artifacts in your working tree. Set to None below (and drop summary_writer # from the NRE_A call) to use sbi's default of ./sbi-logs/ relative to cwd. TUTORIAL_LOG_DIR = Path.home() / "sbi_logs_tutorial" TUTORIAL_LOG_DIR.mkdir(parents=True, exist_ok=True) # Build the classifier. LayerNorm is disabled because jaxonnxruntime # doesn't implement LayerNormalization. No embedding net on theta in this # baseline-revert iteration (see Part 3 markdown for the bisect context). classifier_builder = classifier_nn( model="mlp", norm_layer=nn.Identity, hidden_features=HIDDEN_FEATURES, ) inference_nre = NRE_A( prior=prior, classifier=classifier_builder, summary_writer=SummaryWriter(log_dir=str(TUTORIAL_LOG_DIR)), ) classifier_nre = inference_nre.append_simulations(theta_train, x_train).train( training_batch_size=TRAINING_BATCH_SIZE, max_num_epochs=NUM_EPOCHS, stop_after_epochs=STOP_AFTER_EPOCHS, ) classifier_nre.eval() print("NRE_A training complete")

/var/folders/rg/_j_2xg8n0qgbf0dvg2dyq1_r0000gn/T/ipykernel_75710/3901025436.py:13: UserWarning: Unknown kwargs passed to ClassifierConfig: {'norm_layer'}. These will be forwarded to the underlying builder. If this is unintentional, check for typos.
  classifier_builder = classifier_nn(
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/sbi/inference/trainers/nre/nre_base.py:84: FutureWarning: summary_writer is deprecated. Use tracker instead.
  super().__init__(
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/sbi/inference/trainers/base.py:296: UserWarning: Data has extreme outliers in dimension(s) [0] (beyond 10.0x IQR from quartiles). This may cause precision loss during z-scoring, where distinct values become indistinguishable. Consider removing outliers from your data or setting `z_score_x='none'` (though this may affect training).
  warn_if_invalid_for_zscoring(x)

 Training neural network. Epochs trained: 1

 Training neural network. Epochs trained: 2

 Training neural network. Epochs trained: 3

 Training neural network. Epochs trained: 4

 Training neural network. Epochs trained: 5

 Training neural network. Epochs trained: 6

 Training neural network. Epochs trained: 7

 Training neural network. Epochs trained: 8

 Training neural network. Epochs trained: 9

 Training neural network. Epochs trained: 10

 Training neural network. Epochs trained: 11

 Training neural network. Epochs trained: 12

 Training neural network. Epochs trained: 13

 Training neural network. Epochs trained: 14

 Training neural network. Epochs trained: 15

 Training neural network. Epochs trained: 16

 Training neural network. Epochs trained: 17

 Training neural network. Epochs trained: 18

 Training neural network. Epochs trained: 19

 Training neural network. Epochs trained: 20

 Training neural network. Epochs trained: 21

 Training neural network. Epochs trained: 22

 Training neural network. Epochs trained: 23

 Training neural network. Epochs trained: 24

 Training neural network. Epochs trained: 25

 Training neural network. Epochs trained: 26

 Training neural network. Epochs trained: 27

 Training neural network. Epochs trained: 28

 Training neural network. Epochs trained: 29

 Training neural network. Epochs trained: 30

 Training neural network. Epochs trained: 31

 Training neural network. Epochs trained: 32

 Training neural network. Epochs trained: 33

 Training neural network. Epochs trained: 34

 Training neural network. Epochs trained: 35

 Training neural network. Epochs trained: 36

 Training neural network. Epochs trained: 37

 Training neural network. Epochs trained: 38

 Training neural network. Epochs trained: 39

 Training neural network. Epochs trained: 40

 Training neural network. Epochs trained: 41

 Training neural network. Epochs trained: 42

 Training neural network. Epochs trained: 43

 Training neural network. Epochs trained: 44

 Training neural network. Epochs trained: 45

 Training neural network. Epochs trained: 46

 Training neural network. Epochs trained: 47

 Training neural network. Epochs trained: 48

 Training neural network. Epochs trained: 49

 Training neural network. Epochs trained: 50

 Training neural network. Epochs trained: 51

 Training neural network. Epochs trained: 52

 Training neural network. Epochs trained: 53

 Training neural network. Epochs trained: 54

 Training neural network. Epochs trained: 55

 Training neural network. Epochs trained: 56

 Training neural network. Epochs trained: 57

 Training neural network. Epochs trained: 58

 Training neural network. Epochs trained: 59

 Training neural network. Epochs trained: 60

 Training neural network. Epochs trained: 61

 Training neural network. Epochs trained: 62

 Training neural network. Epochs trained: 63

 Training neural network. Epochs trained: 64

 Training neural network. Epochs trained: 65

 Training neural network. Epochs trained: 66

 Training neural network. Epochs trained: 67

 Training neural network. Epochs trained: 68

 Training neural network. Epochs trained: 69

 Training neural network. Epochs trained: 70

 Training neural network. Epochs trained: 71

 Training neural network. Epochs trained: 72

 Training neural network. Epochs trained: 73

 Training neural network. Epochs trained: 74

 Training neural network. Epochs trained: 75

 Training neural network. Epochs trained: 76

 Training neural network. Epochs trained: 77

 Training neural network. Epochs trained: 78

 Training neural network. Epochs trained: 79

 Training neural network. Epochs trained: 80

 Training neural network. Epochs trained: 81

 Training neural network. Epochs trained: 82

 Training neural network. Epochs trained: 83

 Training neural network. Epochs trained: 84

 Training neural network. Epochs trained: 85

 Training neural network. Epochs trained: 86

 Training neural network. Epochs trained: 87

 Training neural network. Epochs trained: 88

 Training neural network. Epochs trained: 89

 Training neural network. Epochs trained: 90

 Training neural network. Epochs trained: 91

 Training neural network. Epochs trained: 92

 Training neural network. Epochs trained: 93

 Training neural network. Epochs trained: 94

 Training neural network. Epochs trained: 95

 Training neural network. Epochs trained: 96

 Training neural network. Epochs trained: 97

 Training neural network. Epochs trained: 98

 Training neural network. Epochs trained: 99

 Training neural network. Epochs trained: 100

 Training neural network. Epochs trained: 101

 Training neural network. Epochs trained: 102

 Training neural network. Epochs trained: 103

 Training neural network. Epochs trained: 104

 Training neural network. Epochs trained: 105

 Training neural network. Epochs trained: 106

 Training neural network. Epochs trained: 107

 Training neural network. Epochs trained: 108
 Neural network successfully converged after 108 epochs.NRE_A training complete

Part 4 — Export the trained NRE to ONNX¶

The exporter wraps the classifier's forward(theta, x) logit as the HSSM log-likelihood. No Jacobian correction is needed — ratios are invariant to the z-score standardization sbi applies internally.

In [6]:

# User-configurable: where the .onnx file lands. Default is outside the HSSM
# repo so notebook re-runs don't pollute the working tree.
# Override examples:
#   ARTIFACT_DIR = Path("/path/to/my/project/onnx")          # keep nearby
#   ARTIFACT_DIR = Path(tempfile.mkdtemp())                  # ephemeral
ARTIFACT_DIR = Path.home() / "sbi_onnx_tutorial"
ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
nre_onnx_path = ARTIFACT_DIR / "ddm_nre.onnx"

transform_sbi_to_onnx(
    classifier_nre,
    str(nre_onnx_path),
    mode="nre",
    example_theta_dim=4,
    example_x_dim=2,
)
print(f"exported NRE: {nre_onnx_path}  ({nre_onnx_path.stat().st_size:,} bytes)")

# User-configurable: where the .onnx file lands. Default is outside the HSSM # repo so notebook re-runs don't pollute the working tree. # Override examples: # ARTIFACT_DIR = Path("/path/to/my/project/onnx") # keep nearby # ARTIFACT_DIR = Path(tempfile.mkdtemp()) # ephemeral ARTIFACT_DIR = Path.home() / "sbi_onnx_tutorial" ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) nre_onnx_path = ARTIFACT_DIR / "ddm_nre.onnx" transform_sbi_to_onnx( classifier_nre, str(nre_onnx_path), mode="nre", example_theta_dim=4, example_x_dim=2, ) print(f"exported NRE: {nre_onnx_path} ({nre_onnx_path.stat().st_size:,} bytes)")

exported NRE: /Users/afengler/sbi_onnx_tutorial/ddm_nre.onnx  (48,270 bytes)

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/sbi/neural_nets/estimators/base.py:137: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if input_shape != exp_input_shape:
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/sbi/neural_nets/estimators/base.py:107: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if condition_shape != exp_condition_shape:
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/sbi/neural_nets/ratio_estimators.py:98: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if theta_prefix != x_prefix:

Part 4b — Pre-MCMC verification: is the trained classifier any good?¶

Before paying the multi-minute MCMC cost, sanity-check two things:

1. Logit sweep across θ-space. Hold three θ dimensions at their true values and sweep the fourth across its prior range, plotting the summed classifier log-ratio on the observed data. A well-trained NRE shows a sharp peak near the true value with tens-to-hundreds of log units of vertical range. A flat curve (< 5 log units of range) is the smoking gun for "classifier collapsed" — MCMC will produce a posterior equal to the prior, no point continuing.

2. ONNX export round-trip. Compare classifier_nre(theta, x).item() to the exported ONNX file evaluated through onnxruntime on the same input. If they disagree, the export has introduced a bug and MCMC won't be running the model you think it is.

These cells use the in-memory classifier_nre from Part 3 and the exported file from Part 4 — they're cheap, no MCMC required.

In [7]:

# Diagnostic: how much does the NRE logit change as we sweep each theta dim?
sweep_pts = 50
sweep_results = {}
obs_x_t = torch.from_numpy(obs_data[["rt", "response"]].values.astype(np.float32))
theta_center = torch.tensor(TRUE_THETA, dtype=torch.float32)
for dim, name in enumerate(DDM_PARAM_NAMES):
    sweep = torch.linspace(PRIOR_LOW[dim], PRIOR_HIGH[dim], sweep_pts)
    logits = []
    for v in sweep:
        theta = theta_center.clone()
        theta[dim] = v
        theta_row = theta.unsqueeze(0).repeat(len(obs_x_t), 1)
        with torch.no_grad():
            logits.append(classifier_nre(theta_row, obs_x_t).sum().item())
    sweep_results[name] = (sweep.numpy(), np.array(logits))

fig, axes = plt.subplots(1, 4, figsize=(16, 3.5))
for ax, name in zip(axes, DDM_PARAM_NAMES):
    th, lp = sweep_results[name]
    ax.plot(th, lp - lp.max(), "C0-", linewidth=2)
    ax.axvline(
        TRUE_THETA[DDM_PARAM_NAMES.index(name)],
        color="red",
        linestyle="--",
        linewidth=2,
        label="true θ",
    )
    ax.set_xlabel(name)
    ax.set_ylabel("Δ summed log-ratio (= 0 at max)")
    ax.set_title(f"sweep over {name}\n(vertical range = {np.ptp(lp):.2f})")
    ax.legend(fontsize=8)
fig.suptitle(
    "Trained NRE classifier: log-ratio along each θ axis (others held at truth)",
    y=1.02,
)
plt.tight_layout()
plt.show()

# Sanity print: vertical range per dim. < ~5 log units is bad.
print("\nPer-dim vertical range (max log-ratio − min log-ratio):")
for name in DDM_PARAM_NAMES:
    _, lp = sweep_results[name]
    print(f"  {name}: {np.ptp(lp):.2f}")
print("\nInterpretation:")
print("  > 50 log units : strong discriminative signal, classifier well-trained.")
print("  10–50          : moderate signal; should give a usable posterior.")
print("  < 5            : near-flat; classifier collapsed to uninformative.")

# Diagnostic: how much does the NRE logit change as we sweep each theta dim? sweep_pts = 50 sweep_results = {} obs_x_t = torch.from_numpy(obs_data[["rt", "response"]].values.astype(np.float32)) theta_center = torch.tensor(TRUE_THETA, dtype=torch.float32) for dim, name in enumerate(DDM_PARAM_NAMES): sweep = torch.linspace(PRIOR_LOW[dim], PRIOR_HIGH[dim], sweep_pts) logits = [] for v in sweep: theta = theta_center.clone() theta[dim] = v theta_row = theta.unsqueeze(0).repeat(len(obs_x_t), 1) with torch.no_grad(): logits.append(classifier_nre(theta_row, obs_x_t).sum().item()) sweep_results[name] = (sweep.numpy(), np.array(logits)) fig, axes = plt.subplots(1, 4, figsize=(16, 3.5)) for ax, name in zip(axes, DDM_PARAM_NAMES): th, lp = sweep_results[name] ax.plot(th, lp - lp.max(), "C0-", linewidth=2) ax.axvline( TRUE_THETA[DDM_PARAM_NAMES.index(name)], color="red", linestyle="--", linewidth=2, label="true θ", ) ax.set_xlabel(name) ax.set_ylabel("Δ summed log-ratio (= 0 at max)") ax.set_title(f"sweep over {name}\n(vertical range = {np.ptp(lp):.2f})") ax.legend(fontsize=8) fig.suptitle( "Trained NRE classifier: log-ratio along each θ axis (others held at truth)", y=1.02, ) plt.tight_layout() plt.show() # Sanity print: vertical range per dim. < ~5 log units is bad. print("\nPer-dim vertical range (max log-ratio − min log-ratio):") for name in DDM_PARAM_NAMES: _, lp = sweep_results[name] print(f" {name}: {np.ptp(lp):.2f}") print("\nInterpretation:") print(" > 50 log units : strong discriminative signal, classifier well-trained.") print(" 10–50 : moderate signal; should give a usable posterior.") print(" < 5 : near-flat; classifier collapsed to uninformative.")

Per-dim vertical range (max log-ratio − min log-ratio):
  v: 3287.68
  a: 3186.27
  z: 1806.79
  t: 174149.97

Interpretation:
  > 50 log units : strong discriminative signal, classifier well-trained.
  10–50          : moderate signal; should give a usable posterior.
  < 5            : near-flat; classifier collapsed to uninformative.

In [8]:

# Diagnostic: does the exported ONNX match the torch classifier?
import onnxruntime as _ort

_sess = _ort.InferenceSession(str(nre_onnx_path))
_input_name = _sess.get_inputs()[0].name
_test_theta = torch.tensor([[0.5, 1.2, 0.5, 0.25]], dtype=torch.float32)
_test_x = torch.tensor([[0.5, 1.0]], dtype=torch.float32)
_combined = (
    torch.cat([_test_theta, _test_x], dim=-1).squeeze(0).numpy().astype(np.float32)
)

with torch.no_grad():
    _y_torch = float(classifier_nre(_test_theta, _test_x).item())
_y_ort = float(_sess.run(None, {_input_name: _combined})[0])

print(f"torch logit at (θ=true, x=(0.5, +1)):  {_y_torch:+.5f}")
print(f"ORT logit at (θ=true, x=(0.5, +1)):    {_y_ort:+.5f}")
print(f"|Δ|: {abs(_y_torch - _y_ort):.2e}")
print()
if abs(_y_torch - _y_ort) < 1e-4:
    print("→ Export is faithful; any pathology is in the trained classifier itself.")
else:
    print(
        "→ Export disagrees with torch by more than 1e-4 — the exporter is dropping"
        " information on the way through ONNX, regardless of training quality."
    )

# Diagnostic: does the exported ONNX match the torch classifier? import onnxruntime as _ort _sess = _ort.InferenceSession(str(nre_onnx_path)) _input_name = _sess.get_inputs()[0].name _test_theta = torch.tensor([[0.5, 1.2, 0.5, 0.25]], dtype=torch.float32) _test_x = torch.tensor([[0.5, 1.0]], dtype=torch.float32) _combined = ( torch.cat([_test_theta, _test_x], dim=-1).squeeze(0).numpy().astype(np.float32) ) with torch.no_grad(): _y_torch = float(classifier_nre(_test_theta, _test_x).item()) _y_ort = float(_sess.run(None, {_input_name: _combined})[0]) print(f"torch logit at (θ=true, x=(0.5, +1)): {_y_torch:+.5f}") print(f"ORT logit at (θ=true, x=(0.5, +1)): {_y_ort:+.5f}") print(f"|Δ|: {abs(_y_torch - _y_ort):.2e}") print() if abs(_y_torch - _y_ort) < 1e-4: print("→ Export is faithful; any pathology is in the trained classifier itself.") else: print( "→ Export disagrees with torch by more than 1e-4 — the exporter is dropping" " information on the way through ONNX, regardless of training quality." )

torch logit at (θ=true, x=(0.5, +1)):  +0.86791
ORT logit at (θ=true, x=(0.5, +1)):    +0.86791
|Δ|: 1.79e-07

→ Export is faithful; any pathology is in the trained classifier itself.

Part 5 — High-level integration via hssm.HSSM()¶

HSSM's loglik_kind="approx_differentiable" path consumes the .onnx file identically to a LAN-trained network. With model="ddm" HSSM already knows the parameter list and response columns; we just hand it the file.

In [9]:

model_nre = hssm.HSSM(
    data=obs_data,
    model="ddm",
    loglik_kind="approx_differentiable",
    loglik=str(nre_onnx_path),
    p_outlier=0,
)
print(model_nre)

model_nre = hssm.HSSM( data=obs_data, model="ddm", loglik_kind="approx_differentiable", loglik=str(nre_onnx_path), p_outlier=0, ) print(model_nre)

Model initialized successfully.

Hierarchical Sequential Sampling Model
Model: ddm

Response variable: rt,response
Likelihood: approx_differentiable
Observations: 500

Parameters:

v:
    Prior: Uniform(lower: -3.0, upper: 3.0)
    Explicit bounds: (-3.0, 3.0)

a:
    Prior: Uniform(lower: 0.3, upper: 2.5)
    Explicit bounds: (0.3, 2.5)

z:
    Prior: Uniform(lower: 0.0, upper: 1.0)
    Explicit bounds: (0.0, 1.0)

t:
    Prior: HalfNormal(sigma: 2.0)
    Explicit bounds: (0.0, 2.0)

In [10]:

# Verification-budget MCMC. target_accept=0.8 (was 0.9) and max_tree_depth=8
# (caps NUTS at 256 leapfrog steps per draw instead of 1024) bound the per-step
# cost so a pathological surrogate geometry can't produce a multi-hour run.
# max_tree_depth is forwarded to the numpyro NUTS kernel via pymc's
# sample_jax_nuts(nuts_kwargs=...), so it nests under "nuts_kwargs".
# progressbar=True lets you actually see chain progress as it goes.
idata_nre = model_nre.sample(
    sampler="numpyro",
    draws=MCMC_DRAWS,
    tune=MCMC_TUNE,
    chains=MCMC_CHAINS,
    target_accept=0.8,
    progressbar=True,
    nuts_sampler_kwargs={"nuts_kwargs": {"max_tree_depth": 8}},
)

# Verification-budget MCMC. target_accept=0.8 (was 0.9) and max_tree_depth=8 # (caps NUTS at 256 leapfrog steps per draw instead of 1024) bound the per-step # cost so a pathological surrogate geometry can't produce a multi-hour run. # max_tree_depth is forwarded to the numpyro NUTS kernel via pymc's # sample_jax_nuts(nuts_kwargs=...), so it nests under "nuts_kwargs". # progressbar=True lets you actually see chain progress as it goes. idata_nre = model_nre.sample( sampler="numpyro", draws=MCMC_DRAWS, tune=MCMC_TUNE, chains=MCMC_CHAINS, target_accept=0.8, progressbar=True, nuts_sampler_kwargs={"nuts_kwargs": {"max_tree_depth": 8}}, )

Using default initvals. 

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pymc/sampling/jax.py:475: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

  0%|          | 0/1000 [00:00<?, ?it/s]

warmup:   0%|          | 1/1000 [00:00<05:32,  3.00it/s, 1 steps of size 2.34e+00. acc. prob=0.00]

warmup:   4%|▍         | 38/1000 [00:00<00:08, 112.42it/s, 15 steps of size 1.19e-02. acc. prob=0.73]

warmup:   9%|▉         | 93/1000 [00:00<00:03, 241.03it/s, 23 steps of size 2.29e-02. acc. prob=0.77]

warmup:  14%|█▎        | 136/1000 [00:00<00:02, 296.01it/s, 15 steps of size 1.81e-01. acc. prob=0.78]

warmup:  18%|█▊        | 178/1000 [00:00<00:02, 331.93it/s, 15 steps of size 5.39e-01. acc. prob=0.78]

warmup:  23%|██▎       | 234/1000 [00:00<00:01, 397.50it/s, 11 steps of size 4.47e-01. acc. prob=0.78]

warmup:  28%|██▊       | 283/1000 [00:00<00:01, 420.04it/s, 31 steps of size 2.43e-01. acc. prob=0.78]

warmup:  34%|███▍      | 343/1000 [00:01<00:01, 471.53it/s, 7 steps of size 1.01e+00. acc. prob=0.79] 

warmup:  42%|████▏     | 415/1000 [00:01<00:01, 544.78it/s, 7 steps of size 5.88e-01. acc. prob=0.79]

warmup:  47%|████▋     | 472/1000 [00:01<00:00, 530.52it/s, 7 steps of size 4.52e-01. acc. prob=0.79]

sample:  54%|█████▍    | 543/1000 [00:01<00:00, 581.87it/s, 15 steps of size 4.76e-01. acc. prob=0.91]

sample:  62%|██████▏   | 615/1000 [00:01<00:00, 621.95it/s, 7 steps of size 4.76e-01. acc. prob=0.88] 

sample:  68%|██████▊   | 679/1000 [00:01<00:00, 625.63it/s, 7 steps of size 4.76e-01. acc. prob=0.87]

sample:  75%|███████▍  | 746/1000 [00:01<00:00, 636.53it/s, 7 steps of size 4.76e-01. acc. prob=0.88]

sample:  81%|████████  | 811/1000 [00:01<00:00, 639.40it/s, 7 steps of size 4.76e-01. acc. prob=0.88]

sample:  88%|████████▊ | 876/1000 [00:01<00:00, 626.31it/s, 7 steps of size 4.76e-01. acc. prob=0.88]

sample:  94%|█████████▍| 945/1000 [00:01<00:00, 642.88it/s, 15 steps of size 4.76e-01. acc. prob=0.88]

sample: 100%|██████████| 1000/1000 [00:02<00:00, 486.96it/s, 7 steps of size 4.76e-01. acc. prob=0.88]

  0%|          | 0/1000 [00:00<?, ?it/s]

warmup:   3%|▎         | 32/1000 [00:00<00:03, 311.82it/s, 15 steps of size 2.22e-02. acc. prob=0.73]

warmup:   7%|▋         | 73/1000 [00:00<00:02, 364.74it/s, 11 steps of size 1.44e-02. acc. prob=0.76]

warmup:  11%|█         | 111/1000 [00:00<00:02, 368.34it/s, 15 steps of size 2.93e-01. acc. prob=0.77]

warmup:  16%|█▋        | 164/1000 [00:00<00:01, 429.70it/s, 31 steps of size 1.37e-01. acc. prob=0.77]

warmup:  22%|██▏       | 223/1000 [00:00<00:01, 486.12it/s, 5 steps of size 1.80e-01. acc. prob=0.78] 

warmup:  28%|██▊       | 278/1000 [00:00<00:01, 507.42it/s, 7 steps of size 2.76e-01. acc. prob=0.78]

warmup:  34%|███▎      | 336/1000 [00:00<00:01, 528.33it/s, 27 steps of size 4.49e-01. acc. prob=0.78]

warmup:  41%|████      | 406/1000 [00:00<00:01, 581.36it/s, 7 steps of size 4.78e-01. acc. prob=0.79] 

warmup:  47%|████▋     | 473/1000 [00:00<00:00, 607.99it/s, 31 steps of size 1.63e-01. acc. prob=0.78]

sample:  53%|█████▎    | 534/1000 [00:01<00:00, 579.18it/s, 15 steps of size 3.97e-01. acc. prob=0.90]

sample:  59%|█████▉    | 594/1000 [00:01<00:00, 581.60it/s, 11 steps of size 3.97e-01. acc. prob=0.92]

sample:  65%|██████▌   | 653/1000 [00:01<00:00, 569.46it/s, 7 steps of size 3.97e-01. acc. prob=0.92] 

sample:  71%|███████   | 711/1000 [00:01<00:00, 561.89it/s, 3 steps of size 3.97e-01. acc. prob=0.91]

sample:  77%|███████▋  | 768/1000 [00:01<00:00, 558.85it/s, 7 steps of size 3.97e-01. acc. prob=0.91]

sample:  83%|████████▎ | 831/1000 [00:01<00:00, 576.40it/s, 11 steps of size 3.97e-01. acc. prob=0.91]

sample:  89%|████████▉ | 889/1000 [00:01<00:00, 574.81it/s, 11 steps of size 3.97e-01. acc. prob=0.91]

sample:  95%|█████████▌| 951/1000 [00:01<00:00, 585.31it/s, 7 steps of size 3.97e-01. acc. prob=0.91] 

sample: 100%|██████████| 1000/1000 [00:01<00:00, 547.27it/s, 7 steps of size 3.97e-01. acc. prob=0.91]

We recommend running at least 4 chains for robust computation of convergence diagnostics

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

  0%|          | 0/1000 [00:00<?, ?it/s]

 48%|████▊     | 477/1000 [00:00<00:00, 4768.72it/s]

100%|██████████| 1000/1000 [00:00<00:00, 5828.94it/s]

In [11]:

summary_nre = az.summary(idata_nre, var_names=DDM_PARAM_NAMES)
summary_nre

summary_nre = az.summary(idata_nre, var_names=DDM_PARAM_NAMES) summary_nre

Out[11]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
v	0.571	0.051	0.477	0.673	0.002	0.002	457.0	583.0	1.01
a	1.246	0.027	1.194	1.295	0.001	0.001	498.0	482.0	1.01
z	0.478	0.020	0.443	0.514	0.001	0.001	359.0	397.0	1.01
t	0.213	0.023	0.173	0.256	0.001	0.001	423.0	740.0	1.01

In [12]:

az.plot_trace(idata_nre, var_names=DDM_PARAM_NAMES)
plt.tight_layout()
plt.show()

az.plot_trace(idata_nre, var_names=DDM_PARAM_NAMES) plt.tight_layout() plt.show()

Part 5b — Validation: does the NRE classifier peak at the posterior mode?¶

For NRE the logit forward(theta, x) equals log p(x | θ) − log p(x) up to a constant. The θ-independent term cancels when we compare two θ values on the same data, so the summed logit over trials tells us which θ the classifier finds more consistent with the data. A well-trained NRE should rank the posterior region at or above θ values far from it — we confirm that below. A small logit gap between the truth and the posterior mean is expected rather than a red flag: at N=500 trials the posterior has real width, so the mode need not sit exactly on the true θ.

In [13]:

posterior_mean_nre = {p: float(idata_nre.posterior[p].mean()) for p in DDM_PARAM_NAMES}
obs_x_t = torch.from_numpy(obs_data[["rt", "response"]].values.astype(np.float32))


def total_logit(classifier, theta_dict):
    """Sum log-ratio logit over all observed trials at a single theta."""
    theta_row = torch.tensor(
        [[theta_dict[p] for p in DDM_PARAM_NAMES]], dtype=torch.float32
    ).repeat(len(obs_x_t), 1)
    with torch.no_grad():
        return classifier(theta_row, obs_x_t).sum().item()


lt_true = total_logit(classifier_nre, dict(zip(DDM_PARAM_NAMES, TRUE_THETA)))
lt_mean = total_logit(classifier_nre, posterior_mean_nre)

print(f"NRE total logit at true theta:      {lt_true:+.2f}")
print(f"NRE total logit at posterior mean:  {lt_mean:+.2f}")
print(f"Δ (mean − true):                    {lt_mean - lt_true:+.2f}")
print()
if lt_mean > lt_true + 5.0:
    print("→ NRE itself prefers the wrong θ by a large margin.")
    print("  Diagnosis: TRAINING QUALITY.")
elif lt_mean > lt_true:
    print("→ NRE mildly prefers the posterior mean over the truth.")
    print("  Could be marginal-vs-joint, mild miscalibration, or sampling.")
else:
    print("→ NRE prefers the truth; the wrong posterior mean is a sampling")
    print("  artifact (priors / init / mixing), not a training issue.")
print()
print(f"Posterior mean: {posterior_mean_nre}")
print(f"True theta:     {dict(zip(DDM_PARAM_NAMES, TRUE_THETA.tolist()))}")

posterior_mean_nre = {p: float(idata_nre.posterior[p].mean()) for p in DDM_PARAM_NAMES} obs_x_t = torch.from_numpy(obs_data[["rt", "response"]].values.astype(np.float32)) def total_logit(classifier, theta_dict): """Sum log-ratio logit over all observed trials at a single theta.""" theta_row = torch.tensor( [[theta_dict[p] for p in DDM_PARAM_NAMES]], dtype=torch.float32 ).repeat(len(obs_x_t), 1) with torch.no_grad(): return classifier(theta_row, obs_x_t).sum().item() lt_true = total_logit(classifier_nre, dict(zip(DDM_PARAM_NAMES, TRUE_THETA))) lt_mean = total_logit(classifier_nre, posterior_mean_nre) print(f"NRE total logit at true theta: {lt_true:+.2f}") print(f"NRE total logit at posterior mean: {lt_mean:+.2f}") print(f"Δ (mean − true): {lt_mean - lt_true:+.2f}") print() if lt_mean > lt_true + 5.0: print("→ NRE itself prefers the wrong θ by a large margin.") print(" Diagnosis: TRAINING QUALITY.") elif lt_mean > lt_true: print("→ NRE mildly prefers the posterior mean over the truth.") print(" Could be marginal-vs-joint, mild miscalibration, or sampling.") else: print("→ NRE prefers the truth; the wrong posterior mean is a sampling") print(" artifact (priors / init / mixing), not a training issue.") print() print(f"Posterior mean: {posterior_mean_nre}") print(f"True theta: {dict(zip(DDM_PARAM_NAMES, TRUE_THETA.tolist()))}")

NRE total logit at true theta:      +133.90
NRE total logit at posterior mean:  +135.86
Δ (mean − true):                    +1.96

→ NRE mildly prefers the posterior mean over the truth.
  Could be marginal-vs-joint, mild miscalibration, or sampling.

Posterior mean: {'v': 0.5712079136003048, 'a': 1.2460115960254583, 'z': 0.47760079236414615, 't': 0.21269431966332117}
True theta:     {'v': 0.5, 'a': 1.2000000476837158, 'z': 0.5, 't': 0.25}

Part 6 — Ground-truth posterior via HSSM's analytical DDM¶

HSSM ships a closed-form analytical likelihood for the standard DDM (loglik_kind="analytical", the Navarro & Fuss form). Running MCMC against this likelihood on the same observed data produces what we should treat as the gold-standard posterior for this model + dataset: any deviation of the sbi-NRE marginals from these is approximation error in the neural surrogate, not intrinsic posterior width.

The analytical likelihood uses slightly different parameter bounds (a, t unbounded above; otherwise the same) but on the observed data the posterior concentrates regardless of the wider bound.

In [14]:

model_analytical = hssm.HSSM(
    data=obs_data,
    model="ddm",
    loglik_kind="analytical",
    p_outlier=0,
)
idata_analytical = model_analytical.sample(
    sampler="numpyro",
    draws=MCMC_DRAWS,
    tune=MCMC_TUNE,
    chains=MCMC_CHAINS,
    target_accept=0.9,
    progressbar=False,
)
summary_analytical = az.summary(idata_analytical, var_names=DDM_PARAM_NAMES)
summary_analytical

model_analytical = hssm.HSSM( data=obs_data, model="ddm", loglik_kind="analytical", p_outlier=0, ) idata_analytical = model_analytical.sample( sampler="numpyro", draws=MCMC_DRAWS, tune=MCMC_TUNE, chains=MCMC_CHAINS, target_accept=0.9, progressbar=False, ) summary_analytical = az.summary(idata_analytical, var_names=DDM_PARAM_NAMES) summary_analytical

Model initialized successfully.

Using default initvals. 

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pymc/sampling/jax.py:475: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

We recommend running at least 4 chains for robust computation of convergence diagnostics

  0%|          | 0/1000 [00:00<?, ?it/s]

100%|██████████| 1000/1000 [00:00<00:00, 11197.74it/s]

Out[14]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
v	0.479	0.047	0.384	0.564	0.002	0.001	606.0	561.0	1.01
a	1.238	0.028	1.188	1.297	0.001	0.001	618.0	660.0	1.01
z	0.484	0.016	0.455	0.514	0.001	0.000	651.0	716.0	1.00
t	0.237	0.020	0.200	0.272	0.001	0.001	634.0	787.0	1.00

Part 7 — Posterior comparison: analytical vs sbi NRE¶

The keystone comparison. The analytical posterior is the gold standard for this model + data; the sbi-NRE marginals are the approximation we built. Distance between NRE and analytical is the surrogate's approximation error; distance between the analytical posterior and the true θ is intrinsic posterior width (the data simply isn't infinite at N=500 trials).

This notebook reuses the exact TRUE_THETA, N_OBS, and prior ranges from bayesflow_lre_integration.ipynb, so you can run the two notebooks side by side to compare the sbi-NRE and BayesFlow-LRE surrogates on identical data.

In [15]:

fig, axes = plt.subplots(1, 4, figsize=(16, 4))
for ax, name, true_val in zip(axes, DDM_PARAM_NAMES, TRUE_THETA):
    samples_ana = idata_analytical.posterior[name].values.flatten()
    samples_nre = idata_nre.posterior[name].values.flatten()
    ax.hist(
        samples_ana,
        bins=30,
        alpha=0.5,
        label="analytical (truth)",
        color="C2",
        density=True,
    )
    ax.hist(samples_nre, bins=30, alpha=0.5, label="sbi NRE", color="C1", density=True)
    ax.axvline(true_val, color="red", linestyle="--", linewidth=2, label="true θ")
    ax.set_xlabel(name)
    ax.set_title(f"posterior over {name}")
    ax.legend(fontsize=8)
fig.suptitle("DDM posterior recovery: analytical (gold) vs sbi NRE", y=1.02)
plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 4, figsize=(16, 4)) for ax, name, true_val in zip(axes, DDM_PARAM_NAMES, TRUE_THETA): samples_ana = idata_analytical.posterior[name].values.flatten() samples_nre = idata_nre.posterior[name].values.flatten() ax.hist( samples_ana, bins=30, alpha=0.5, label="analytical (truth)", color="C2", density=True, ) ax.hist(samples_nre, bins=30, alpha=0.5, label="sbi NRE", color="C1", density=True) ax.axvline(true_val, color="red", linestyle="--", linewidth=2, label="true θ") ax.set_xlabel(name) ax.set_title(f"posterior over {name}") ax.legend(fontsize=8) fig.suptitle("DDM posterior recovery: analytical (gold) vs sbi NRE", y=1.02) plt.tight_layout() plt.show()

Summary and deferred work¶

We trained an sbi NRE_A classifier on synthetic DDM data — a moderate MLP (hidden_features=100, norm_layer=nn.Identity so the graph stays ONNX-exportable, no embedding net on the parameters) on 1M (θ, x) pairs (one simulation per θ) — exported it to ONNX via lanfactory.onnx.transform_sbi_to_onnx, and ran MCMC through HSSM's existing loglik_kind="approx_differentiable" pipeline. The resulting posterior is compared against HSSM's analytical DDM posterior on the same data — that's our gold-standard reference.

A more ambitious configuration — NRE_B with atomic contrastive estimation (num_atoms=20), an FCEmbedding on θ, and a larger classifier — collapsed to a near-constant log-likelihood (the chains explored the whole prior with no concentration). This tutorial uses the simpler, validated NRE_A/MLP baseline (see Part 3); scaling it up from there is left as follow-up work.

Why not NLE in this tutorial?

We originally planned an NLE section too. Vanilla NLE with a MAF flow produces qualitatively wrong posteriors on DDM data because rt is continuous but choice is discrete (∈ {−1, +1}); the flow can't represent that structure or the hard support boundary rt > t_nd. The correct sbi method is MNLE (Mixed Neural Likelihood Estimator), which factorizes p(rt, choice | θ) = p(choice | θ) · p(rt | choice, θ). But MNLE's CategoricalMassEstimator uses torch.searchsorted for value-to-index lookup, which torch.onnx.export does not support — blocking the ONNX export path until a SearchSorted ONNX-op handler is added to jaxonnxruntime. The same gap blocks Neural Spline Flows; a single ~50-line upstream PR adding that SearchSorted handler would unlock both NSF flows and MNLE in one stroke.

Where to look next

• LANfactory's Exporting sbi Models guide — supported-architecture matrix, known constraints, troubleshooting.
• The BayesFlow LRE tutorial (bayesflow_lre_integration.ipynb) — the same DDM with a different SBI library, for cross-toolkit comparison.
• LAN tutorials (main_tutorial.ipynb) — the original LANfactory workflow this integration builds on top of.