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

[Jiashun97]

Hi Jan,

Thanks for your implementation of the conditioned potential function! I have reproduced the tutorial code with my LCA model. However, I still cannot manage to vary certain parameters across conditions. In the function condition_on_theta there is an argument called dims_global_theta. Now I set it to [0,1,2,3,4] because I have 5 parameters in the LCA model and there are 2 experimental conditions. If I want to only make the first parameter free to vary across conditions, shall I set dims_global_theta=[1,2,3,4]? The ultimate goal is to obtain 6 (instead of 5) parameter estimates in the inference stage, in which the first parameter should have two values for the two conditions respectively and the other four are fixed across the two conditions. In your comments of the LikelihoodBasedPotential it says

dims_global_theta: The indices of the columns in theta that will be sampled, i.e., that not conditioned.

But this lead to an assertion error saying that dims_global_theta must match the number of parameters to sample. Or we should pass the first parameter to local_theta? I have also attached my code. It would be great if you can help with this. I really appreciate your time and contribution!
SBI with multiple experimental conditions 2.ipynb.zip

Best regards,
Jiashun

[janfb]

Hi Jiashun,

great that you have reproduced the tutorial!

The dims_global_theta should be a list with a column indices that correspond to the global variables, i.e., those parameters that are not conditions. The local_theta should then contain the values of the local variables, i.e., the conditions. Thus, if you have 5 parameters in total in overall theta and two of those are conditions (let's say the ones at index 1 and 3), then you would pass dims_global_theta=[0, 2, 4, 5] and set local_theta=theta[:, [1, 3]] (or others values defined in your experiments).

Does this make sense?

If this does not work, can you please share the error with stack trace here?
(sorry, I don't have capacity to look at the zipped code at the moment).

Cheers,
Jan

[Jiashun97]

Hi Jan,

Thanks for your reply. Now the code works for me! I was also wondering if I can also use .map() method for the mnle_posterior? If I just use the same argument as before, I got an assertion error.

mnle_posterior.set_default_x(x_o)
mnle_map = mnle_posterior.map(show_progress_bars=True,num_iter=1000)

`AssertionError Traceback (most recent call last)
Cell In[81], line 2
1 mnle_posterior.set_default_x(x_o)
----> 2 mnle_map = mnle_posterior.map(show_progress_bars=True,num_iter=1000)

File ~/Desktop/LCA/sbi/inference/posteriors/mcmc_posterior.py:982, in MCMCPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
927 def map(
928 self,
929 x: Optional[Tensor] = None,
(...)
937 force_update: bool = False,
938 ) -> Tensor:
939 r"""Returns the maximum-a-posteriori estimate (MAP).
940
941 The method can be interrupted (Ctrl-C) when the user sees that the
(...)
980 The MAP estimate.
981 """
--> 982 return super().map(
983 x=x,
984 num_iter=num_iter,
985 num_to_optimize=num_to_optimize,
986 learning_rate=learning_rate,
987 init_method=init_method,
988 num_init_samples=num_init_samples,
989 save_best_every=save_best_every,
990 show_progress_bars=show_progress_bars,
991 force_update=force_update,
992 )

File ~/Desktop/LCA/sbi/inference/posteriors/base_posterior.py:273, in NeuralPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
271 if self._map is None or force_update:
272 self.potential_fn.set_x(self.default_x)
--> 273 self._map = self._calculate_map(
274 num_iter=num_iter,
275 num_to_optimize=num_to_optimize,
276 learning_rate=learning_rate,
277 init_method=init_method,
278 num_init_samples=num_init_samples,
279 save_best_every=save_best_every,
280 show_progress_bars=show_progress_bars,
281 )
282 return self._map

File ~/Desktop/LCA/sbi/inference/posteriors/base_posterior.py:198, in NeuralPosterior._calculate_map(self, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars)
195 else:
196 raise ValueError
--> 198 return gradient_ascent(
199 potential_fn=self.potential_fn,
200 inits=inits,
201 theta_transform=self.theta_transform,
202 num_iter=num_iter,
203 num_to_optimize=num_to_optimize,
204 learning_rate=learning_rate,
205 save_best_every=save_best_every,
206 show_progress_bars=show_progress_bars,
207 )[0]

File ~/Desktop/LCA/sbi/utils/sbiutils.py:935, in gradient_ascent(potential_fn, inits, theta_transform, num_iter, num_to_optimize, learning_rate, save_best_every, show_progress_bars, interruption_note)
933 while iter_ < num_iter:
934 optimizer.zero_grad()
--> 935 probs = potential_fn(theta_transform.inv(optimize_inits)).squeeze()
936 loss = -probs.sum()
937 loss.backward()

File ~/Desktop/LCA/sbi/inference/potentials/base_potential.py:124, in CallablePotentialWrapper.call(self, theta, track_gradients)
119 """Call the callable potential function on given theta.
120
121 Note, x_o is re-used from the initialization of the potential function.
122 """
123 with torch.set_grad_enabled(track_gradients):
--> 124 return self.potential_fn(theta=theta, x_o=self.x_o)

File /var/folders/5g/vqmn12yd48z6l6jf2s2yy62r0000gn/T/ipykernel_53021/924198916.py:147, in LikelihoodBasedPotential.condition_on_theta..conditioned_potential(theta, x_o, track_gradients)
143 def conditioned_potential(
144 theta: Tensor, x_o: Optional[Tensor] = None, track_gradients: bool = True
145 ) -> Tensor:
146 assert (
--> 147 len(dims_global_theta) == theta.shape[1]
148 ), "dims_global_theta must match the number of parameters to sample."
149 global_theta = theta[:, dims_global_theta]
150 x_o = x_o if x_o is not None else self.x_o

AssertionError: dims_global_theta must match the number of parameters to sample.
`

Best regards,
Jiashun

[janfb]

Great to hear that this is working now!

You should be able to do .map() with an MNLE-based posterior as well. So this seems to be a bug. Could you please create an issue where you also post the code and the error stack like here? Thanks 🙏

[Jiashun97]

Great to hear that this is working now!

You should be able to do .map() with an MNLE-based posterior as well. So this seems to be a bug. Could you please create an issue where you also post the code and the error stack like here? Thanks 🙏

Sure! Thanks for your quick reply.