Questions about fitting Stroop-task models in HSSM

[gvnpolzin]

Hi,

I’m currently using HSSM to fit hierarchical sequential sampling models to a Stroop task, and I was hoping to get some guidance on several modeling and sampling issues I’ve encountered.

I have been comparing the ddm, ddm_sdv, full_ddm, gamma_drift, and gamma_drift_angle model families. My main goal is to model condition-related differences in response-time and accuracy distributions while allowing parameters such as drift rate, boundary separation, and nondecision time to vary hierarchically across participants and sessions.

I have a few questions:

1. Divergences and number of cores
   I have noticed that changing the number of cores/chains appears to substantially change the number of divergences I obtain, even when the model specification is otherwise unchanged. I am using a fixed random seed and, for differentiable models, NumPyro/NUTS. For example, if I use 4 cores, 4 chains with 200 draws, there are significant divergences compared to just 1 core, 1 chain. Is this behavior expected? Should changing the number of cores materially affect the divergence behavior of a model? Are there recommended settings for obtaining reproducible diagnostics across runs, particularly for hierarchical HSSM models?

2. Rapid model testing
   I am testing a fairly large number of hierarchical model specifications before committing to long sampling runs.
   What workflow would you recommend for rapidly checking whether a model is correctly specified and computationally viable before running a full model? Is there a recommended combination of short tuning/draw runs, prior predictive checks, or another diagnostic that is appropriate for a quick test of an HSSM model?

3. gamma_drift and gamma_drift_angle
   I am also experimenting with the newer gamma_drift and gamma_drift_angle models. I would particularly appreciate clarification on the interpretation of the c, shape, and scale parameters and how they should be used when modeling a conflict task such as Stroop. At the moment, I am seeing posterior-predictive RT distributions in which the congruent and incongruent simulated distributions can look nearly identical, even when I specify a condition effects on v. I expected the incongruent distribution to be shifted relative to the congruent distribution. I am trying to determine whether this reflects an issue with my model specification, or an incorrect assumption about how the gamma-drift model should behave. I would also be interested in any recommended examples that demonstrate how gamma_drift or gamma_drift_angle should be specified and interpreted.

I’m happy to provide model formulas, package versions, and sampler output if that would be useful.

Thank you very much for your time and for developing HSSM.

Best,
Gabi

[EItanm1999]

Hi Gabi! I can take the divergences question and full_ddm.

First, LANs are faster than numerical integration and can use NUTS sampling, so that's already a big time save. Right now full_ddm uses numerical integration, which is slow and has no gradients, so it can't use NUTS. A LAN is a neural network trained to approximate the likelihood. It's much faster and differentiable, so NUTS works.

The number of cores and chains shouldn't change whether the model is well-behaved. A fixed seed also won't make a 1-chain run and a 4-chain run comparable. NumPyro splits the seed into a separate key per chain, so the single chain in your 1-chain run isn't one of the four chains in your 4-chain run. Four chains also give you four chances to find a poorly behaved region of the posterior, and with only 200 draws each, one bad chain's divergences will overpower other results. A single clean chain is less trustworthy, since it basically just had fewer chances to find a problem area.

For reproducible diagnostics, fix the seed and the number of chains, then compare like with like. For a hierarchical LAN fit I typically use 4 chains, tune=500, draws=500, sampler="numpyro". If you still see divergences there, the first thing I'd try is link_settings="log_logit" in the HSSM(...) constructor. By default, subjects can wander outside the parameter bounds. log_logit transforms the parameter domain to the real number line, so that can't happen. That caused the most divergences/headaches in my hierarchical fits (though it may not be your problem). It would also help to look at divergences per chain (idata.sample_stats.diverging.sum(("draw",))). If one chain produces most of them, then the chain is stuck but the model isn't necessarily wrong everywhere. Feel free to post your formulas and traces, and we can look at whether the chains agree.

A LAN for full_ddm is in progress. It will be much faster than the numerical-integration likelihood. I don't have a firm date, but my guess is a month or two (no promises). A LAN for non-decision-time variability should land sooner, as it's going through processing, but basically done. Whichever you use, I'd recommend fitting the variability parameters (sv, sz, st) at the group level rather than hierarchically. They're weakly identified per subject, and letting them vary by participant tends to cost more in sampling trouble than it buys. If only one of them misbehaves (sz often does), keep just that one at the group level.

For getting through possibilities quickly, use a process that eliminates large batches at each step:

1. For each candidate model, simulate data across a range of parameter values and ask whether it can produce the patterns you care about. On one project, this ruled out a whole model class by reproducing one feature of the data but totally missing on error RTs, and no fit would have fixed that.
2. Use Variational Inference with a few starting points. If different starts land in different places, that's an early warning that something isn't identified, before you've spent any sampling time. (VI needs a differentiable likelihood, so this works for the LAN and analytical models but not the current full_ddm)
3. Once you've ruled out the most obvious bad fits, run a small NUTS test, e.g. 2 chains, tune=100, draws=100. You're not looking at estimates, just sampler health. If the step size collapses toward zero or tree depth is pinned at the maximum (10 by default) on every draw, the run is broken. That usually shows up within the first few dozen iterations, so you can kill it rather than wait hours to find out. You can watch this live in the NumPyro progress bar (step size and number of steps per draw), or check afterwards with idata.sample_stats.step_size and idata.sample_stats.tree_depth.
4. After you know it has a chance to work and fit the data, simulate data from known parameters, fit it, and check that you get them back. This tells you whether a parameter like c or scale is identified with your number of trials.

PPCs are a great way to see if you're on the right track, but they and the traces tell you different things. Traces (and r_hat, divergences, tree depth) tell you whether the sampler explored the posterior. PPCs tell you whether that posterior describes your data. You need both. I've had fits with clean traces where the PPC showed a misfit. I've also seen chains that all agreed with each other because they were stuck at the same parameter bound. When chains converge to the same place, check that the place makes sense.

[gvnpolzin]

Thank you so much for the detailed explanation! This is very helpful, especially the distinction between sampler diagnostics and PPCs. I'll try using 4 chains with consistent sampling settings and link_settings="log_logit" to see whether that improves convergence. I'll also look more closely at divergences per chain, tree depth, and step size rather than relying on whether an individual run completes successfully.