HSSM Seems to Return Priors as Posterior

[theonlydvr]

I've been working on implementing an RLDDM in HSSM following some advice from @krishnbera which so far has worked great! One of the issues I've run into though has been speed so I've tried working on setting up a pytensor version of the code so I can take advantage of the faster samplers and a GPU. To test this approach out, I've started by working with a simple RL-only case rather than the full RLDDM. However, whenever I try sampling this model, I seem to just get the priors back. I've included my code and some other model details below. Any help/advice would be greatly appreciated!

def update_Q(choice, reward, alpha, reset, Qs):
    if reset:
        Qs = pt.set_subtensor(Qs[0], 0.5)
        Qs = pt.set_subtensor(Qs[1], 0.5)
    else:
        pe = reward - Qs[(choice + 1) // 2]
        Qs = pt.set_subtensor(Qs[(choice + 1) // 2], Qs[(choice + 1) // 2] + alpha * pe)
    return Qs

def pytensor_llik(data, rl_alpha, rl_beta, response_s, feedback_s, split_by):
    # Convert to tensors
    responses_ = data[:,1].astype('int32')
    feedback_s_ = pt.as_tensor_variable(feedback_s, dtype="int32")
    response_s_ = pt.as_tensor_variable(response_s, dtype="int32")
    split_by_ = pt.as_tensor_variable(split_by, dtype="int32")
    
    # Compute Q-values
    reset = pt.concatenate([pt.zeros(1), pt.neq(pt.diff(split_by_), 0)], axis=0)
    Qs = pt.as_tensor_variable([0.5, 0.5], dtype="float64")
    Qs, _ = pytensor.scan(fn=update_Q, sequences=[response_s_, feedback_s_, rl_alpha, reset], outputs_info=[Qs])
    
    # Compute trial-wise log-likelihood
    Qs = Qs * pt.shape_padright(rl_beta,1)
    logp_actions = Qs - pt.logsumexp(Qs, axis=1, keepdims=True)
    return logp_actions[pt.arange(responses_.shape[0]), (responses_ + 1) // 2]

model_config_combined = {
    'list_params': ['rl.alpha', 'rl.beta'],
    'bounds': {
        'rl.alpha': (0, 1.0),
        'rl.beta': (0, np.inf)
    },
    'backend': 'pytensor',
    'extra_fields': ["response_s", "feedback_s", "split_by"],
    'choices': [-1, 1]
}

model = hssm.HSSM(
                data,
                model='2AB_softmax',
                model_config=model_config_combined,
                loglik=pytensor_llik,
                loglik_kind="analytical",
                # link_settings="log_logit",
                p_outlier=None,
                lapse=None,
                )

res = model.sample(draws=1200, tune=1200, chains=1, cores=4, idata_kwargs = {'log_likelihood': True}, discard_tuned_samples=False)

[AlexanderFengler]

Hi @theonlydvr,

I don't have an immediate answer looking at this, need to actually look at the likelihood.
I did see something like this before, and in that case there was a mix-up of columns so that the connection between parameters and output didn't respect the expected relation at all --> parameters didn't really affect anything, so you would get back the prior.

However can't confirm this is the case here from just looking at it. Will get back.

[theonlydvr]

That seems like it could be the source. Would that sort of thing happen if the rt column had no connection to the parameters (as would be the case in this pure RL test model)? If it would be helpful for checking things, I can also send over the test data.

[theonlydvr]

You were right, there was a subtle decoupling between the parameters and output (I am not used to some of the low level pytensor syntax). I've attached an updated version of the RL code that seemed to work very well for parameter recovery and was compatible with the NUTS sampler. Hopefully the transition to the full RLDDM version will be relatively straightforward!

from pytensor.ifelse import ifelse

def update_Q(choice, reward, alpha, reset, Qs):
    pe = reward - Qs[(choice + 1) // 2]
    q0 = ifelse(pt.eq(reset, 1), pt.as_tensor_variable(0.5, dtype="float64"), Qs[0])
    q1 = ifelse(pt.eq(reset, 1), pt.as_tensor_variable(0.5, dtype="float64"), Qs[1])
    q0 = ifelse(pt.eq(reset, 0) & pt.eq((choice + 1) // 2, 0), q0 + alpha * pe, q0)
    q1 = ifelse(pt.eq(reset, 0) & pt.eq((choice + 1) // 2, 1), q1 + alpha * pe,  q1)
    Qs = pt.set_subtensor(Qs[0], q0)
    Qs = pt.set_subtensor(Qs[1], q1)
    return Qs

def pytensor_llik(rl_data, rl_alpha, rl_beta, response_s, feedback_s, split_by):
    # Convert to tensors
    responses_ = rl_data[:,1].astype('int32')
    feedback_s_ = pt.as_tensor_variable(feedback_s, dtype="int32") # feedback on preceding trial
    response_s_ = pt.as_tensor_variable(response_s, dtype="int32") # response on preceding trial
    split_by_ = pt.as_tensor_variable(split_by, dtype="int32") # session ID
    
    # Compute Q-values
    reset = pt.concatenate([pt.ones(1), pt.neq(pt.diff(split_by_), 0)], axis=0)
    Qs = pt.as_tensor_variable([0.5, 0.5], dtype="float64")
    Qs, _ = pytensor.scan(fn=update_Q, sequences=[response_s_, feedback_s_, rl_alpha, reset], outputs_info=[Qs])
    
    # Compute trial-wise log-likelihood
    Qs = Qs * pt.shape_padright(rl_beta,1)
    logp_actions = Qs - pt.logsumexp(Qs, axis=1, keepdims=True)
    return logp_actions[pt.arange(responses_.shape[0]), (responses_ + 1) // 2]

[krishnbera]

Hi Evan, this is great. Do you mean this code was directly compatible with the NUTS sampler? Or did you supply a pytensor Op for gradients separately?

[theonlydvr]

It's directly compatible with the NUTS sampler without a pytensor Op. I was able to specify it as an "analytical" model during setup with no issues.