Compositional Amortized Inference for Hierarchical Bayesian Models

References: - Arruda et al. (2026) — Compositional amortized inference for large-scale hierarchical Bayesian models - Geffner et al. (2023) — Compositional score modeling for simulation-based inference

What This Tutorial Covers

Hierarchical Bayesian models are powerful tools in cognitive science and many other fields, but they face a fundamental scalability challenge: inference must be repeated every time data grows — more subjects, more trials, larger studies. Standard amortized inference (training a neural network to approximate the posterior) requires training data that match the exact study size, which makes generalizing to larger datasets difficult.

Compositional amortized inference resolves this by decomposing the posterior into pieces that can be independently amortized and then composed at test time — without retraining. The key insight is:

\[p(\boldsymbol{\eta} \mid \mathbf{Y}_{1:J}) \propto p(\boldsymbol{\eta})^{1-J} \prod_{j=1}^{J} p(\boldsymbol{\eta} \mid \mathbf{Y}_j)\]

Each posterior term can be trained using a single network, and then combined via a compositional sampling procedure. This allows a network trained on small datasets (e.g., 1 or 5 subjects) to be repurposed for inference over thousands of subjects at test time.

The Example: Hierarchical Evidence Accumulation Model

We use the Evidence Accumulation Model (EAM) from mathematical psychology as our running example (see Habermann et al. (2025) — Amortized Bayesian Multilevel Models). The EAM is a widely-used model of two-alternative forced-choice reaction time tasks. It captures both response times and choice accuracy in a principled mechanistic framework.

The hierarchical structure has two levels:

Level	Parameters	Description
Group	\(\mu_\nu, \mu_{\log\alpha}, \mu_{\log t_0}, \sigma_\nu, \sigma_{\log\alpha}, \sigma_{\log t_0}, \beta\)	Population-level tendencies and variability
Subject	\(\nu_p, \alpha_p, t_{0,p}\)	Individual drift rates, thresholds, and non-decision times

By the end of this miniature tutorial, you will know how to: 1. Define a hierarchical generative model with BayesFlow 2. Train a CompositionalDiffusionModel for group-level parameters 3. Scale inference to different numbers of subjects at test time via compositional_sample

1. Setup and Imports

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import beta as beta_dist, norm as norm_dist

import bayesflow as bf
import keras

from numba import njit

2. Hierarchical Evidence Accumulation Model (EAM)

2.1 The Evidence Accumulation Model

The EAM models a decision as a noisy accumulation of evidence over time. A particle starts at position \(\beta \cdot \alpha\) (the relative starting point scaled by the boundary) and drifts with rate \(\nu\) and diffusion noise \(\sigma \, dW_t\) until it hits either \(0\) (lower boundary, choice = 0) or \(\alpha\) (upper boundary, choice = 1). The response time is the first-passage time plus a non-decision offset \(t_0\) for sensory and motor latency.

The four EAM parameters are:

Parameter	Symbol	Interpretation
Drift rate	\(\nu\)	Speed and direction of evidence accumulation
Boundary separation	\(\alpha\)	Response caution — larger values mean slower but more accurate responses
Non-decision time	\(t_0\)	Sensory encoding + motor execution time, not part of decision
Starting point	\(\beta \in [0,1]\)	Prior bias toward one response option

The simulation below implements the Euler–Maruyama discretization of this stochastic differential equation. We use numba to JIT-compile the inner trial loop to increase simulation speed.

2.2 Hierarchical Prior Structure

Subjects are modeled as exchangeable draws from a population. The generative process is:

Group level (shared across all subjects): \[\mu_\nu \sim \mathcal{N}(0.5, 0.3^2), \quad \mu_{\log\alpha} \sim \mathcal{N}(0, 0.05^2), \quad \mu_{\log t_0} \sim \mathcal{N}(-1, 0.3^2)\] \[\log\sigma_\nu \sim \mathcal{N}(0, 1), \quad \log\sigma_{\log\alpha} \sim \mathcal{N}(0, 1), \quad \log\sigma_{\log t_0} \sim \mathcal{N}(-1, 0.3^2)\] \[\beta_\text{raw} \sim \mathcal{N}(0, 1), \quad \beta = F^{-1}_{\text{Beta}(50,50)}(\Phi(\beta_\text{raw}))\]

Subject level (per-subject draws from the group): \[\nu_p \sim \mathcal{N}(\mu_\nu, \sigma_\nu^2), \quad \log\alpha_p \sim \mathcal{N}(\mu_{\log\alpha}, \sigma_{\log\alpha}^2), \quad \log t_{0,p} \sim \mathcal{N}(\mu_{\log t_0}, \sigma_{\log t_0}^2)\]

The starting point \(\beta\) is shared across subjects (a global bias), while drift rate and thresholds vary. Log-normal priors on \(\alpha\) and \(t_0\) ensure positivity without truncation.

The beta_raw reparameterization maps an unconstrained Gaussian to the \([0,1]\) Beta-distributed prior via the probability integral transform.

@njit
def simulate_ddm_trial(nu, alpha, t0, beta, dt=1e-3, scale=1.0, max_time=10.0):
    """
    Simulates one realization of the diffusion process given
    a set of parameters and a step size `dt`.

    Returns:
    --------
    (x, c) - a tuple of response time (y - float) and a
        binary decision (c - int)
    """

    # Inits (process starts at relative starting point)
    y = beta * alpha
    rt = t0
    const = scale * np.sqrt(dt)

    # Loop through process and check boundary conditions
    while (alpha >= y >= 0) and rt <= max_time:
        # Perform diffusion equation
        z = np.random.randn()
        y += nu * dt + const * z

        # Increment step counter
        rt += dt

    if y >= alpha:
        c = 1.0
    else:
        c = 0.0
    return c, rt


@njit
def _simulate_ddm(nu, alpha, t0, beta, n_trials):
    """Simulates all subjects and trials in Numba-compiled code."""
    n_subjects = nu.shape[0]
    data = np.zeros((n_subjects, n_trials, 2))

    for j_subject in range(n_subjects):
        for i_trial in range(n_trials):

            data[j_subject, i_trial] = simulate_ddm_trial(nu[j_subject], alpha[j_subject], t0[j_subject], beta)

    return data


def simulate_ddm(nu, alpha, t0, beta, n_subjects=1, n_trials=30):
    """Python wrapper around the Numba-compiled DDM simulator."""
    nu = np.broadcast_to(np.asarray(nu, dtype=np.float64), (n_subjects,))
    alpha = np.broadcast_to(np.asarray(alpha, dtype=np.float64), (n_subjects,))
    t0 = np.broadcast_to(np.asarray(t0, dtype=np.float64), (n_subjects,))

    data = _simulate_ddm(nu, alpha, t0, beta, n_trials)
    if n_subjects == 1:
        data = data[0]
    return dict(sim_data=data)

def score_log_norm(x, m, s):
    return -(x-m) / s**2


def beta_from_normal(z, a, b):
    u = norm_dist.cdf(z)
    x = beta_dist.ppf(u, a, b)  # Beta inverse CDF
    return x


global_prior = {
    'mu_nu': (0.5, 0.3),
    'mu_alpha': (0.0, 0.05),
    'mu_t0': (-1.0, 0.3),
    'log_sigma_nu': (-1.0, 1.0),
    'log_sigma_alpha': (-3.0, 1.0),
    'log_sigma_t0': (-1.0, 0.3),
    'beta_raw': (0.0, 1.0),
}

def sample_hierarchical_priors(n_subjects=1):
    """
    Returns a dict with group params and per subject params.
    """
    # Group level
    mu_nu = np.random.normal(global_prior["mu_nu"][0], global_prior["mu_nu"][1])
    mu_alpha = np.random.normal(global_prior["mu_alpha"][0], global_prior["mu_alpha"][1])
    mu_t0 = np.random.normal(global_prior["mu_t0"][0], global_prior["mu_t0"][1])

    log_sigma_nu = np.random.normal(global_prior["log_sigma_nu"][0], global_prior["log_sigma_nu"][1])
    log_sigma_alpha = np.random.normal(global_prior["log_sigma_alpha"][0], global_prior["log_sigma_alpha"][1])
    log_sigma_t0 = np.random.normal(global_prior["log_sigma_t0"][0], global_prior["log_sigma_t0"][1])

    beta_raw = np.random.normal(global_prior["beta_raw"][0], global_prior["beta_raw"][1])
    beta = beta_from_normal(beta_raw, a=50, b=50)

    # Subject level
    nu = np.random.normal(mu_nu, np.exp(log_sigma_nu), size=n_subjects)
    alpha = np.exp(np.random.normal(mu_alpha, np.exp(log_sigma_alpha), size=n_subjects))
    t0 = np.exp(np.random.normal(mu_t0, np.exp(log_sigma_t0), size=n_subjects))

    return {
        # group
        "mu_nu": mu_nu,
        "mu_alpha": mu_alpha,
        "mu_t0": mu_t0,
        "log_sigma_nu": log_sigma_nu,
        "log_sigma_alpha": log_sigma_alpha,
        "log_sigma_t0": log_sigma_t0,
        "beta_raw": beta_raw,
        "beta": beta,
        # subjects
        "nu": nu,
        "alpha": alpha,
        "t0": t0,
    }

2.3 The Prior Score Function

Compositional inference requires the score function of the prior, i.e., \(\nabla_\theta \log p(\theta)\). For a Gaussian prior \(\mathcal{N}(\mu, \sigma^2)\), the score is simply: \[\nabla_x \log p(x) = -\frac{x - \mu}{\sigma^2}\]

The score is time-weighted by \((1 - t)\) following the diffusion bridge formulation in Arruda et al. (2026), where \(t \in [0, 1]\) is the diffusion time. At \(t=0\) (start of diffusion) the prior score has full weight; at \(t=1\) (end) it vanishes. Optionally, the prior score function can have a time argument and you can define your one time dependence.

def prior_global_score(x):
    mu_nu = x["mu_nu"]
    mu_alpha = x["mu_alpha"]
    mu_t0 = x["mu_t0"]
    log_sigma_nu = x["log_sigma_nu"]
    log_sigma_alpha = x["log_sigma_alpha"]
    log_sigma_t0 = x["log_sigma_t0"]
    beta_raw = x["beta_raw"]

    parts = {
        "mu_nu": score_log_norm(mu_nu, m=global_prior["mu_nu"][0], s=global_prior["mu_nu"][1]),
        "mu_alpha": score_log_norm(mu_alpha, m=global_prior["mu_alpha"][0], s=global_prior["mu_alpha"][1]),
        "mu_t0": score_log_norm(mu_t0, m=global_prior["mu_t0"][0], s=global_prior["mu_t0"][1]),
        "log_sigma_nu": score_log_norm(log_sigma_nu, m=global_prior["log_sigma_nu"][0], s=global_prior["log_sigma_nu"][1]),
        "log_sigma_alpha": score_log_norm(log_sigma_alpha, m=global_prior["log_sigma_alpha"][0], s=global_prior["log_sigma_alpha"][1]),
        "log_sigma_t0": score_log_norm(log_sigma_t0, m=global_prior["log_sigma_t0"][0], s=global_prior["log_sigma_t0"][1]),
        "beta_raw": score_log_norm(beta_raw, m=global_prior["beta_raw"][0], s=global_prior["beta_raw"][1]),
    }
    return parts

simulator_hierarchical = bf.make_simulator([sample_hierarchical_priors, simulate_ddm])

_ = simulator_hierarchical.sample(1)

3. Building the BayesFlow Workflows

We need two workflows, one for each level of the hierarchy:

3.1 Global Workflow (Group-level Parameters)

Summary statistics: Each subject’s trial data is reduced to 10 response-time quantiles per response and the proportion of response 1, yielding 21 hand-crafted summary statistics.

Inference network: The DiffusionModel is the key component enabling compositional inference. It learns to approximate the per-subject posterior score \(\nabla_{\boldsymbol{\eta}_t} \log p({\boldsymbol{\eta}_t} \mid \mathbf{Y}_j)\) using a diffusion-based score model. During training this is just a standard Diffusion Model (see our Diffusion Model Tutorial). At inference time, scores from many individuals are summed and integrated via an SDE solver to produce samples from the composed posterior.

param_names_global = list(global_prior.keys())
pretty_param_names_global = [
    r'$\mu_\nu$', r'$\mu_\alpha$', r'$\mu_{t_0}$',
    r'$\log \sigma_\nu$', r'$\log \sigma_\alpha$', r'$\log \sigma_{t_0}$',
    r'$\beta_\text{raw}$'
]

adapter = (
    bf.adapters.Adapter()
    .convert_dtype("float64", "float32")
    .concatenate(param_names_global, into="inference_variables")
    .rename("sim_data", "summary_variables")
)

workflow_global = bf.CompositionalWorkflow(
    adapter=adapter,
    simulator=simulator_hierarchical,
    summary_network=bf.networks.DeepSet(depth=1),
    inference_network=bf.networks.DiffusionModel(
        prediction_type="velocity",
        noise_schedule="cosine",
        subnet_kwargs={"widths": (128,)*3}
    ),
    standardize="all"
)

# Uncomment if training needed
# history = workflow_global.fit_online(
#     epochs=500,
#     num_batches_per_epoch=100,
#     batch_size=32,
# )

test_data_single = simulator_hierarchical.sample(100)

test_posterior = workflow_global.sample(
    num_samples=300,
    conditions={'sim_data': test_data_single['sim_data']},
)

fig = bf.diagnostics.recovery(
    estimates=test_posterior,
    targets=test_data_single,
    variable_names=pretty_param_names_global
)

fig = bf.diagnostics.calibration_ecdf(
    estimates=test_posterior,
    targets=test_data_single,
    variable_names=pretty_param_names_global
)

4. Compositional Inference

This is where the method’s scalability shines. We now simulate 50 independent subjects.

test_data = simulator_hierarchical.sample(30, n_subjects=50, n_trials=30)

How Compositional sampling works

Instead of passing all new subjects at once (which the network was not trained for), compositional_sample partitions subjects into mini-batches of M=mini_batch_size. For each mini-batch, the trained CompositionalDiffusionModel evaluates the approximate posterior score contributions and then the total score is: \[\nabla_{\boldsymbol{\eta}} \log p_t(\boldsymbol{\eta} \mid \mathbf{Y}_{1:J}) = d(t) \cdot \bigl[(1-t)(1-J) \nabla_{\boldsymbol{\eta}} \log p(\theta) + \frac{J}{M}\sum_{m=1}^{M} \hat{s}_m(\boldsymbol{\eta} \mid \mathbf{Y}_m)\bigr]\]

This summed score is then used by a diffusion solver.

Key arguments: - compute_prior_score — injects the analytically known prior score into the solver - mini_batch_size — number of subjects processed per score evaluation (trades memory for accuracy)

Compositional sampling accumulates errors. The compositional_bridge_d1 and compositional_bridge_d0 arguments can help with calibration of the posterior by scaling the score over time. This is the exponential function from Arruda et al. (2026) that smoothly transitions from compositional_bridge_d1 at time 1 to compositional_bridge_d0 at time 0. Setting both to the same value (e.g., 0.1) applies a constant scaling factor.

global_posterior = workflow_global.compositional_sample(
    num_samples=300,
    conditions={'sim_data': test_data['sim_data']},
    method="two_step_adaptive",
    steps='adaptive',
    compute_prior_score=prior_global_score,
    batch_size=10,
    mini_batch_size=4,
)

fig = bf.diagnostics.recovery(
    estimates=global_posterior,
    targets=test_data,
    variable_names=pretty_param_names_global
)

Inspecting single marginal posteriors

idx = 1
f, axarr = plt.subplots(2, 4, figsize=(12, 6))

for i, (k, ax) in enumerate(zip(global_posterior.keys(), axarr.flat)):

    sns.histplot(global_posterior[k][idx].squeeze(), ax=ax, color="#263078", alpha=0.4)
    sns.despine(ax=ax)

    ax.axvline(test_data[k][idx].squeeze(), color="black")


f.tight_layout()