PyMC's treatment of shape versus deterministic data, when a random variable's parameter is vector-valued

Question

I'm working on a problem with PyMC3 that makes me think I need to better understand how it deals with random variables whose parameters are vector-valued.

Data description and problem setup

I have $N$ samples of binomial data, $(n_i, k_i)$ where $i = 1...N$, that is, $N$ experiments where for the $i$th experiment $n_i$ trials were performed and $k_i$ of those were successes (so $0 \leq k_i \leq n_i$).

For each of the $N$ data points, I also have a feature $x_i$, which happens to be a real number between 0 and 1. But it's known for each of the $N$ data points and fixed.

I'd like to model the binomial data with a Binomial distribution, and then model the distribution's probability parameter $p$ as Beta-distributed, which is the Binomial's conjugate prior.

I have reason to believe that the Beta prior's parameters, $\alpha = m_\alpha x + b_\alpha$ and $\beta = m_\beta x + b_\beta$, that is, a parameter of the Beta is a linear function of $x_i$.

So each of my binomial data points can be said to $$ k_i \sim Binomial(n_i, p_i), $$ where $$p_i \sim Beta(\alpha_i, \beta_i)$$ and where $$\alpha_i = m_\alpha x_i + b_\alpha$$ and $$\beta_i = m_\beta x_i + b_\beta$$

Synthetic data generation

The following generates some sample data:

import numpy as np
from scipy.stats import beta, binom


Ndata = 1000

feature = np.random.rand(Ndata) # x's
trials = np.random.randint(1, 20, (Ndata,)) # n's

alphaSlope = 10.0    # m_alpha
alphaIntercept = 5.0 # b_alpha

betaSlope = -3.0.    # m_beta
betaIntercept = 8.0  # b_beta

obs = [binom(n, beta(al, be).rvs(1)[0]).rvs(1)[0] 
       for al, be, n in zip(
           alphaSlope * feature + alphaIntercept,
           betaSlope * feature + betaIntercept,
           trials)] # k's

PyMC model 1

I model this in PyMC (version 3.8) as:

import theano.tensor as tt
import pymc3 as pm


with pm.Model() as model:
    # weights: slope, intercept
    walpha = pm.Uniform('walpha', lower=-20, upper=20, shape=2)
    wbeta = pm.Uniform('wbeta', lower=-20, upper=20, shape=2)

    # parameters of Beta
    a = (walpha[0] * feature + walpha[1]).clip(1, 5000)
    b = (wbeta[0] * feature + wbeta[1]).clip(1, 5000)

    # prior and likelihood
    p = pm.Beta('p', alpha=a, beta=b, testval=0.7)
    x = pm.Binomial('x', n=trials, p=p, observed=obs)

    trace = pm.sample(10000, tune=2000, cores=5)

(Before proceeding, I understand I could have modeled this using PyMC's BetaBinomial random variable. I don't want to use this because in my actual problem, the Binomial isn't parameterized by probability $p$ but rather $p^\delta$, a nonlinear exponentiation.)

But the above PyMC model produces incorrect results: using Pandas to describe the MCMC trace with pm.trace_to_dataframe(trace).describe() reveals both $\alpha$ and $\beta$'s slope and intercept parameters to be far from true values:

          walpha__0     walpha__1      wbeta__0      wbeta__1             p
count  50000.000000  50000.000000  50000.000000  50000.000000  50000.000000
mean      19.870350     19.942757     13.847049     14.402839      0.589303
std        0.128837      0.057044      0.725568      0.465387      0.005422
min       18.562779     19.464832     10.930696     12.495120      0.563810
25%       19.819790     19.920913     13.356722     14.090275      0.585660
50%       19.909900     19.960104     13.842200     14.396608      0.589340
75%       19.962834     19.983450     14.335798     14.713298      0.593011
max       19.999999     19.999999     17.108659     16.592652      0.612058

PyMC's trace plot shows how bad the results are (pm.traceplot(trace)):

What exactly is the line, p = pm.Beta('p', alpha=a, beta=b, testval=0.7), doing? Note that a and b here are vector-valued (1000 by 1, since I have 1000 points of data), each a linear combination of two random variables, so perhaps I need to specify the shape of p manually? This is my second PyMC model:

PyMC model 2

This is the exact same model as above except I add a shape parameter to the Beta random variable, explicitly stating that it should be a 1000-long vector of a thousand different random variables:

with pm.Model() as model:
    # weights: slope, intercept
    walpha = pm.Uniform('walpha', lower=-20, upper=20, shape=2)
    wbeta = pm.Uniform('wbeta', lower=-20, upper=20, shape=2)

    # parameters of Beta
    a = (walpha[0] * feature + walpha[1]).clip(1, 5000)
    b = (wbeta[0] * feature + wbeta[1]).clip(1, 5000)

    # prior and likelihood
    p = pm.Beta('p', alpha=a, beta=b, testval=0.7, shape=len(obs)) # !!!
    x = pm.Binomial('x', n=trials, p=p, observed=obs)

    trace = pm.sample(10000, tune=2000, cores=5)

This model runs much slower (~200 draws per second, versus model 1's ~1500 draws per second), but produces reasonable results:

          walpha__0     walpha__1      wbeta__0      wbeta__1          p__0  \
count  50000.000000  50000.000000  50000.000000  50000.000000  50000.000000   
mean      13.071932      6.050287     -3.331585      9.634910      0.579830   
std        2.793833      1.244391      1.643823      1.603906      0.090586   
min        2.112455      2.620591    -11.579191      5.225143      0.208103   
25%       11.097158      5.167857     -4.386282      8.506066      0.518849   
50%       12.977088      5.926347     -3.280822      9.479486      0.581282   
75%       14.968447      6.794506     -2.206138     10.582560      0.642689   
max       19.995363     13.397899      2.877322     18.900355      0.880307   

The true values for walpha__0 walpha__1 wbeta__0 wbeta__1 are, recall, 10, 5, -3, and 8. (I'm a bit surprised that with 1000 data samples, and 50'000 MCMC iterations, it couldn't get closer to the true walpha__0 so maybe there's something weird going on.)

Model 3: BetaBinomial

As mentioned above, I can't use PyMC's BetaBinomial random in my actual application because I need a nonlinearly-transformed Beta prior on the Binomial parameter, but for just checking the above result, I tried it:

with pm.Model() as model:
    # weights: slope, intercept
    walpha = pm.Uniform('walpha', lower=-20, upper=20, shape=2)
    wbeta = pm.Uniform('wbeta', lower=-20, upper=20, shape=2)

    # parameters of Beta
    a = (walpha[0] * feature + walpha[1]).clip(1, 5000)
    b = (wbeta[0] * feature + wbeta[1]).clip(1, 5000)

    # prior and likelihood
    x = pm.BetaBinomial('x', observed=obs, n=trials, alpha=a, beta=b)

    trace = pm.sample(12000, tune=6000, cores=5)

With BetaBinomial, PyMC runs at ~600 draws per second, faster than model 2 above but slower than model 1. I didn't need to specify any shape parameter but it appears to have found the same results as model 2 above:

          walpha__0     walpha__1      wbeta__0      wbeta__1
count  50000.000000  50000.000000  50000.000000  50000.000000
mean      13.049388      6.060485     -3.359202      9.649070
std        2.737753      1.250730      1.657210      1.600937
min        3.351191      2.292545    -11.563000      4.695184
25%       11.122850      5.171321     -4.418419      8.514347
50%       12.953331      5.943487     -3.293241      9.501761
75%       14.914511      6.806098     -2.238418     10.626369
max       19.999812     13.462526      3.996577     18.123882

Questions

My main question is, am I correctly specifying my model by including a shape parameter in p = pm.Beta('p', alpha=a, beta=b, testval=0.7, shape=len(obs))? I think this is doing the right thing because an explicit BetaBinomial random variable returns extremely similar results, but I also think this could be improved because it runs so much slower than BetaBinomial (which I cannot use in my actual example). Is there a better way of expressing this?

Secondarily, what exactly does PyMC do with p = pm.Beta('p', alpha=a, beta=b, testval=0.7), that is, model 1's random variable where the alpha and beta parameters are vectors but p is apparently a scalar random variable? What kind of model is that statistically fitting? Is it perhaps just looking at the first elements of the alpha and beta vectors, thereby giving weird results for the slope and intercept terms?

Any intuition about how PyMC handles random variables with vector parameters, and the impact of the shape argument in those cases, would be helpful.

A Jupyter Notebook with all code used is available on Gist.

Answer 1

There is a lot to address here, and while I appreciate you providing so much detail, I will only focus on some of the most important issues I see.

So, the first thing I notice from looking at your model structure is that your model is perhaps misspecified. For one, the beta distribution is defined only for positive $\alpha, \beta$. However, your parameterization allows for them to be negative. I would suggest you model the mean of the beta as follows

$$ \operatorname{logit}\mu(x) = m_0 + m_1x$$

And then parameterize the beta distribution as

$$ p_i \sim \operatorname{Beta}\left(\dfrac{\mu(x)}{\kappa}, \dfrac{1-\mu(x)}{\kappa}\right)$$

You can either put a prior on $\kappa$ or model it in the same way as we modeled the mean. Here is an example of what I mean. In this model, I model both the $\mu$ and $\kappa$ as functions of the covariates

X = np.c_[np.ones_like(feature), feature]
n = X.shape[0]

with pm.Model():

    coef_mu = pm.Normal('coef_mu',0,1, shape = 2) #2 coefficients 
    coef_kappa = pm.Normal('coef_kappa',0,1, shape = 2) #2 coefficients 

    mu = pm.invlogit(pm.math.dot(X,coef_mu))
    kappa = pm.invlogit(pm.math.dot(X,coef_kappa))

    p = pm.Beta('p',mu/kappa, (1-mu)/kappa, shape=n)

    Y = pm.Binomial('Y',n=trials, p=p, observed=obs)

    trace = pm.sample(chains = 8)

This model samples 8000 samples on a 2017 dual core iMac in about 2.5 minutes with 0 divergences. The only warning that appears is that the effective number of samples for some parameters is smaller than 25%, but this isn't so bad as Stan only warns when the fraction of effective samples is smaller than 1%. You could do some clever algebra to perhaps see if you recover your initial parameter values, which I won't do here.

As to your other questions...

My main question is, am I correctly specifying my model by including a shape parameter in p = pm.Beta('p', alpha=a, beta=b, testval=0.7, shape=len(obs))

Yes and no. Again, your priors allow for negative $\alpha, \beta$ (if I understand the .clip method correctly). You have the shape of the parameters correct, but it might be better to model them in a way which preserves the boundaries.

Secondarily, what exactly does PyMC do with p = pm.Beta('p', alpha=a, beta=b, testval=0.7)

a and b are theano tensors of shape (1000,). Each element of those tensors defines a beta distribution. This expression essentially boils down to the assumption that the $p$ for observation $i$ comes from a beta distribution with parameters a[i] and b[i]. You can see now why bounding these parameters is so important.