Sampling parameters from a custom prior

Misc.
1

Hi!

I’m trying to use a custom, pre-generated prior in my model but it seems that the model training doesn’t sample from the prior at all.

The omega is [200, 1] and they are the parameters I want to fit. When I replace the Empirical distribution with a Normal(0., 1.0) I get ok results and nice distributions for each omega. However when I use the Empirical, each omega have a constant value in different MCMC samples.

Any idea why this is the case? I’m not sure if Empirical is the correct way to do this, but it was the only thing that was even close to what I wanted to do.

Thanks a lot for answering!


predictive = pyro.poutine.block(predictive, hide_all=True)
A = torch.tensor(A, dtype=torch.float32).cuda()
y_data = torch.tensor(y_data, dtype=torch.float32)

f_gpu = torch.tensor(f, dtype=torch.float32)

# Define the model in Pyro
def inverse_model(A, y_data):
    n = A.size(1)  # Number of features
    
    # Prior distribution for omega, sampled from a tensor shaped [200, 5000]
    prior_samples = prior['obs'].T.cuda()  
    prior_weights = torch.ones(prior_samples.shape).cuda()
    omega = pyro.sample('omega', dist.Empirical(prior_samples, prior_weights).expand([n]).to_event(1))
    
    sigma = pyro.sample("sigma", dist.Uniform(0., 0.05)).to(A.device)

    # Linear model
    mu = A @ omega
    

    with pyro.plate("data", y_data.shape[0]):
        pyro.sample("obs", dist.Normal(mu, sigma), obs=y_data)
    
    return omega

inverse_nuts_kernel = pyro.infer.NUTS(inverse_model, init_strategy=pyro.infer.autoguide.init_to_sample())
mcmc = pyro.infer.MCMC(inverse_nuts_kernel, num_samples=10, warmup_steps=10)

with pyro.validation_enabled():
    res = mcmc.run(A, y_data)
2

Instead of using Empirical distribution, you can use the Categorical distribution and use enumeration: Inference with Discrete Latent Variables — Pyro Tutorials 1.9.1 documentation

3

Hi!
Thanks for the tip!

If I understood correctly, I would want to create indexing variables for my prior (x and y in the examples) and sample those? I tried to do it like this:

@config_enumerate
def inverse_model(A, y_data):
    n = A.size(1)  # Number of features
    
    # Prior distribution for omega, sampled from the empirical distribution
    prior_samples = pyro.param('prior', prior_bnn.T.cuda())
    
    print(f'prior_samples shape: {prior_samples.shape}')

    x = pyro.sample("x", dist.Categorical(torch.ones(200)))
    y = pyro.sample("y", dist.Categorical(torch.ones(5000)))

    with pyro.plate("z_plate", 200):
        omega = Vindex(prior_samples)[...,:,x] 

    sigma = pyro.sample("sigma", dist.Uniform(0., 0.05)).to(A.device)

    # Linear model
    print(f'omega shape: {omega.shape}')
    mu = A @ omega

    # Likelihood (sampling distribution) of observations
    with pyro.plate("data", y_data.shape[0]):
        pyro.sample("obs", dist.Normal(mu, sigma), obs=y_data)
    
    return omega

but I can’t understand why I get the output below. It seems that the omega is changing its shape and I don’t understand why it does that.

Warmup:   0%|          | 0/1010 [00:00, ?it/s]
prior_samples shape: torch.Size([200, 5000])
omega shape: torch.Size([200])
prior_samples shape: torch.Size([200, 5000])
omega shape: torch.Size([200, 1, 200])

RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [200, 200] but got: [200, 1].
Trace Shapes:                  
 Param Sites:                  
        prior          200 5000
Sample Sites:                  
       x dist                 |
        value      200   1    |
       y dist                 |
        value 5000   1   1    |
 z_plate dist                 |
        value          200    |
   sigma dist                 |
        value                 |
5

Updated question:

I did the enumeration with a sequential plate:

    n = A.size(1)  # Number of features
    prior_samples = pyro.param('prior', prior_bnn.cuda()) # shape [5000, 200]
    
    sigma = pyro.sample("sigma", dist.Uniform(0., 0.05)).to(A.device)

    x = []
    omega = torch.zeros([n, 1])
    for ii in range(0, n):
        x.append(pyro.sample(f"x_{ii}",
                            dist.Categorical(torch.ones(5000)),
                            infer={"enumerate": "sequential"}))
        omega[ii] = prior_samples[x[ii], ii]
    assert omega.shape == (n, 1)

This generates me a random sample from the prior for each step ii. These are then collected to the variable omega. Omega is the variable that I would ultimately want to get as a result.

I multiply the omega with a matrix A and then compare the result to my data in another plate:

    mu = torch.matmul(A, omega)    
    with pyro.plate("data", y_data.shape[0]):
        pyro.sample("obs", dist.Normal(mu.squeeze(-1), sigma), obs=y_data)
    return omega

I’m able to run the code with SVI, but the problem is that when in Predictive I use return_sites=["_RETURN"], The omegas follow the initial prior instead of the desired output. This means that the model isn’t actually learning the parameters for the n Categorical distributions x_ii, since with more predicted omegas the mean goes towards the mean of the prior.
I think now the model is just giving samples of the sigma parameter. How could I get it to also sample omega (and x_ii)?

Here is also the full model:

A = torch.tensor(A, dtype=torch.float32).cuda()    # shape [100, 200]
y_data = torch.tensor(y_data, dtype=torch.float32) # shape [100, 1]
f_gpu = torch.tensor(f, dtype=torch.float32) # shape [100, 1]


@config_enumerate
def inverse_model(A, y_data):
    
    n = A.size(1)  # Number of features
    prior_samples = pyro.param('prior', prior_bnn.cuda()) # shape [5000, 200]
    
    sigma = pyro.sample("sigma", dist.Uniform(0., 0.05)).to(A.device)

    x = []
    omega = torch.zeros([n, 1])
    for ii in range(0, n):
        x.append(pyro.sample(f"x_{ii}",
                            dist.Categorical(torch.ones(5000)),
                            infer={"enumerate": "sequential"}))
        omega[ii] = prior_samples[x[ii], ii]
    assert omega.shape == (n, 1)
    
    # Want to compare this to y_data
    # and get omega as the solution
    mu = torch.matmul(A, omega)
        
    with pyro.plate("data", y_data.shape[0]):
        pyro.sample("obs", dist.Normal(mu.squeeze(-1), sigma), obs=y_data)
    return omega


def guide(A, y_data):
    n = A.size(1)  # Number of features

    q_x_probs = pyro.param('q_x_probs', torch.ones(5000) / 5000,
                           constraint=constraints.simplex)

    x = []
    for ii in range(0, n):
        x.append(pyro.sample(f"x_{ii}",
                            dist.Categorical(q_x_probs),
                            infer={"enumerate": "sequential"}))
    
    pyro.sample("sigma", dist.Uniform(0., 0.05))
6

You will want to use infer_discrete to get samples for discrete latent variables. There are a couple of tutorials for it: Search — Pyro Tutorials 1.9.1 documentation