Fixing parameters of the guide (e.g., loc) and optimizing the rest (e.g., scale)

Misc.
1

I have a Bayesian neural network implemented as a Pyro model (see bottom). Through some other means (e.g., MAP estimation) I have identified a mode in the posterior p(w|D) of the model. My goal is to essentially estimate the posterior around this point as a multivariate Gaussian, focusing on capturing the local curvature by estimating scale / covariance parameters. So, is there a way to specify a guide for the model as a Gaussian with the mean parameters fixed and the scale / covariance parameters as the variance parameters to learn? I am a little new to pyro so if someone could provide some guidance as to how to construct the guide of the model as I have described or tell me if this is/isn’t possible, I’d much appreciate it. Thanks!

Also, I understand that this is conceptually very similar to the Laplace approximation. I’ve had trouble utilizing the Laplace approximation due to matrix singularity in the hessian of the posterior and decided to try the approach described above.

def bayesian_cnn(x, y=None):
    
    all_params = pyro.sample("all_params", dist.MultivariateNormal(torch.zeros(1283), torch.eye(1283)))
    
    #perform the forward pass
    x = torch.relu(
        F.conv1d(x, 
                 weight = all_params[:20].reshape([2,1,10]),
                 bias = all_params[20:22].reshape([2]),
                 stride = 1,
                 padding = 'same'))
    
    x = torch.flatten(x,start_dim=1,end_dim=-1)
    x = torch.relu(
        F.linear(x,
                 weight = all_params[22:822].reshape([20,40]) ,
                 bias = all_params[822:842].reshape([20]))
        )
    x = torch.relu(
        F.linear(x,
                 weight = all_params[842:1242].reshape([20,20]),
                 bias = all_params[1242:1262].reshape([20]))
        )
    
    x = F.linear(x,
                 weight = all_params[1262:1282].reshape([1,20]) ,
                 bias = all_params[1282:1283].reshape([1]) )
    
    
    mu = x.squeeze()
    sigma = pyro.sample("sigma", dist.Uniform(0., 1))
    

    with pyro.plate("data", x.shape[0]):
        obs = pyro.sample("obs", dist.Normal(mu, sigma**2), obs=y)
        
        
        return obs
2

I think you can look at the param store and freeze the parameters to optimize there. see e.g. Bayesian Regression - Introduction (Part 1) — Pyro Tutorials 1.9.1 documentation