How to increase weights of conditional log probability

I am using state space models for time series modeling using the following model

X_t = g(X_{t-1}) + N(0,\sigma^2) \\ y^i(t) \sim Poisson(f(z^i(t),X_t))

where i corresponds to ith sample at time t, g and f are non-linear models and \sigma is part of parameters to be estimated along with f and g. y^i(t) is a scaler observed value and for each i and t, there corresponds a feature vector z^i(t). The goal is to estimate g, f and \sigma. Note that for each time point t there are 400k samples, and there are total 100 time points.
For modeling state space I am using scan function

    x0_ar = pm.Normal("xo_ar", 0, sigma=1, initval = init_ar, shape=(latent_factors), rng=rng)
    sigmas_Q_ar = pm.InverseGamma('sigmas_Q_ar', alpha=3,beta=0.5, shape= (latent_factors), rng=rng)
    Q_ar = pt.diag(sigmas_Q_ar) 
    
    def step_simple(x, A, Q, bais_latent):
        innov = pm.MvNormal.dist(mu=0.0, tau=Q)    
        next_x = g(x) + innov 

        return next_x, collect_default_updates( [next_x]) 

    ar_states_pt, ar_updates = pytensor.scan(step_simple, 
                                              outputs_info=[x0_ar], 
                                              non_sequences=[A_ar, Q_ar, bais_latent],
                                              n_steps=T, 
                                              strict=True)
    assert(ar_updates)

    model_minibatch.register_rv(ar_states_pt, name='ar_states_pt', initval=pt.zeros((T, latent_factors)))
    ar_states = pm.Deterministic("ar", pt.concatenate([x0_ar.reshape((1,latent_factors)), ar_states_pt], axis=0))
    obs = pm.Poisson('obs', f(ar_states,z), observed=data_out, total_size=424560)

ar_states goes as input to f function. When I run the inference using ADVI, the estimates of ar_states do not follow the state space dynamics and they just fit the training data in whatever way. The reason being that ar_states should maximize the two likelihoods, one is for the observed data and on is for the state space dynamics. Due to large number of samples (>400k), likelihood for observed data just dominates the state space dynamics. Is there a way to extract conditional log probability of states and weigh them such that optimizer equally maximize the two log probabilities?

Hi @ricardoV94 and @jessegrabowski, any suggestions for the above question?

Is there a way to use logprob.conditional_logp#? I am thinking maybe I add something like 1000*Potential(conditional_logp(X_t)) to the model where 1000 is the weight to the conditional logp but conditional_logp takes value variable.

You can retrieve the value variable from model.rvs_to_values[X_t]

Thanks for your response. So I added the below line to my model

error = pm.Potential(“error”, pm.logprob.conditional_logp({ar_states_pt: model_minibatch.rvs_to_values[ar_states_pt]}))

but I got the below error
model = modelcontext(model)
var.name = model.name_for(name)
model.potentials.append(var)
model.add_named_variable(var, dims)
AttributeError: ‘dict’ object has no attribute ‘name’

conditional_logp returns a dictionary. If you are working with a single variable you may want to just use logp.

Got you, I have few questions based on the above discussion-

  1. So from what I understand, logp and conditional_logp would give the same thing if used for just one variable, is that right? I want to make sure they both are the same because in my case I am trying to extract conditional logp of X_t given X_t-1. In my model, only y is observed and the rest of the variables are estimated.
  2. In the above model, I have obs = pm.Poisson(‘obs’, f(ar_states,z), observed=data_out, total_size=424560). If I add error = pm.Potential(pm.logp((X_t, model.rvs_to_values[X_t]))) then it will override the Poisson term, right? or the model will create a potential for Poisson and add it to the error potential?

It will not override anything, it’s just added to the total model logp.

Here is a simple example and some code so you can investigate if things are doing what you want them to:

import numpy as np
import pymc as pm
from pymc.model.fgraph import clone_model

with pm.Model( ) as m1:
  x = pm.Exponential("x", 1)
  y = pm.Poisson("y", x, observed=[0, 1, 2])

with clone_model(m1) as m2:
  # Just add one to whatever the model joint-probability is
  pm.Potential("potential", pm.math.ones((1,)))

print(m1.point_logps())  # {'x': -1.0, 'y': -3.69}
print(m2.point_logps())  # {'x': -1.0, 'y': -3.69, 'potential': 1.0}
np.testing.assert_allclose(
    m1.compile_logp()({"x_log__": 0}) + 1,
    m2.compile_logp()({"x_log__": 0}),
)          
1 Like