Imputed variables are not auto-transformed

[npschafer]

I’m looking to do imputation of partially observed random variables with positive support. Attempting to do so results in sampling values outside the support of the distribution and causes various downstream steps to fail, e.g., finding the MAP and variational inference.

import pymc3 as pm
from numpy.ma import masked_values

MASK_VALUE = -999.
observed_values = np.array([1., MASK_VALUE], dtype=float)
with pm.Model() as model:
    pm.Gamma("partially_observed", alpha=3., beta=1/0.1, observed=masked_values(observed_values, value=MASK_VALUE))
    map_estimate = pm.find_MAP()
print(map_estimate)
print(model.test_point)

{'partially_observed_missing': array([-0.7])}
{'partially_observed_missing': array([0.3])}

I originally raised this on the PyMC3 Discourse (https://discourse.pymc.io/t/imputation-for-partially-observed-random-variables-with-positive-support/6844) and @ckrapu suggested that I create an issue.

Versions and main components

• PyMC3 Version: '3.9.3'
• Theano Version: '1.0.5' (not sure if this is the correct version as things are somewhat confusing with having Theano, Theano-PyMC, and Aesara all co-existing)
• Python Version: 3.6.12
• Operating system: Linux
• How did you install PyMC3: pip

Thanks for all the great work and for your help with resolving the issue!

[ricardoV94]

I can replicate.

This is a very unusual example in that there is nothing to infer other than the missing value (other gamma parameters are fixed). If you add an unobserved parameter it seems things work fine:

with pm.Model() as model: 
    alpha = pm.Gamma('alpha', mu=3, sigma=.001) 
    pm.Gamma("partially_observed", alpha=alpha, beta=1/0.1, observed=masked_values(observed_values, value=MASK_VALUE)) 
    print(pm.find_MAP())

{'alpha_log__': array(1.09861256), 'partially_observed_missing': array([0.19999973]), 'alpha': array(3.00000081)}

If you just model this as an unobserved distribution it's also fine:

with pm.Model() as model: 
     pm.Gamma('not_observed', alpha=3, beta=1/0.1)
     print(pm.find_MAP())

{'not_observed_log__': array(-1.60943705), 'not_observed': array(0.20000017)}

Looks like a real edge case.

Did you have a more realistic example where this issue also appeared?

[npschafer]

Thanks for helping me to look into it @ricardoV94. Indeed, I ran into this in a much more complicated model with a real application. In my application, having all of the variables float freely makes the model unidentifiable. I need to condition on the value of one of them, which then sets the scale for the rest through its interaction with another set of random variables. I don't think that this problem is specific to MAP, as I initially ran into the problem while trying to perform variational inference. I chose to illustrate the problem with a simple example because that is what is usually considered most useful for debugging purposes.

[ricardoV94]

Could you give a minimal example with at least one free parameter (other than the missing distribution)?

If I had to guess it has to do with imputed distribution auto transformation getting lost / forgotten somewhere

[ricardoV94]

Nevermind I think I found the problem.

[ricardoV94]

So the problem is that the imputed variables are not being auto-transformed and when an invalid value is evaluated it will return a gradient of zero for the cost function, which the scipy optimizer will take as perfect convergence.

The same abnormal result can be obtained by doing this:

 with pm.Model() as model: 
     pm.Gamma("partially_observed", alpha=3., beta=1/0.1, transform=None) 
     print(pm.find_MAP())                                                                        

{'partially_observed': array(-0.7)}

This also means that the normal NUTS sampler will struggle with inputed variables, as seen by the number of divergences when running the original model:

with pm.Model() as model: 
    pm.Gamma("partially_observed", alpha=3., beta=1/0.1, observed=masked_values(observed_values, value=MASK_VALUE)) 
    trace = pm.sample() 

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 2 seconds.
There were 135 divergences after tuning. Increase `target_accept` or reparameterize.
There were 116 divergences after tuning. Increase `target_accept` or reparameterize.
There were 127 divergences after tuning. Increase `target_accept` or reparameterize.
There were 161 divergences after tuning. Increase `target_accept` or reparameterize.