Weight Initialization

In this lesson, you'll learn how to find good initial weights for a neural network. Weight initialization happens once, when a model is created and before it trains. Having good initial weights can place the neural network close to the optimal solution. This allows the neural network to come to the best solution quicker.

To see how different weights perform, we'll test on the same dataset and neural network. That way, we know that any changes in model behavior are due to the weights and not any changing data or model structure. We'll instantiate at least two of the same models, with different initial weights and see how the training loss decreases over time.

Sometimes the differences in training loss, over time, will be large and other times, certain weights offer only small improvements.

We'll train an MLP to classify images from the Fashion-MNIST database to demonstrate the effect of different initial weights. As a reminder, the FashionMNIST dataset contains images of clothing types; classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']. The images are normalized so that their pixel values are in a range [0.0 - 1.0). Run the cell below to download and load the dataset.

get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")
seaborn.set(style="whitegrid",
            rc={"axes.grid": False,
                "font.family": ["sans-serif"],
                "font.sans-serif": ["Open Sans", "Latin Modern Sans", "Lato"],
                "figure.figsize": (10, 8)},
            font_scale=1)

Obtain one batch of training images.

dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy()

Plot the images in the batch, along with the corresponding labels.

fig = pyplot.figure(figsize=(12, 10))
fig.suptitle("Sample FASHION Images", weight="bold")
for idx in np.arange(20):
    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
    ax.imshow(np.squeeze(images[idx]), cmap='gray')
    ax.set_title(classes[labels[idx]])

We've defined the MLP that we'll use for classifying the dataset.

Let's start looking at some initial weights.

Below, we are using helpers.compare_init_weights to compare the training and validation loss for the two models we defined above, model_0 and model_1. This function takes in a list of models (each with different initial weights), the name of the plot to produce, and the training and validation dataset loaders. For each given model, it will plot the training loss for the first 100 batches and print out the validation accuracy after 2 training epochs. Note: if you've used a small batch_size, you may want to increase the number of epochs here to better compare how models behave after seeing a few hundred images.

We plot the loss over the first 100 batches to better judge which model weights performed better at the start of training. I recommend that you take a look at the code in helpers.py to look at the details behind how the models are trained, validated, and compared.

Run the cell below to see the difference between weights of all zeros against all ones.

Initialize two NN's with 0 and 1 constant weights.

model_0 = Net(constant_weight=0)
model_1 = Net(constant_weight=1)

Put them in list form to compare.

model_list = [(model_0, 'All Zeros'),
              (model_1, 'All Ones')]

ModelLabel = Tuple[nn.Module, str]
ModelLabels = Collection[ModelLabel]

def plot_models(title:str, models_labels:ModelLabels):
    """Plots the models

    Args:
     title: the title for the plots
     models_labels: collections of model, plot-label tuples
    """
    figure, axe = pyplot.subplots()
    figure.suptitle(title, weight="bold")    
    axe.set_xlabel("Batches")
    axe.set_ylabel("Loss")

    for model, label in models_labels:
        loss, validation_accuracy = helpers._get_loss_acc(model, train_loader, valid_loader)
        axe.plot(loss[:100], label=label)
    legend = axe.legend()
    return

Plot the loss over the first 100 batches.

plot_models("All Zeros vs All Ones",
            ((model_0, "All Zeros"),
             (model_1, "All ones")))

After 2 Epochs:
Validation Accuracy
    9.475% -- All Zeros
   10.175% -- All Ones
Training Loss
    2.304  -- All Zeros
  1914.703  -- All Ones

As you can see the accuracy is close to guessing for both zeros and ones, around 10%.

The neural network is having a hard time determining which weights need to be changed, since the neurons have the same output for each layer. To avoid neurons with the same output, let's use unique weights. We can also randomly select these weights to avoid being stuck in a local minimum for each run.

A good solution for getting these random weights is to sample from a uniform distribution.

The general rule for setting the weights in a neural network is to set them to be close to zero without being too small. A good practice is to start your weights in the range of \([-y, y]\) where \(y=1/\sqrt{n}\) (\(n\) is the number of inputs to a given neuron).

Let's see if this holds true; let's create a baseline to compare with and center our uniform range over zero by shifting it over by 0.5. This will give us the range [-0.5, 0.5).

weights_init_uniform_center = partial(weights_init_uniform, -0.5, 0.5)

Unlike the uniform distribution, the normal distribution has a higher likelihood of picking number close to it's mean. To visualize it, let's plot values from NumPy's np.random.normal function to a histogram.

np.random.normal(loc=0.0, scale=1.0, size=None)

Outputs random values from a normal distribution.

• loc: The mean of the normal distribution.
• scale: The standard deviation of the normal distribution.
• shape: The shape of the output array.

figure, axe = pyplot.subplots()
figure.suptitle("Standard Normal Distribution", weight="bold")
grid = seaborn.distplot(numpy.random.normal(size=[1000]))

Let's compare the normal distribution against the previous, rule-based, uniform distribution.

The normal distribution should have a mean of 0 and a standard deviation of \(y=1/\sqrt{n}\)

def weights_init_normal(m: nn.Module) -> None:
    '''Takes in a module and initializes all linear layers with weight
       values taken from a normal distribution.'''

    classname = m.__class__.__name__
    if classname.startswith("Linear"):    
        m.weight.data.normal_(mean=0, std=1/numpy.sqrt(m.in_features))
        m.bias.data.fill_(0)
    return

create a new model with the rule-based, uniform weights

model_uniform_rule = Net()
model_uniform_rule.apply(weights_init_uniform_rule)

create a new model with the rule-based, NORMAL weights

model_normal_rule = Net()
model_normal_rule.apply(weights_init_normal)

compare the two models

plot_models('Uniform vs Normal',
            ((model_uniform_rule, 'Uniform Rule [-y, y)'), 
             (model_normal_rule, 'Normal Distribution')))

The normal distribution gives us pretty similar behavior compared to the uniform distribution, in this case. This is likely because our network is so small; a larger neural network will pick more weight values from each of these distributions, magnifying the effect of both initialization styles. In general, a normal distribution will result in better performance for a model.

Let's quickly take a look at what happens without any explicit weight initialization.

model_normal_rule = Net()
model_normal_rule.apply(weights_init_normal)
model_default = Net()
model_rule = Net()
model_rule.apply(weights_init_uniform_rule)

plot_models("Default vs Normal vs General Rule", (
    (model_default, "Default"),
    (model_normal_rule, "Normal"),
    (model_rule, "General Rule")))

They all sort of look the same at this point.