How to make mixture model code more concise

Questions v5
1

I am trying to implement a mixture model of logistic regressions. A long version is working, but I am trying to make the model more concise and I am struggling to get it right, since I don’t have much experience with pymc. Specifically, I haven’t yet figured out how to use shapes appropriately, which is probably responsible for my error.

I ran this model on fake data I created, to see whether the mixtures fit the parameters. It worked well with the following code, i.e. it discovered for all three faked users appropriate mixed parameters.

k = 2

with pm.Model() as long_model:
    w_1 = pm.Dirichlet('w_1', a=np.ones(k))
    w_2 = pm.Dirichlet('w_2', a=np.ones(k))
    w_3 = pm.Dirichlet('w_3', a=np.ones(k))
    
    betas = pm.Normal('betas', mu=np.zeros((5, k)), sigma=2, shape=(5, k))
    mixed_betas_1 = pm.NormalMixture('mixed_betas_1', w=w_1, mu=betas, shape=(5,))
    mixed_betas_2 = pm.NormalMixture('mixed_betas_2', w=w_2, mu=betas, shape=(5,))
    mixed_betas_3 = pm.NormalMixture('mixed_betas_3', w=w_3, mu=betas, shape=(5,))
    
    bern_1 = pm.Bernoulli('bern_1', p=pm.invlogit(first_X @ mixed_betas_1), observed=first_y)
    bern_2 = pm.Bernoulli('bern_2', p=pm.invlogit(second_X @ mixed_betas_2), observed=second_y)
    bern_3 = pm.Bernoulli('bern_3', p=pm.invlogit(third_X @ mixed_betas_3), observed=third_y)
    inference = pm.sample()

I came up with the following short version:

k = 2
user_count = 3

with pm.Model() as short_model:
    w = pm.Dirichlet('w', a=np.ones((user_count, k)))
    
    betas = pm.Normal('betas', mu=np.zeros((5, k)), sigma=2, shape=(5, k))
    stacked_mu = pm.math.stack(user_count * [betas], axis=1)
    mixed_betas = pm.NormalMixture('mixed_betas', w=w, mu=stacked_mu, shape=(5, user_count))
    
    stacked_ps = pm.math.stack([pm.math.dot(user_X[:, :, i], mixed_betas[:, i]) for i in range(user_count)], 1)
    bern = pm.Bernoulli('bern', p=pm.invlogit(stacked_ps), observed=user_y)
    inference = pm.sample()

The error message includes:

ValueError: Shape mismatch: A.shape[1] != x.shape[0]
Apply node that caused the error: CGemv{inplace}(AllocEmpty{dtype='float64'}.0, TensorConstant{1.0}, TensorConstant{[[0.304049..83757717]]}, Subtensor{::, int64}.0, TensorConstant{0.0})
Toposort index: 14
Inputs types: [TensorType(float64, (None,)), TensorType(float64, ()), TensorType(float64, (10000, 3)), TensorType(float64, (None,)), TensorType(float64, ())]
Inputs shapes: [(10000,), (), (10000, 3), (5,), ()]
Inputs strides: [(8,), (), (120, 40), (24,), ()]
Inputs values: ['not shown', array(1.), 'not shown', array([ 0.8743515 , -0.69598828, -1.01642807,  0.88118696, -0.63481115]), array(0.)]
Outputs clients: [[InplaceDimShuffle{0,x}(CGemv{inplace}.0)]]

I struggle to understand how I messed up the shapes and any pointers would be appreciated. Are there other ways to make the code more concise.

2

Getting the shape right in mixture distribution is tricky, usually my approach is to build it step by step and keep checking model.initial_point() and model.point_logps(), for example in this case I will start with:

k = 2
user_count = 3

with pm.Model() as short_model:
    w = pm.Dirichlet('w', a=np.ones((user_count, k)))
    
    betas = pm.Normal('betas', mu=np.zeros((5, k)), sigma=2, shape=(5, k))
    mixed_betas = pm.NormalMixture('mixed_betas', w=w, mu=betas[:, None, :], shape=(5, user_count))

short_model.initial_point()
short_model.point_logps()

Note here I add an extra dim to beta by doing betas[:, None, :], this should be more efficient than pm.math.stack.

Then I will add the remaining of the random variables:

import aesara.tensor as at

with short_model:
    # The shape of the operation is (i, 5, user_count) * (5, user_count) -> (i, user_count)
    # it is easier with einsum, but aesera does not have that yet.
    ps = pm.math.sum(user_X * mixed_betas, axis=1)
    bern = pm.Bernoulli('bern', p=pm.invlogit(ps), observed=user_y)

Let me know if this work on your real data.

3 Likes
3

I was about to write that it didn’t work, but then I realised that the axes of my user_X data were swapped relative to what your implementation assumed. Works well now! Thank you for your help!

2 Likes