Speeding Up Sampling for an RLDDM

[theonlydvr]

I've been able to get an RLDDM working with HSSM; however, I've been having substantial issues with slow sampling (especially during tuning). It generally seems like the slow sampling occurs because of an excessive number of steps at each iteration (frequently hitting the max tree depth). This is especially problematic for the RLDDM because it needs to use the PyMC scan function which leads to relatively slow function evaluations. From my understanding, this kind of behavior normally occurs when the priors are somewhat misspecified. I was wondering if there was any way in HSSM currently to use the priors that were previously implemented in HDDM? The original HDDM priors used non-Gaussian distributions for the DDM parameters and I can't see any way to specify non-Gaussian, hierarchical regressions for parameters in the HSSM framework (I know Bambi can definitely handle them). Is there any way to use the original HDDM priors in HSSM?

Alternatively, does anyone have suggestions for potentially improving the sampling speed or reduce the number of steps per iteration? I can share any code that might be helpful!

[theonlydvr]

Adjusting the likelihood function to compute the Q-values for each session/subject in parallel in the scan function massively sped things up. I'm still curious about adjusting the priors but the model is properly converging with enough tuning steps now.

[AlexanderFengler]

@theonlydvr great to hear.

Few comments on your original statements:

1. You should be able to define priors as you wish, also for hierarchical models. Could you let us know where you reached a limit?
2. If you use the nuts_numpyro sampler, at least you forgo evaluating the pytensor scan function, which (I have seen this) can lead to huge speed-ups. I am assuming you are already doing that.

Another suggestion: Possibly try VI. It can massively help with speed, and posterior approximations do tend to be quite reasonable if parameter trade-offs are not huge.

Are you interested in helping us include RLDDM functionality properly into HSSM?

[theonlydvr]

The issue I was having was figuring out how to define a regression for a DDM term as a non-linear regression. So for example, the boundary separation would be defined by a gamma distribution with a generalized linear model. I couldn't find how to do something like that in the current documentation.

I did use the nuts_numpyro sampler which substantially sped things up. I think my code is currently converging for a 31 subject, 4 session, 200 trial (per session) dataset in under 2 hours when running on the GPU.

I would definitely be interested in including RLDDM functionality properly into HSSM! My current code is relatively straightforward but I'll probably need some guidance about how the team would like the implementation structured for general users because it required me to work-around a couple of the default features in the package.