Numpyro: bnn.py example with another artificial dataset

[milongo]

Hi,

I'm trying train a Bayesian Neural Network in a toy dataset of the form
y = xsin(x) + 0.3e_1 + x * 0.3e_2 with e_1, e_2 distributed as Gaussians with 0 mean and unit variance.
I'm using the following code: http://pyro.ai/numpyro/bnn.html

with the following changes to the data generating function:

# create artificial regression dataset
def get_data(n_train, d_x):
    d_y = 1  # create 1d outputs
    onp.random.seed(0)
    x_train = np.linspace(0, 10, n_train)
    # x_train = np.expand_dims(x_train, -1)
    eps_1 = onp.random.normal(loc=0, scale=1, size=(len(x_train)))
    eps_2 = onp.random.normal(loc=0, scale=1, size=(len(x_train)))
    y_train = x_train * np.sin(x_train) + 0.3 * eps_1 + 0.3 * x_train * eps_2
    x_test = np.linspace(-4, 14, 5000)

    x_train = np.expand_dims(x_train, -1)
    y_train = np.expand_dims(y_train, -1)
    x_test = np.expand_dims(x_test, -1)
    return x_train, y_train, x_test

With the following result:

Everything else is exactly the same. Could I have some pointers to get a BNN to train on this toy dataset?

Thanks

[fehiepsi]

I think that you can try the following enhancements:

• normalize y: otherwise, it is hard to capture the variance of y.
• add bias for all layers, e.g.

b1 = numpyro.sample("b1", dist.Normal(0, 1), sample_shape=(D_H,))
z1 = nonlin(b1 + np.matmul(X, w1))

• change hidden size
• change the number of hidden layers

I think a small NN can train this dataset, so does BNN. If you still can not make it work, please let me know. In the worst case, you can try Gaussian Process. :D

Btw, we have a forum to discuss about issues like this. It would be great to post the question there, so the discussions don't get lost :).

[milongo]

Thanks for the answer! I'll try normalizing and adding biases. I'm trying to reproduce a paper experiment in which they train a BNN with a 1 layer MLP with 50 neurons, so I think I'm limited in terms of architecture. I'll let you know how it goes.

[milongo]

Normalizing worked! Off-topic question, I'm guessing the --num-warmup parameter controls the number of iterations, but no matter what number I put, it runs for 2000 more iterations than that.
Is there a reason for this? Additionally, it doesn't seem to scale linearly? Running --num-warmup 1000 vs 5000 iterations makes 3000 vs 7000 iterations, but it in the first case its finished in about 20 minutes, whereas the second case is not even at 10% after 20 minutes.

[fehiepsi]

it runs for 2000 more iterations than that.

We have num_samples to control the number of samples we want to collect from MCMC (after num_warmup adaptation steps). So basically, MCMC runs num_warmup + num_samples steps.

in the first case its finished in about 20 minutes

I think that adding bias to z1, z2, z3 can make the inference more stable. But your result is really slow. Could you share your system information, including numpyro/jax/jaxlib versions?

[martinjankowiak]

@milongo note that the example input data has a constant feature of all 1s. this basically adds "bias" units to the nn. if you augment your input data with 1s, you should get more reasonable performance.

[fehiepsi]

a constant feature of all 1s. this basically adds "bias" units

Oh, nice trick!

@milongo I just took a quick run on your dataset and it just took me 44s with default args. Could you please check again your timing?

bnn_plot.pdf

[milongo]

@fehiepsi I should have mentioned I increased the number of neurons to 50.
Using the default arguments is indeed very fast. But it seems like messing with them quickly increases the time...

For instance, trying these arguments

parser.add_argument("-n", "--num-samples", nargs="?", default=2000, type=int)
parser.add_argument("--num-warmup", nargs='?', default=5000, type=int)
parser.add_argument("--num-chains", nargs='?', default=1, type=int)
parser.add_argument("--num-data", nargs='?', default=500, type=int)
parser.add_argument("--num-hidden", nargs='?', default=50, type=int)
parser.add_argument("--device", default='cpu', type=str, help='use "cpu" or "gpu".') 

Results in about 10 steps of 7000 in about 4 minutes.

This is what I get with 50 neurons, 500 samples of the dataset and everything else the same. After normalizing, too! This is quite similar to the paper I'm reproducing.

[fehiepsi]

num_hidden=50 is pretty large for MCMC to take care of (with this, the latent variable w2 has more than 2500 dimensions). So I think the slowness is acceptable... Does the paper you try to reproduce mention about timing?

[milongo]

Nothing about timing.. But there are other datasets with more samples and bigger dimension. In any case, BNNs are a baseline and not their main contribution so I think it's fine for now. Thanks for your help!

[milongo]

Hi,

I'm sorry if this not the correct place to post this but I tried signing up in the forum and got some error.
I'm revisiting this problem but trying a 1 layer MLP as follows:

def model(x_train, y_train, d_h):
    d_x, D_Y = x_train.shape[1], 1

    # sample first layer (we put unit normal priors on all weights)
    w1 = numpyro.sample("w1", dist.Normal(np.zeros((d_x, d_h)), np.ones((d_x, d_h))))  # d_x d_h
    b1 = numpyro.sample("b1", dist.Normal(0, 1), sample_shape=(d_h,))
    z1 = b1 + nonlin(np.matmul(x_train, w1))
    # z1 = nonlin(np.matmul(x_train, w1))   # N d_h  <= first layer of activations

    # sample second layer
    w2 = numpyro.sample("w2", dist.Normal(np.zeros((d_h, D_Y)), np.ones((d_h, D_Y))))  # d_h d_h
    b2 = numpyro.sample("b2", dist.Normal(0, 1), sample_shape=(D_Y,))
    z2 = b2 + np.matmul(z1, w2)  # N d_h  <= second layer of activations

    # sample final layer of weights and neural network output
    # w3 = numpyro.sample("w3", dist.Normal(np.zeros((d_h, D_Y)), np.ones((d_h, D_Y))))  # d_h D_Y
    # z3 = np.matmul(z2, w3)  # N D_Y  <= output of the neural network

    # we put a prior on the observation noise
    prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0))
    sigma_obs = 1.0 / np.sqrt(prec_obs)

    # observe data
    numpyro.sample("y_train", dist.Normal(z2, sigma_obs), obs=y_train)

As before, the data is normalized.

And this is what I'm getting with 10 hidden neurons, 2000 num_samples and 3000 num_warmup

Why are the results so dramatically different when using 1 layer vs 2 layers? Or am I making a dumb mistake somewhere?

[fehiepsi]

Could you put b1 inside the nonlin to see if it helps?

[milongo]

That seems to have worked... Interesting. I'm not used to seeing the biases inside non-linearities. Is there a reason for this?

[fehiepsi]

Usually, biases go inside non-linearities. I haven't seen biases be outside in practice. If it is outside, then bias + tanh(.) is a different activation function, which should deserve a different name.

[milongo]

Dang, you are right. I got confused there! Thanks a lot for your help again!