How to create a Neural Linear Model?

Hi! I’m new to the Bayesian framework, and currently I’m trying to quantify uncertainty in one of my neural networks. I’m working through this paper: Bayesian Classifiers With Out Of Distribution Uncertainty where a Bayesian Logistic Regression model is required.

The thing is, I don’t how where to start with this. Any help would be appreciated, as my undergraduate honours thesis depends on this :sweat_smile:!

I would start by reading the NumPyro Bayesian regression tutorial. Presumably you would want to do something very similar, except your inputs are features from a pretrained neural network rather than raw data.

1 Like

Thanks for your quick reply! I created the following model below, but I am not getting the results I expected (which may be due to the other, unrelated parts of my methodology). Just to make sure I have everything right in terms of how I am using numpyro to follow the previously attached figure, does the following model make sense?

def logistic_regression(output_of_neural_network, labels=None, L):
    output_of_neural_network = phi_theta(x)
    y | x = cat(softmax(W^T*phi_theta(x))) (if we are distinguishing between the different L classes)
    Since we want a bayesian logistic regression model that distinguishes between in and out of distribution data,
    I assume that we can re-write above as:
    b | x = bernoulli(softmax(W^T*phi_theta(x)))
    where b is either 0 for an in distribution class or 1 for an out of distribution class
    W ~ P(W)
    W = numpyro.sample('weights', dist.Normal(jnp.zeros(L), jnp.ones(L)))
    intercept = numpyro.sample('intercept', dist.Normal(0.0, 5.0))
    softmax = jnn.softmax(W*output_of_neural_network)
    logits = jnp.sum(softmax + intercept, axis=-1)
    return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)     # b | x
#     return numpyro.sample('obs', dist.Categorical(probs=softmax), obs=labels) # y | x

The line

softmax = jnn.softmax(W*output_of_neural_network)

is performing an elementwise multiplication of your weights and features, rather than a matrix multiplication (jnp.matmul), and you are adding intercept after the softmax nonlinearity rather than before. You also don’t need to use softmax to compute logits, which are scalars and in log-space.

Note that this model will not train your neural network’s weights, which you must do separately beforehand.

1 Like

Ah right, yep, I figured it out. It works! Thanks again for the help! :grinning: