import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from scipy.stats import norm
RNG = np.random.default_rng(0)
torch.manual_seed(0)
model_label = ["Direct difference", "Hyperbolic", "Hyperboloid"]
colors = ["#4477AA", "#EE6677", "#228833"] Model comparison with a neural network — a worked SBI example
3 competing cognitive models of intertemporal choice. Given one participant’s 27 binary answers on the MCQ, we want to know which model generated the data
simulate → label → train a classifier → read the posterior off the softmax
0. Setup
1. Kirby MCQ
27 fixed questions, a smaller-sooner reward against a larger-later one, e.g. “$54 today, or $55 in 117 days?” Every participant produces a length-27 vector of 0/1 (0 = chose sooner, 1 = chose later)
# Kirby MCQ items: [smaller-sooner amount, larger-later amount, larger-later delay]
# The smaller-sooner option is always immediate (delay 0)
MCQ_STIM = np.array([
[54, 55, 117], [55, 75, 61], [19, 25, 53], [31, 85, 7], [14, 25, 19],
[47, 50, 160], [15, 35, 13], [25, 60, 14], [78, 80, 162], [40, 55, 62],
[11, 30, 7], [67, 75, 119], [34, 35, 186], [27, 50, 21], [69, 85, 91],
[49, 60, 89], [80, 85, 157], [24, 35, 29], [33, 80, 14], [28, 30, 179],
[34, 50, 30], [25, 30, 80], [41, 75, 20], [54, 60, 111], [54, 80, 30],
[22, 25, 136], [20, 55, 7],
], dtype=float)
def unpack_stim(stim):
SS, LL = stim[:, 0], stim[:, 1]
if stim.shape[1] == 3:
tSS, tLL = np.zeros_like(SS), stim[:, 2]
else:
tSS, tLL = stim[:, 2], stim[:, 3]
return SS, LL, tSS, tLL
def sigmoid(x):
return 0.5 * (1 + np.tanh(np.clip(x, -500, 500) / 2))
print(f"{MCQ_STIM.shape[0]} items")2. Three theories of how delay discounts value
Each model a simulator (parameters in, binary responses 0/1 choices out). Each model is vectorized over simulations (the first axis of params is one draw per simulated subject)
| Model | Value / rule | Parameters |
|---|---|---|
| Direct difference | compares amount vs. delay, weighted by w |
u, v, w |
| Hyperbolic | \(V = A / (1 + k\,t)\) | k, m |
| Hyperboloid | \(V = A / (1 + k\,t^{s})\) — hyperbolic is the \(s=1\) case | k, s, m |
def r_direct_difference(params, stim):
"""Direct-difference model (Dai & Busemeyer). params columns: u, v, w"""
u, v, w = params[:, [0]], params[:, [1]], params[:, [2]]
SS, LL, tSS, tLL = unpack_stim(stim)
A = LL ** u - SS ** u # advantage of LL on the amount dimension
# tSS is 0 for immediate items; define 0**v := 0 (immediate = no delay cost).
pos = tSS > 0
tSS_v = np.where(pos, np.power(np.where(pos, tSS, 1.0), v), 0.0)
B = tLL ** v - tSS_v # disadvantage of LL on the delay dimension
d = w * A - (1 - w) * B
s = np.maximum(np.sqrt(w * (1 - w)) * np.abs(A + B), 1e-12)
p = norm.cdf(d / s)
return (RNG.random(p.shape) < p).astype(np.float32)
def r_hyperbolic(params, stim):
""" hyperbolic discounting + logistic choice. params: k, m"""
k, m = params[:, [0]], params[:, [1]]
SS, LL, tSS, tLL = unpack_stim(stim)
d = LL / (1 + k * tLL) - SS / (1 + k * tSS) # value difference
p = sigmoid(d / m) # choice probability
return (RNG.random(p.shape) < p).astype(np.float32)
def r_hyperboloid(params, stim):
""" hyperboloid params: k, s, m"""
k, s, m = params[:, [0]], params[:, [1]], params[:, [2]]
SS, LL, tSS, tLL = unpack_stim(stim)
d = LL / (1 + k * tLL ** s) - SS / (1 + k * tSS ** s)
p = sigmoid(d / m)
return (RNG.random(p.shape) < p).astype(np.float32)3. Priors
We draw each model’s parameters from priors, not fixed values
def draw_priors(n):
"""Returns one (n, n_params) array per model, in the order of model labels"""
# Direct difference: u, v, w
w = sigmoid(RNG.normal(0, 1.5, n))
v = RNG.normal(0.7, 0.3, n)
u = RNG.normal(0.7, 0.3, n)
dd = np.stack([u, v, w], axis=1)
# Hyperbolic: k, m
k_c = np.exp(RNG.normal(-5, 3, n))
m_c = np.exp(RNG.normal(1, 5, n))
hyp = np.stack([k_c, m_c], axis=1)
# Hyperboloid: k, s, m
k_b = np.exp(RNG.exponential(0.1, n))
s_b = np.exp(RNG.normal(0.9, 0.2, n))
m_b = np.exp(RNG.exponential(20.0, n))
hbd = np.stack([k_b, s_b, m_b], axis=1)
return dd, hyp, hbd 4. Build the labelled dataset and split it
Simulate n subjects from each model, stack them, and tag every row with the index of the model that produced it.
- train : the data the network trains on
- validation : evaluated during training to check if we’re generalizing (tuning occurs by monitoring validation loss)
- test : drawn and set aside, it is only touched during final held out evaluation (confusion matrix, calibration)
All three came from same priors so we are not checking for robustness to a mis-specified prior
def simulate_dataset(n_per_model, stim=MCQ_STIM):
dd, hyp, hbd = draw_priors(n_per_model)
X = np.concatenate([ # 27 choices per subject
r_direct_difference(dd, stim),
r_hyperbolic( hyp, stim),
r_hyperboloid( hbd, stim),
], axis=0).astype(np.float32)
y = np.repeat([0, 1, 2], n_per_model).astype(np.int64) # generating-model label
return X, y
X_np, y_np = simulate_dataset(20_000) # train
Xval_np, yval_np = simulate_dataset(12_000) # val
Xtest_np, ytest_np = simulate_dataset(12_000) # test
print("train:", X_np.shape, " val:", Xval_np.shape, " test:", Xtest_np.shape)
print("LL choice rate per model:", [round(float(X_np[y_np == k].mean()), 3) for k in range(3)])5. Build neural network
A small feed-forward net (100 → 50 → 25, tanh). The final layer outputs raw logits, nn.CrossEntropyLoss applies log-softmax internally.
We softmax during prediction time, in posterior()
class ModelComparisonNet(nn.Module):
def __init__(self, n_inputs, n_outputs=3):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_inputs, 100), nn.Tanh(),
nn.Linear(100, 50), nn.Tanh(),
nn.Linear(50, 25), nn.Tanh(),
nn.Linear(25, n_outputs),
)
def forward(self, x):
return self.net(x)
@torch.no_grad()
def posterior(self, x):
return torch.softmax(self(x), dim=-1)6. Train with cross-entropy
def train(model, X, y, X_val, y_val, epochs=60, batch_size=1000,
lr=1e-3, weight_decay=1e-5, verbose=True):
opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
loss_fn = nn.CrossEntropyLoss()
n = len(X)
for epoch in range(epochs):
model.train()
perm = torch.randperm(n)
for i in range(0, n, batch_size):
idx = perm[i:i + batch_size]
opt.zero_grad()
loss = loss_fn(model(X[idx]), y[idx]) # logits vs. labels
loss.backward()
opt.step()
if verbose and (epoch % 10 == 0 or epoch == epochs - 1):
model.eval()
with torch.no_grad():
acc = (model(X_val).argmax(1) == y_val).float().mean().item()
print(f"epoch {epoch:3d} | val acc {acc:.3f}")
return model
X = torch.from_numpy(X_np); y = torch.from_numpy(y_np)
Xv = torch.from_numpy(Xval_np); yv = torch.from_numpy(yval_np)
Xt = torch.from_numpy(Xtest_np); yt = torch.from_numpy(ytest_np)
model = ModelComparisonNet(X.shape[1])
train(model, X, y, Xv, yv, epochs=60)
model.eval()
probs = model.posterior(Xt).numpy() # posterior model probs on held-out test
preds = probs.argmax(1)
test_acc = (preds == ytest_np).mean()
print(f"\nheld-out test accuracy: {test_acc:.3f}")7. Does it work? Model recovery (on the held-out test set)
def confusion(y_true, y_pred, normalize=True):
C = np.zeros((3, 3))
for t, p in zip(y_true, y_pred):
C[t, p] += 1
return C / C.sum(1, keepdims=True) if normalize else C
C = confusion(ytest_np, preds) # held-out test set
fig, ax = plt.subplots(figsize=(5.6, 4.8))
im = ax.imshow(C, cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(3), model_label, rotation=20, ha="right")
ax.set_yticks(range(3), model_label)
ax.set_xlabel("Predicted model"); ax.set_ylabel("Generating model")
ax.set_title("Model recovery (test set)")
for i in range(3):
for j in range(3):
ax.text(j, i, f"{C[i, j]:.2f}", ha="center", va="center",
color="white" if C[i, j] > 0.55 else "black", fontsize=13, fontweight="bold")
fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="P(predicted | generating)")
ax.tick_params(length=0); fig.tight_layout(); plt.show()8. Check calibration
Bin predictions by their confidence (the max softmax value) and check the accuracy in each bin.
def calibration(probs, y_true, n_bins=10):
conf = probs.max(1)
correct = (probs.argmax(1) == y_true)
edges = np.linspace(1/3, 1.0, n_bins + 1)
out = []
for lo, hi in zip(edges[:-1], edges[1:]):
m = (conf >= lo) & (conf < hi)
if m.sum():
out.append((conf[m].mean(), correct[m].mean(), int(m.sum())))
return np.array(out)
cal = calibration(probs, ytest_np) # held-out test set
fig, ax = plt.subplots(figsize=(5.2, 5.0))
ax.plot([1/3, 1], [1/3, 1], "--", color="gray", lw=1.5, label="perfect calibration")
ax.scatter(cal[:, 0], cal[:, 1], s=30 + 260 * cal[:, 2] / cal[:, 2].max(),
color="#4477AA", alpha=0.9, edgecolor="white", zorder=3)
ax.set_xlabel("Predicted confidence (max softmax)")
ax.set_ylabel("Empirical accuracy")
ax.set_xlim(0.3, 1.02); ax.set_ylim(0.3, 1.02); ax.set_aspect("equal")
ax.legend(loc="upper left", frameon=False); ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout(); plt.show()9. Apply the trained net to a fresh dataset
You can dowload data from the many available datasets on Kirby MCQ and try it out. Just point csv path to the downloaded file and make sure rename columns and fields where appropriate or reformat (long vs wide)
The network will give each participant a full posterior over the three models e.g. [0.55, 0.05, 0.40] Collapsing that to a single label (argmax) throws the uncertainty away and biases any population summary. So here:
- keep the whole posterior vector per participant (the stacked bars)
- estimate the population composition by averaging those vectors, not counting argmax winners
If no real data csv is found, the code falls back to a balanced simulated demo batch
import os
CSV_PATH = "SimulatedMCQData.csv" # point to your real data
n_each = 200 # for the demo fall back, 200 per model, 1/3 truth
if os.path.exists(CSV_PATH):
ids, Xreal = load_mcq_csv(CSV_PATH)
y_demo = None
print(f"loaded {len(ids)} participants from {CSV_PATH}")
else:
print(f"{CSV_PATH} not found, simulating a demo batch ")
Xreal, y_demo = simulate_dataset(n_each)
ids = np.arange(len(Xreal))
# One forward pass -> a full posterior over models for every participant
real_probs = model.posterior(torch.from_numpy(Xreal)).numpy() # shape (N, 3)# Per participant posterior vector
order = np.argsort(real_probs.argmax(1) + real_probs.max(1))
P = real_probs[order]; m = len(P)
fig, ax = plt.subplots(figsize=(11, 3.6))
bottom = np.zeros(m)
for k in range(3):
ax.bar(range(m), P[:, k], bottom=bottom, width=1.0, color=colors[k], label=model_label[k])
bottom += P[:, k]
ax.set_xlim(-0.5, m - 0.5); ax.set_ylim(0, 1)
ax.set_xlabel(f"Participant (n = {m}, sorted for readability)")
ax.set_ylabel("P(model | choices)")
ax.legend(ncol=3, loc="lower center", bbox_to_anchor=(0.5, 1.01), frameon=False)
ax.spines[["top", "right"]].set_visible(False); fig.tight_layout(); plt.show()# Population composition: average the posteriors, don't count argmax
argmax_prop = np.bincount(real_probs.argmax(1), minlength=3) / len(real_probs)
mean_posterior = real_probs.mean(axis=0)
print(f"{'':24s}{'DD':>8}{'Hyp':>8}{'Hbd':>8}")
if y_demo is not None: # only known for the demo batch
truth = np.bincount(y_demo, minlength=3) / len(y_demo)
print(f"{'true composition':24s}" + "".join(f"{v:8.3f}" for v in truth))
print(f"{'argmax counts (biased)':24s}" + "".join(f"{v:8.3f}" for v in argmax_prop))
print(f"{'averaged posteriors':24s}" + "".join(f"{v:8.3f}" for v in mean_posterior))Recap
- simulate → label → classify
- softmax = posterior Training with cross-entropy makes the outputs calibrated model probabilities, verified on the diagonal above
- held-out evaluation Fit on training data, tune during validation, and test on unseen data that way recovery and calibration are entirely out of sample
- errors are informative The hyperbolic/hyperboloid confusion is the net reporting genuine model similarity (nesting)
Note: the 27 items are fixed, so a plain MLP works. For variable-length choice sequences, swap the MLP for a permutation-invariant set encoder (DeepSet) or an LSTM