MNLE with experimental conditions

[Jiashun97]

Hi there,

I am trying to apply MNLE in DDMs that need to allow some parameters free to vary across experimental conditions. I noticed that in the sbi/tutorials/Example_01_DecisionMakingModel.ipynb, it provides a simulator wrapper to do so. If I understand it correctly, it makes all parameters free across conditions. However, in some of the behavioral experiments, only one or two parameters are free and the others are fixed. I was wondering if it is possible to implement this in MNLE. Thanks for your help!

Best regards,
Jiashun

[Jiashun97]

I got the answer from Jan: "the emulator is trained on both the parameter and the conditioning variables. At inference time you build a wrapper that fixes the conditions via the potential function wrapper. You could use the same mechanism to fix some of the actual parameters as well at inference time."

[Jiashun97]

Thanks for your answer. I know how to fix a parameter to a specific value in the potential function wrapper. However, the key idea in my analysis is to force the parameter to be the same in two conditions. Suppose we have a five-parameter competing accumulator (CA) model: dI, I, k, a, ter are the parameters. The experiment has a difficulty manipulation that only affect dI (input difference), resulting in two exp. conditions. The behavioral data of one subject will have three columns: RTs, choices, and conditions. In the analysis, I am trying to build a CA model that has six parameters: dI1 (for condition 1), dI2 (for condition 2), I, k, a, ter. This means the dI is free to vary across the condition while the others are fixed (because the experimental manipulation was supposed to only change dI).

In the inference stage of MNLE, it needs to take the RTs, choices, and conditions into account, returning the MAP estimate (or posterior samples) for dI1, dI2, I, k, a, ter. My intuition is that we should also specify the relationship of the parameters before training the network. Maybe we need to change the underlying network structure to achieve this?

I've attached a Jupyter notebook with code for one condition and the implementation of the simulator wrapper. I would greatly appreciate any help with this issue, as fitting decision-making models to experimental data like this is critically important.

SBI with multiple experimental conditions.ipynb.zip

Best regards,
Jiashun

[janfb]

Hi @Jiashun97

thanks for sharing your question here - and sorry for the delayed response.

I am not sure I fully understand your setup. The initial CA model has 5 parameters: $dl$, $l$, $k$, $a$ and $ter$. However, $dl$ is actually an experimental condition that varies across the experiments. In this setting, you would train MNLE on a simulator that has all five parameters as arguments and returns RTs and choices. You would sample parameters from a proposal distribution to generate training data for MNLE.

Q1: Is $dl$ a continuous variable, that just happens to take only two values in the conditions?

At inference time, you would then define a prior only over $l$, $k$, $a$ and $ter$ as those are the parameters you want to do inference over. You would use the Potential Function wrapper to fix $dl$: e.g., for each subject you have N trials (rows in the tensor / dataframe) and RTs, choices, and condition as columns. Then you would pass the N trials to the potential function wrapper to condition an both data and conditions, and run MCMC or VI to obtain posterior samples.

Q2: Am I missing something in the setup?

Comment: I had a brief look at your notebook and it seems in cell 8 or nine where you define the wrapper to simulate both parameters and conditions the indexing into the theta tensor seems to overlap:

# define a simulator wrapper in which the experimental condition are contained
# in theta and passed to the simulator.
def sim_wrapper(theta):
    # simulate with experiment conditions
    return simul_LCA5(
        # we assume the first two parameters are beta and rho
        parameters=theta[:, :5],
        # we treat the third concentration parameter as an experimental condition
        # add 1 to deal with 0 values from Categorical distribution
        cond=theta[:, 2:] + 1,
    )

i.e., parameters takes the first five columns, and cond takes all columns after index 2.
This looks suspicious because I would expect theta and cond to be exclusive.
Also, it looks like in simul_LCA5, the cond argument is not even used.

Cheers,
Jan

[Jiashun97]

Dear Jan,

Thanks for your reply.

Q1: Is dl a continuous variable, that just happens to take only two values in the conditions?

dI is not an experiment condition. It is a parameter in the CA model that depends on the experimental condition. The goal is that in the inference stage, the behavioral data have three columns: RTs, choices, and experimental conditions. Then fitting the CA model to the data using MNLE, resulting in parameter estimates for dI1 (for exp. condition 1), dI2 (for exp. condition 2), I, k, a, ter. Only dI has two values because only dI is assumed to have different values for different the experiment conditions, and all the others are assumed to be the same across conditions. Therefore, we definitely need a prior for dI.

In this case, how should I construct simul_LCA5() and sim_wrapper()?

Shall I use dI1, dI2, I, k, a, ter as input for simul_LCA5() and it outputs RTs, choices, and conds. In this setting, one set of parameter values results in two rows of data like this:

Or I can use conds, dI, I, k, a, ter as input for simul_LCA5() and it outputs RTs, choices? In this case, how to specify the relationship between conds and parameters?

If I didn't explain it clearly, what I need is just a "depends_on" argument in the HDDM package. https://hddm.readthedocs.io/en/latest/howto.html#estimate-parameters-for-different-conditions

I appreciate your time answering my questions.

Best regards,
Jiashun

[janfb]

Hi @Jiashun97 ,

thank for sharing the details, and sorry, it always take a while until I find the time to look into this.

I now started implementing a feature that allows passing a batch of different condition-values that correspond to a batch of observed i.i.d x_o, see #1331. It will take some days until this is reviewed and merged, but we hope to make a new release before the end of the year.

With this new feature it will be possible to do the following:

• you have a simulator that has parameters and conditions that depend on each other in some way (that's up to you, it's just how the model is set up)
• during training, the parameters and conditions are treated as one thing, all contained in theta. Thus, it's a python function simulator(theta: Tensor) -> Tensor that takes theta = [parameters, conditions] and returns x=[rts, choices].
• you define a proposal over theta, e.g., a prior over the parameters and proposal distributions for the conditions. The proposals can be discrete or continuous.
• then you sample many theta ~ proposal and simulate corresponding data x ~ simulator(theta)
• then you train MNLE using (theta, x)

Only at inference time, you start to distinguish between parameters and conditions. Your observed data is now given (e.g., for a specific subject) as a table of i.i.d. (rts, choices, conditions). Therefore, you now want to fix conditions and do inference only over the parameters. To do so, you

• first define a potential for the learned MNLE estimator, the full proposal and the current batch of iid x_o (only rts and choices)
• and then obtain a "conditioned" potential, passing the corresponding batch of conditions, and the dimensions (columns in theta) you want to do inference over.

potential_fn = LikelihoodBasedPotential(estimator, proposal, x_o)
conditioned_potential_fn = potential_fn.condition_on(conditions, dims_to_sample=[0, 1])

This is the crucial step. You can then use the conditioned potential function to perform MCMC as usual.

You need to repeat this step for each subject or whenever the x_o changes. x_o and the conditions always need to match.

You can find the full code snippet here: https://github.com/sbi-dev/sbi/blob/multiple-conditions-in-iid-sampling/tutorials/Example_01_DecisionMakingModel.ipynb

Note that this is still WIP, the exact API might change a bit. But I hope this helps you figuring out how to set up your specific case.

Best wishes,
Jan