Basic Conceptual overview of the HSSM Ecosystem.

HSSM provides a platform for theoreticians and computational scientists to contribute (and stress test) new models and for experimentalists to test their empirical data against a larger, expanding array of such models. The resulting cycle contributes to accelerate translation between theoretical and experimental contributions to the field of computational cognitive neuroscience.

Overview of the HSSM ecosystem.

The core ecosystem consists of three interoperable Python packages (left column) that span the end-to-end workflow from cognitive model simulation to scientific analysis via Bayesian parameter inference (right panel). ssm-simulators provides fast forward simulators for a broad, extensible set of cognitive process models, and produces the training data required for likelihood surrogate approximation. LANFactory trains lightweight neural networks — likelihood approximation networks (LANs), choice probability networks (CPNs), and omission probability networks (OPNs) — using JAX or PyTorch, and exports them as ONNX files that can be uploaded to a shared HuggingFace repository for community reuse. HSSM is the user-facing hub: it retrieves the relevant simulator and likelihood, combines them with user-supplied behavioral data and a formulaic model description, constructs the model via PyMC and Bambi, and supports full Bayesian inference (MCMC, VI, MAP) alongside posterior predictive plotting, diagnostics, and model comparison. The right panel traces the full pipeline from forward simulation through likelihood amortization through to the inverse problem of parameter inference and downstream scientific analysis.

Architecture of the HSSM package.

HSSM is designed around two complementary APIs at its core: the high-level API (HSSM class, RLSSM class, etc.) constructs hierarchical Bayesian models automatically from user-supplied data, model choice, formulae, and priors; the low-level API exposes pre-assembled PyMC random variables and supports embedding HSSM likelihoods in entirely custom PyMC models. The likelihood layer accommodates multiple sources, including surrogates produced via LANFactory, BayesFlow, or sbi and deployed through JAX, ONNX, or PyTensor. Forward simulators are drawn from ssm-simulators for prior and posterior predictive sampling; users can supply their own. The assembled model is a standard PyMC distribution and can be fit with a wide range of inference backends — gradient-based MCMC via NUTS, NumPyro, or BlackJax; variational inference; MAP estimation; or external samplers. Outputs are compatible with ArviZ diagnostics and HSSM’s plotting utilities, and models can be persisted via native save/load methods. See Figure 3—figure supplement 1 for the full internal architecture. Figure 3—figure supplement 1. HSSM Conceptual Architecture: Detailed

Simplified task structure for the “cavanagh_theta” dataset.

Training pairs (top) and example test-phase pairings with corresponding stim values (bottom).

Loading a dataset that is included in the HSSM package.

Getting the list of models currently supported in HSSM

Illustrating a specified HSSM model.

We provide a custom string for a simple readout of prior choices and the cognitive model parameters. Inheriting from the PyMC ecosystem, we can also get a simple graphical model depiction of our model via the graph method.

HSSM natively returns ArviZ idata objects which hold our traces from a model fit.

We can then access the full plotting capability of ArviZ to illustrate results. We showcase the basic trace plot above.

Schematic of the angle DDM.

Reinforcement Learning - Sequential Sampling Models (RLSSM).

RLSSM is a canonical example of combining across-trial cognitive processes such as reinforcement learning (RL) alongside within-trial decision processes such as DDMs and other SSMs. The DDM trial-by-trial parameters are modified via RL, e.g., to capture how the value of decision variables may evolve with experience or feedback in the task. The model produces a joint distribution of choice and RT distribution. More generally, the across-trial RL process could be swapped with other dynamic cognitive processes such as Bayesian updating, latent state models of attentional processes or modulatory models of cognitive control. Similarly, the within-trial DDM process can be replaced with any suitable decision models (e.g., Angle DDM, race models such as LBA to account for n>2 choice behaviors, etc). Likelihoods for the entire joint RL-SSM process can be differentiable and used with either analytic or likelihood-free methods (e.g., LANs).

Specifying a hierarchical model with regression backends on the “a” and “v” parameters of the “angle” DDM with HSSM.

HSSM allows multiple approaches to specifying such regressions. Here we show the recommended procedure, via the include argument and a list of hssm.Param objects.

Posterior predictive checks for the RLSSM model.

Here the Q learning rule is used and the decision process is Angle (linear collapsing bound DDM), but either process can be replaced with other learning rules or SSMs. Top: early-trial choice accuracy, shown as the probability of choosing the high-reward option across binned condition trials. Empirical trajectories are shown with solid lines and markers, while posterior predictive simulations are shown with dashed lines; thin translucent lines show individual simulated datasets. Bottom: condition-wise signed response-time distributions, where upper-boundary/best-option choices are positive and lower-boundary/worst-reward choices are negative. Across AB, CD, and EF reward-contrast conditions, the model captures both the learning-related separation in choice accuracy and the broad structure of the response-time distributions.

Basic RLSSM model configuration in HSSM.

The RLSSMConfig object ties together the choice of decision process (here, the angle DDM with an approximate differentiable likelihood), the learning process (here, delta rule with an analytical update), and the data columns required by the learning rule.

Fitting an RLSSM model with the numpyro backend, and inspecting the posterior summary.

The hierarchical regression interface, NUTS samplers, and ArviZ-compatible outputs are all inherited unchanged from the core HSSM class.

Basic posterior predictive plots for 2-choice models.

The plot_predictive method allows for large stylistic and computational flexibility. We highlight the predictive_group argument, which specifies whether we would like plot the prior or the posterior predictive.

Quantile probability plots for visualizing model fits to data, comparing a true and misspecified model

This figure illustrates how QP plots can identify features of the data that the model may predict or miss. Here data is generated synthetically from a model with linearly collapsing boundaries, and then fit with that same model (left) or with the simple DDM (right). The ellipses, reflect the uncertainty (HDI) of where the data are expected to lie for both choice probability in each condition (x-axis) and corresponding RT quantiles (y-axis). For the true model, the data are well within the expected uncertainty of the model for both choice and RT quantiles in each condition. For the misspecified model (right), the model misses some of the RT quantiles, in particlar error are expected to be faster than observed empirically. While synthetic, to maintain harmony with the rest of our examples, the data here is otherwise structured as the “cavanagh_theta” dataset.

HSSM ships with quantile probability plots (Ratcliff and Tuerlinckx, 2002), via the plot_quantile_probability method.

The plot has various computational and visual knobs that can be adjusted to specialize it to a particular scenario. We highlight the get quantile_by argument, which establishes at which level quantiles are being computed. If e.g. (as illustrated) the argument is set to “participant_id”, we compute quantiles first per participant, then aggregate. The ellipses reflect the uncertainty of the quantiles and choice probabilities in line with the chosen aggregation process.

Anatomy of the model cartoon plot.

The plot is built in layers. (a) The structural elements of the cognitive process model are overlaid on the response time axis: here, the decision boundaries, non-decision time (dashed line), drift trajectory (with boundary crossing marked by a red star), and the choice-conditioned response time histogram evaluated at the posterior mean. (b) Posterior uncertainty is added by overlaying many samples from the posterior predictive distribution (grey shading). (c) The observed data are added in blue, allowing direct visual comparison with the posterior predictive. Each element can be toggled independently via keyword arguments.

Model cartoon plot for a two-choice SSM.

The plot_model_cartoon() function visualizes the fitted cognitive process model alongside the observed data. Blue histograms show observed response time distributions (plotted on the negative axis when the choice was −1 and on the positive axis otherwise). Overlaid in grey are individual sampled decision trajectories from the fitted model, with boundary crossings marked in red. The black curve traces the mean posterior predictive distribution. Dashed vertical lines indicate the non-decision time. Users can toggle trajectories, posterior uncertainty, axis layout and various other stylistic elements via the keyword arguments.

Low level API.

This listing shows how we can use random variables that are shipped with HSSM directly inside custom PyMC models. Advanced users can use this lower level interface to build much more ambitious models than feasible via the current base HSSM class.

External samplers.

This listing showcase how we can compile the log-likelihood function of a given HSSM model (we are accessing the underlying PyMC model directly) and downstream use it to serve as the basic likelihood to be called through entirely different MCMC sampler libraries. The zeus library serves as a convenient example to showcase how we can seamlessly exit the entire PyMC ecosystem while making HSSM do useful groundwork for us nevertheless.

Architecture of the ssm-simulators package.

The package provides two core functionalities: forward simulation of cognitive process models via the Simulator class, and a modular training data generation pipeline that combines configurable likelihood estimators and sampling strategies to produce structured datasets for downstream likelihood approximation (e.g. via LANFactory, connectors to the sbi (Tejero-Cantero et al., 2020) and the BayesFlow (Radev et al., 2020) libraries exist). Figure 8—figure supplement 1. ssm-simulators Conceptual Architecture: Detailed

Simulation of the angle DDM using ssm-simulators

Basic training data generation for the “angle” DDM with ssm-simulators

ssm-simulators CLI interface for training data generation from YAML configuration files.

Architecture of the LANFactory package.

The package trains surrogate likelihood networks (LANs, CPNs, OPNs) from training data produced by ssm-simulators, using either JAX or PyTorch backends. Trained networks are converted to the ONNX format and uploaded to HuggingFace for deployment within HSSM or elsewhere. Figure 9—figure supplement 1. LANFactory conceptual architecture, detailed.

LANFactory CLI interface.

Three core commands (i) torchtrain (or the JAX equivalent) to train the network (ii) transform-onnx to generate the ONNX files and (iii) upload-huggingface to upload the network to our HuggingFace database.

LANFactory Model Card.

A model card is associated with every uploaded model in the HuggingFace database.

Feature comparison of likelihood-based software packages for cognitive process modeling.

· = fully supported, ◦ = partial or limited support, – = not supported. “Hierarchical Models” refers to full hierarchical estimation with group and subject parameters. “Trial-by-trial covariates” refers to incorporation of continuous trial-level measures (e.g., neural activity, learning signals) as predictors of model parameters. “Full LMER” refers to lmer-style formula-based specification of between- and within-subject effects on model parameters, including random slopes and intercepts. “RL-SSM integration” refers to built-in support for combining reinforcement learning with sequential sampling models. “Surrogate models” refers to support for models lacking closed-form likelihoods via learned likelihood approximations (e.g., neural networks). “User Extensible” refers to the ability for users to add new models or likelihood functions without modifying the package source code. BayesFlow, which targets amortized posterior estimation and surrogate likelihood training via a fundamentally different paradigm, is discussed in the text but omitted from this comparison.

A more complete view on the architecture of the HSSM package.

HSSM serves as the user-facing hub of the ecosystem and is designed around two complementary APIs (blue block). The high-level API centers on the HSSM class and consumes model and parameter configurations — including regression setup and optionally prior choices — to construct a full hierarchical Bayesian model automatically. The low-level API exposes the underlying pre-assembled PyMC distributions, allowing advanced users to build entirely custom PyMC models while retaining access to HSSM’s cognitive process model likelihoods. The likelihood layer accommodates multiple sources — analytical closed-form expressions, neural network surrogates loaded from ONNX (via either JAX or PyTensor), and blackbox user-provided functions — alongside forward simulators drawn from ssm-simulators for prior and posterior predictive sampling (users can provide their own if desired). The assembled likelihood node is a standard PyMC distribution and can be fit with a wide range of inference backends (yellow block): gradient-based MCMC via PyMC NUTS, NumPyro, or BlackJax; variational inference (ADVI, FullRank ADVI); MAP estimation; or external samplers. Outputs are compatible with ArviZ diagnostics and HSSM’s proprietary plotting utilities, and models can be persisted via native save/load methods.

Architecture of the ssm-simulators package.

The package provides two core functionalities: forward simulation of cognitive process models via the Simulator class, and a modular training data generation pipeline that combines configurable likelihood estimators and sampling strategies to produce structured datasets for downstream likelihood approximation. The package has a modular organization across level of functionality. Simulators can be constructed natively via ssm-simulators or e.g. via the PyDDM backend (Shinn et al., 2020) (with a vision toward extension). Training data generators are designed to combine different types of likelihood estimators as training signals for downstream networks (currently implemented are likelihoods based on PyDDM (Shinn et al., 2020) and our native approach following (Fengler et al., 2021).

Detailed architecture of the LANFactory package.

The package is configured through a central YAML file, optionally invoked via the jaxtrain and torchtrain CLI commands. The configuration drives three parallel substructures: network configuration (architecture choice and target network type — LAN, OPN, or CPN), training configuration (batch size, optimizer, and other hyperparameters), and data loading (parallelized, shuffled feature/label data, typically produced by ssm-simulators). The training subsystem provides parallel JAX and PyTorch implementations via the JaxMLP and TorchMLP class abstractions, both feeding into a unified training loop. Trained weights are exported to ONNX and uploaded to HuggingFace for downstream consumption by HSSM or other tools. Throughout training, an optional MLflow integration logs metrics, stores artifacts, and tracks data lineage to support reproducible network development.