Intro
This blog post is based on the third part of Andrej Karpathy’s makemore series. For part 2 have a look at my previous blog post.
We are building on Andrej’s name prediction model makemore that we started to develop in part 1.
Intro
Before we move on to understand more complex and larger networks such as Gated Recurrent Units (GRUs) or Recurrent Neural Networks (RNNs) it is important to have a good understanding of the activations in a neural net and the gradients flowing backwards and how they behave. So we are stuck a little longer with MLPs and gradient descent based techniques.
Recap part 2
The starting point for the lecture is the same as before. We import the libraries, read the words from the file containing the names and build our character to integer translation dictionairies.
import torchimport torch.nn.functional as Fimport matplotlib.pyplot as plt%matplotlib inline
words = open("/home/arend/JUPYTERENV/notebooks/makemore/names.txt", 'r').read().splitlines()words[:8]
# build the vocabulary of characters and mappings to/from integerschars = sorted(list(set(''.join(words))))stoi = {s:i+1 for i, s in enumerate(chars)}stoi['.'] = 0itos = {i:s for s, i in stoi.items()}vocab_size = len(itos)
We build the train, dev and test split
block_size = 3# Build the datasetdef build_dataset(words): X, Y = [], [] for w in words: context = [0] * block_size for ch in w + '.': ix = stoi[ch] X.append(context) Y.append(ix) context = context[1:] + [ix] X = torch.tensor(X) Y = torch.tensor(Y) print(X.shape, Y.shape) return X, Y#X, Y = build_dataset(words)import randomrandom.seed(42)random.shuffle(words)n1 = int(0.8*len(words))n2 = int(0.9*len(words))Xtr, Ytr = build_dataset(words[:n1])Xdev, Ydev = build_dataset(words[n1:n2])Xte, Yte = build_dataset(words[n2:])
and we initialize the neural net. Note that we parameterized the number of hidden layers and the number of embeddings to get away from ‘hard-coded’ numbers.
# MLP revisitedn_embd = 10 # the dimensionality of the character embedding vectorsn_hidden = 200 # the number of neurons in the hidden layer of the MLPg = torch.Generator().manual_seed(2147483647)C = torch.randn((vocab_size, n_embd), generator = g)W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) b1 = torch.randn(n_hidden, generator = g) W2 = torch.randn((n_hidden, vocab_size), generator = g) # through the multiplication with 0.1 we are 'closer to zero' for the output of the hidden layer# which means smaller logits which means a more even distribution in the probability functionb2 = torch.randn(vocab_size, generator = g) bngain = torch.ones((1, n_hidden))bnbias = torch.zeros((1, n_hidden))bnmean_running = torch.zeros((1,n_hidden))bnstd_running = torch.ones((1,n_hidden))parameters = [C, W1, b1, W2, b2]print(sum(p.nelement() for p in parameters))for p in parameters: p.requires_grad = True
And this is the ‘forward/backward’ pass gradient descent optimization
# same optimization as last timemax_steps = 200000batch_size = 32lossi = []for i in range(max_steps): # minibatch construct ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g) # generates numbers between 0 and # of rows of X-row index Xb, Yb = Xtr[ix], Ytr[ix] # batch X, Y # forward pass emb = C[Xb] # embed the characters in the vector embcat = emb.view(emb.shape[0], -1) # concatenate the embeddings # Linear layer hpreact = embcat @ W1 + b1 # hidden layer pre-activation - don't need bias if we apply batch normalization # Non-linearity h = torch.tanh(hpreact) logits = h @ W2 + b2 # output layer loss = F.cross_entropy(logits, Yb) # Only consider the correct characters of the selected rows Yb #print(loss.item()) #backward pass for p in parameters: p.grad = None loss.backward() # update lr = 0.1 if i < 100000 else 0.01 # step learning rate decay for p in parameters: p.data += -lr * p.grad # track stats if i % 10000 == 0: # print every once in a while print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}') lossi.append(loss.log10().item())
We calculate the loss over the training and dev data set. The Python decorator @torch.no_grad() indicates that whatever happens in the function split_loss does not require a gradient. In Python, a decorator is a function that takes another function as input and extends its behavior without modifying the original function’s code. It essentially wraps the original function with additional functionality. Not having to maintain a graph of gradients removes unnecessary computational overhead. Also, note the interesting Python construct in loss calculation … we create a dictionary and index into it in a single assignment. Key into the dictionary is the type of dataset as a string (‘train’, ‘val’, ‘test’), value is the tuple (input set, predicted output) for each dataset in the split.
torch.no_grad() # decorator disables gradient trackingdef split_loss(split): x, y = { 'train': (Xtr, Ytr), 'val': (Xdev, Ydev), 'test': (Xte, Yte), }[split] # Dictionary with a fixed key ('train', 'val', 'test') and the training/dev/test as tuple value emb = C[x] # (N, block_size, n_embd) embcat = emb.view(emb.shape[0], -1) h = torch.tanh(embcat @ W1 + b1) logits = h @ W2 + b2 loss = F.cross_entropy(logits, y) print(split, loss.item()) split_loss('train')split_loss('val')
The sampling of the model was described in detail at the end of the previous post (inference). Results are not perfect but better than the bigram model. This is our starting point.
Initializing with smaller values
One of the problems here is that we are starting with a loss that is way too high which means that the network is improperly configured at initialization.
In the first iteration we have a loss >27 that then comes down quickly. When training NNs you typically have an idea of what loss to expect at initialization depending on loss function and problem setup. Initially we should have an even probability distribution that assigns the same probability to all 27 characters. So the probability for any character would be 1/27 which means an average negative log likelihood of 3.2958 (ln(1/27)). If we look at a small example we see why. We initialize the logits randomly with a normal distribution that we multiply by factor 10 generating relatively large logits.
If we calculate the probability (we manually calculate the softmax here calculating the exponent of the logit with index 2 by the sum of the exponents over all indices) we see that the probability for index two is extremely small (6.52e-12) because the corresponding logit was initialized to a very small number (-17.4797). If we take the negative log of this small number we end up with a very high loss. Andrej calls this ‘confidently wrong’. At initialization time we want the logits to be closer to zero i.e. the probabilities are closer to an even distribution. We multiply the W2 weights by 0.01.
# through the multiplication with 0.1 we are 'closer to zero' for the output of the hidden layer# which means smaller logits which means a more even distribution in the probability functionW2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01
We can initialize b2 to zero as we don’t want to add a bias of random numbers. After doing that we see that we start with a smaller loss
Note that we don’t want to set weights to zero. It is better to use small numbers to introduce entropy to ‘break symmetry’. So if we rerun this now with the smaller weights for W2 the loss looks like this
We don’t see the hockey stick any longer and we end up with a lower overall loss as we are now spending more cycles optimizing rather than optimizing weights.
Improving the activation
The other problem here is the activation i.e. the output of the hidden linear layer that is the input into the non-linearity.
hpreact = embcat @ W1 + b1 # hidden layer pre-activationh = torch.tanh(hpreact)
If we look at h after the first iteration we see a lot of values close to 1.
This makes sense as we have a wide range of inputs for the tanh. All larger values (>2.5 / <-2.5) generate an output close to 1 or -1.
We can visualize which values in the matrix h (<size batch size> x <hidden layer size>) are close to 1.
We have too much white in our example as the preactivation is too saturated, i.e. the spread of the values in the distribution is too high. We have too many input values for tanh that are so far away from 0 that the gradient of the tanh of that input is 0. With a 0 gradient backpropagation does not work as it relies on a small change in the weight having an impact on the loss. If there is no impact we cannot converge. A completely white column in this matrix would mean that the entire neuron in W1 that matches that column does not contribute to the network’s learning. Adjusting the weights will have no impact as the gradient for the output is 0. This is called a ‘dead neuron’. This observation also holds true for other activation functions that have flat parts e.g. sigmoid, relu. Saturation can happen at initialization and at optimization and you can create dead neurons in the optimization process if you have a steep gradient. To reduce the saturation want lower absolute input values. We multiply W1 by 0.2 and the bias by 0.01. Andrej says that going by his experience adding a small amount of entropy through the bias helps with the training process. Initializing with these values we see a narrower input distribution and less saturation. The input now looks like this …
Instead of a -15 to +15 range for the input we are now at -3 to 3. The output of tanh now looks like this
and we have hardly any saturation
If we run over all full iterations we get this loss
This exercise demonstrates the impact of initialization. As our model is very simple it is forgiving and converges even though the initialization was initially very bad. More complex models are less forgiving and proper initialization is more important. The other thing is that in our example we came up with magic numbers. What we want is a more systematic way to determine these.
A systematic approach to initializing correctly
To get a sense of how we can avoid the expansion that leads to saturation of the non-linearity let’s feed normally distributed input values (x) into a randomly initialized matrix w whose values are also normally distributed. Then we look at how the width of the input distribution (measured by the standard deviation) changes for the dot product of input and matrix.
Matrix x has 1000 normally distributed rows with 10 elements each. w is a 10 x 200 matrix. We concatenate all rows of input and output through ‘view(-1)’ and draw a histogram.
# 1000 random input vectors size 10 - 200 results for each vector - result 1000 x 200x = torch.randn(1000, 10)w = torch.randn(10, 200) y = x @ wprint(x.mean(), x.std())print(y.mean(), y.std())plt.figure(figsize=(20,5))plt.subplot(121)plt.hist(x.view(-1).tolist(), 50, density=True)plt.subplot(122)plt.hist(y.view(-1).tolist(), 50, density=True)
With the means and standard deviations for input and output
We see that the resulting output y has a a wider normal distribution than the input x. We don’t want that in a neural net. We want similar activations roughly unit gaussian (mean μ = 0 and standard deviation σ = 1). Multiplying w with a constant impacts the width of the output distribution. What is the value that preserves the gaussian unit standard deviation through the linear layer?
If you work through the variance of the multiplication it turns out that we can preserve the normal distribution by dividing all values in w by the square root of the number of rows/inputs (in our example 10), also called the fan in. So if we initialize with
w = torch.randn(10, 200) / 10**0.5
we see that the normal distribution is preserved
So far so good. We now know how to preserve the unit Gaussian input. But how can we initialize the weights so that the activations maintain reasonable values throughout a network with linear layers and non-linearities?
This paper from Kaiming He et al. investigates this. Kaiming He et al. study Convolutional Neural Networks (CNNs) with relu and prelu non-linearities but the analysis is similar for tanh. The tanh or relu non-linearities following the linear operation squash/compress the distribution. In order to counter this squashing/narrowing effect we need to boost the input to the non-linearities through a constant. Relu clamps half of the distribution so the authors suggest 2/sqrt(fan in). For a tanh nonlinearity the paper suggests that the ideal multiplier is a gain of 5/3 divided by the square root of the fan-in. If you want to dig deeper and understand where 5/3 is coming from, I wrote up my derivation here:
The paper also studies the backpropagation. Kaiming He et al. found that proper initialization of the forward pass also properly initializes the backward pass. Pytorch offers these initialization methods.
In our example we initialize manually with
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * 5/3 / ((n_embd * block_size) ** 0.5)
And if we train the neural net we end up in roughly the same place. But we got there with a more systematic approch that did not use any magic numbers we needed to find through experimentation.
Up to about 9 years ago you had to be very careful with your activations and gradients. Everything was very fragile. Nowadays it has become less important to 100% correctly initialize these networks as there are residual corrections, normalization layers and other optimizers. One of these modernizations is batch normalization introduced by a team at Google in 2015.
Batch normalization
As mentioned above, at initialization time we want hpreact (the input into the non-linearity) to be roughly unit Gaussian (normal distribution with a mean of 0 and a standard deviation of 1). We don’t want satuation of the tanh output and at the same time we want a range of input values for the tanh that is not too narrow. So the idea suggested in this paper
is that at initialization time we normalize/standardize (we consider these as interchangeable here) the output of the hidden layer before activation (applying the non-linearity) over all elements of the batch. Normalization means that for each value of the matrix before activation we subtract the mean over all columns and divide that value by the standard deviation. This standardization is equivalent to calculating the z-score. The epsilon in the below formula for the normalization was introduced to prevent a division by zero in case the variance is zero.
For the normalization we need the mean and the standard deviation, which is also the root of the variance. We can get the mean over all batches (collapsing the rows – dimension 0) with
hpreact.mean(0, keepdim=True)
and the standard deviation with
hpreact.std(0, keepdim=True)
So we can normalize the preactivation layer with
hpreact = (hpreact - hpreact.mean(0, keepdim=True)) / hpreact.std(0, keepdim=True)
All the operations are differentiable, which allows us to calculate the backward pass.
But we only want to force this normalization at initialization time. We don’t want to prevent the network from learning. Blindly standardizing is too restrictive. We want the mean and std deviation to be able to move as the network learns. So we introduce a scaling factor. We take the input values and multiply by a gain and offset by a value (‘scale and shift’) for the final output of the layer.
So we introduce a multiplier bngain (the equivalent of γ in the formula) and a bias (equivalent to β). These parameters will also be trained through the backpropagation and allow the Gaussian distribution to widen/narrow and move. These values apply to the distribution over each batch and are of the size of the hidden layer.
As we want to force the unit Gaussian distribution at initialization time we make γ 1 and β 0.
bngain = torch.ones((1, n_hidden))bnbias = torch.zeros((1, n_hidden))
and add it to the calculation for the hidden layer
hpreact = bngain * (hpreact - hpreact.mean(0, keepdim=True)) / hpreact.std(0, keepdim=True) + bnbias
To make sure that gain and bias are optimized as the network learns we need to include them in the list of parameters that we adjust with each iteration.
parameters = [C, W1, W2, b2, bngain, bnbias]
And we also need to adjust the validation to include gain and bias.
The results are comparable to the results we achieved previously. Reason is that the batch normalization ‘does not do much’ as we already initialized with the Kaiming init and the neural network is relatively simple with only one hidden layer. Once you have a deeper neural net it will be hard to maintain ‘unit Gaussian’ through the neural network. ‘Sprinkling’ batch normalization through a complex network is much easier than trying to figure out parameters to maintain unit Gaussian.
One cost of batch normalization is that it starts coupling the formerly independent samples through the mean and variance calculated across all samples in each batch. hpreact (and therefore h) changes slightly based on the other samples within that batch. This slightly couples the samples. So you get a jitter for h and logits. In a strange way this is a good thing … you can think of it as a ‘regularizer’ … noise that prevents overfitting and regularizing the neural net.
If we train the model based on batches of coupled samples, how can we use the optimized model to predict based on a single context? How do we feed a single example and expect sensible results?
For inference we need to go through the same steps as in the forward pass. hpreact.mean and hpreact.std are calculated for a batch. For inference we cannot just apply the values for the latest batch, we need to apply the mean and std for the entire data set. One solution is to calculate mean and std for the whole model which introduces an extra step e.g.
with torch.no_grad(): emb = C[Xtr] embcat = emb.view(emb.shape[0], -1) hpreact = embcat @ W1 bnmean = hpreact.mean(0, keepdim=True) bnstd = hpreact.std(0, keepdim=True)
Another option that avoids this extra step is to keep track of the running mean and std for each batch to avoid that extra step, in other words, we calculate an averaged mean and std over all means for each batch/iteration. We initialize the mean to 0 and the std to 1 – as we Kaiming init the initial mean/std will be roughly 0/1. We initialize the running values
bnmean_running = torch.zeros((1,n_hidden))bnstd_running = torch.ones((1,n_hidden))
and calculate them in each iteration, updating the current running mean by a small factor to limit the impact of outliers for a batch
Note that mean and running are not trained through gradient descent. They are just updated on an ongoing basis. When we compare with the explicit calculation we see that the two values are very close
and the loss is identical.
with the calculation of mean and std over the full dataset, and
with the calculation averaging over the running batches.
Note that with this method bias b1 is ‘wasteful’ as we are adding it, then calculate the mean and then subtract the mean. So whatever bias we add is subtracted later. When we use batch normalization there is no point in adding a bias in the previous layer.
We use batch normalization to control the statistics of activation in a neural net. It is typically applied after the layer that has multiplications, for example a linear or a convolutional (to be explained) layer.
An example of that is Resnet that follows the same pattern. Convolutional Layer – Batch Normalization – Non-Linearity (typical RelU). Convolutional networks calculate Wx+b for regions in a matrix equivalent to sections of an image. It follows the same pattern: 1. Convolution (linear layer) 2. Normalization 3. Non-linearity.
Pytorch has functions for the linear layer that work ‘out of the box’.
By default the factor applied at initialization is one over square root of fan in (without the 5/3 for the tanh). It does initialize with a linear distribution though, not with a Gaussian distribution.
You can also (re-)initialize with a normal distribution
There is also a batch normalization layer in Pytorch
num_features is needed to initialize the gain, bias and the running buffers in our example that would be 200 i.e. the number of hidden layers. Epsilon, as discussed, prevents the division by zero and can typically stay at the default value. The momentum are the constants for calculating the running means and std. They are using .1, we are using .001. Smaller batch sizes should use a smaller momentum as you want slower convergence to avoid ‘thrashing’ of the mean and std. affine = True means that we also have gain and bias as learnable parameters. Should always be true. track_running_stats determines if we should keep track of the running mean and std. You may want to skip running_stats if you want to calculate them at the end over the whole matrix so that the batch normalization does not waste effort keeping track of these values.
Summary
We need to understand the importance of gradients and activations in neural networks … which becomes increasingly important as the size and complexity of networks grows.
- You can end up with hockey stick losses, meaning the network does wasteful work, if you initialize with the wrong values. We multiplied the output layer W2 by 0.01 to be closer to 0 which means closer to an even distribution.
- We want roughly gaussian activations through the network – we saw how to initialize the neural net parameters to preserve a unit Gaussian distribution through the Kaiming initialization of W1. The idea is to compensate for the ‘squashing’ the non-linearity introduces by increasing the values before the expansion through the non-linearity. For deeper networks it is hard to set weights and biases to ensure a uniform distribution.
- We introduced batch normalization, one of many normalization layer types – the idea is that if we want a unit Gaussian activation, we need to center and normalize our data at each layer to avoid saturation through the non-linearity or just going through a small range entering the non-linearity. This is possible as the ‘centering’ operation is differentiable (subtract mean and divide by std). But now we need gain and bias as trainable parameters. For inference and to calculate the loss we also need to calculate mean and std over the whole layer – to avoid the extra step we approximate these two values as we iterate over the minibatches – this layer causes a lot of bugs as it couples the examples in a batch (“I shot myself in the foot over and over again“). Other types of normalization as layer normalization are now more popular. All of this is very important in recurring neural nets.
Making our code ‘Pytorch like‘
We structured the code into modules so that we can build neural networks like we would build them in Pytorch.
First we are creating a linear layer that takes the number of input (fan in) and output layers and whether we want to use a bias. We are initializing the weights dividing by the square root of the fan in as we discussed above. __call__() calculates the dot product of the matrix and the vector x that we pass. parameters returns the weights and biases.
class Linear: def __init__(self, fan_in, fan_out, bias=True): self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in**0.5 self.bias = torch.zeros(fan_out) if bias else None def __call__(self, x): self.out = x @ self.weight if self.bias is not None: self.out += self.bias return self.out def parameters(self): return [self.weight] + ([] if self.bias is None else [self.bias])
As in Pytorch we call the batch norm layer BatchNorm1d.
class BatchNorm1d: # Similar to Python batch1d layer def __init__(self, dim, eps=1e-5, momentum=0.1): self.eps = eps # to prevent division by zero self.momentum = momentum self.training = True # Behavior is different for training vs evaluation mode - # parameters trained with backprop self.gamma = torch.ones(dim) self.beta = torch.zeros(dim) # buffers - updated for each iteration self.running_mean = torch.zeros(dim) self.running_var = torch.ones(dim) def __call__(self, x): if self.training: xmean = x.mean(0, keepdim=True) # batch mean xvar = x.var(0, keepdim=True, unbiased=True) # batch variance else: xmean = self.running_mean xvar = self.running_var xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize unit variance self.out = self.gamma * xhat + self.beta # update the buffers if self.training: with torch.no_grad(): self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar return self.out def parameters(self): return [self.gamma, self.beta]
We implement the same parameters with the exception of device and and data type. We are implementing a switch indicating whether we are in training or inference mode.
self.training = True
When we are training we are calculating mean and std over each batch, when we are in inference mode we want to use the running mean and std as we don’t have a batch to calculate over
if self.training:
xmean = x.mean(0, keepdim=True) # batch mean
xvar = x.var(0, keepdim=True, unbiased=True) # batch variance
else:
xmean = self.running_mean
xvar = self.running_var
And we need to update the running mean and gradient (also called buffers) in training. As these parameters are tracked without the need for any optimization we can set no_grad to avoid unnecessary computational overhead.
if self.training:
with torch.no_grad():
self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
In the spirit of the paper we are calculating the variance and as we standardize we divide by the square root of the variance which is the same as the standard deviation
xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize unit variance
The tanh is pretty straightforward and has the same parameter as the pytorch equivalent
class Tanh: def __call__(self, x): self.out = torch.tanh(x) return self.out def parameters(self): return []
Now that we have layers we can stack them up into a list.
n_embd = 10n_hidden = 100g = torch.Generator().manual_seed(2147483647)C = torch.randn((vocab_size, n_embd), generator=g) # Initialize embeddinglayers = [ # Create a list of layers - again, a list can comprise objects of different types Linear(n_embd * block_size, n_hidden), Tanh(), Linear(n_hidden, n_hidden), Tanh(), Linear(n_hidden, n_hidden), Tanh(), Linear(n_hidden, n_hidden), Tanh(), Linear(n_hidden, n_hidden), Tanh(), Linear(n_hidden, vocab_size), ]with torch.no_grad(): # last layer, make less confident - remember that smaller values in the last layer # mean that we are closer to an even distribution layers[-1].weight *= 0.1 for layer in layers[:-1]: if isinstance(layer, Linear): layer.weight *= 5/3 # 5/3 # We have so far initialized with 1/sqrt(fan_in) - now we are adding the gain parameters = [C] + [p for layer in layers for p in layer.parameters()]print(sum(p.nelement() for p in parameters)) # number of parameters in totalfor p in parameters: p.requires_grad = True
We have six layers
layers = [ Linear(n_embd * block_size, n_hidden), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, vocab_size), BatchNorm1d(vocab_size),]
As we discussed at the very beginning we reduce the value of the weights in the last layer to come closer to an even distribution and avoid the initial hockey stick. For all layers other than the last one we initialize with the Kaiming initialization. We have a 5/3 gain and the weights are already divided by the square root of the fan in when they are intialized.
with torch.no_grad(): # last layer, make less confident - remember that smaller values in the last layer # mean that we are closer to an even distribution #layers[-1].weight *= 0.1 layers[-1].gamma *= 0.1 # all other layers: apply again for layer in layers[:-1]: if isinstance(layer, Linear): layer.weight *= 5/3 # 5/3 # We have so far initialized with 1/sqrt(fan_in) - now we are adding the gain
The list of parameters to optimize are the embeddings and the weights and biases (and gain and biases for the batch norm layers) over all layers. Note the double comprehension used in the list. Finally we tell Pytorch that all parameters require a gradient.
parameters = [C] + [p for layer in layers for p in layer.parameters()]print(sum(p.nelement() for p in parameters)) # number of parameters in totalfor p in parameters: p.requires_grad = True
Now we run the optimization as we are used to.
The difference to the previous optimization is that rather than explicitly calculating the output of the layers we iterate over the list of layers and the calculations are performed by the __call__ functions for each layer. It helps to look at the parameters for each layer.
The first parameter is the embedding followed by W1, followed by b1 etc. So at a high level we enter the input ‘triple’ letters for each mini-batch (in our example size 32), we embed and concatenate so that the input is in our example 32 x 30 (10 embedding values for each triple). Then we pass this through the layers and calculate the cross-entropy for each mini-batch.
Note that for the backward pass we need to tell pytorch explicitly to retain the gradient
Activation statistics – Diagnostic tools
Now we uncomment the break in the optimization and look at the outputs from each tanh layer after the first iteration. We want to get a sense of the distribution of tanh output as we go through the layers. We calculate mean, std and saturation (# of elements with |tanh|> .97) over all the outputs in each layer.
plt.figure(figsize=(20, 4))legends = []for i, layer in enumerate(layers[:-1]): # excluding the output layer if isinstance(layer, Tanh): t = layer.out print(t.shape) print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100)) hy, hx = torch.histogram(t, bins=100, density=False) plt.plot(hx[:-1].detach(), hy.detach()) legends.append(f'layer {i} ({layer.__class__.__name__}')plt.legend(legends);plt.title('activation distribution')
Note that there is no dimension specified for mean and standard deviation and we calculate over all values of the 32×100 matrix.
For the initial pass we are always getting a high saturation. This is because the input values are a unit gaussian distribution that takes relatively large values ‘on the wings’. After the first pass we always have the tanh that ‘squashes’ the values and limits the range of input values for the next layer. As we have the gain set to 5/3 the saturation goes down and stabilizes.
When we look at the gradient we see a reasonable gradient distribution. What we want is roughly the same distribution over all layers.
With an initialization factor that is too small e.g. 1
we don’t compensate for the tanh squashing and end up with a high density of values around 0 and zero saturation
If the value is to high, in our example 2, we end up with very high saturation
Side note: Why do we not get more values around 0 as the distribution of the input peaks around 0? The reason is that the bins for the output layer are over the output range of tanh. An output bin around 0 (in blue in the image) maps to a much narrower input bin than an output bin closer to 1 (in red) where the tanh is close to flat. The larger input interval for the higher output intervals more than compensates for the larger number of values around zero in the preactivation layer.
So without the use of batch normalization we have to be careful with the initialization.
Andrej also points out that we need the non-linear layer as otherwise the entire network just behaves like a linear layer
Now we are adding another plot that is important to look at: The gradients for the weights of the matrices in the linear layers and the scale of the gradient relative to the value of the matrix element. The dimension check ensures that we are only capturing ‘matrix’ weights and ignore the batch layer parameters and anything else.
# visualize histogramsplt.figure(figsize=(20, 4)) # width and height of the plotlegends = []for i,p in enumerate(parameters): t = p.grad if p.ndim == 2: print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std())) hy, hx = torch.histogram(t, density=True) plt.plot(hx[:-1].detach(), hy.detach()) legends.append(f'{i} {tuple(p.shape)}')plt.legend(legends)plt.title('weights gradient distribution');
The gradient to data ratio gives a sense of the scale of the gradient relative to the scale of the values. If the gradient is too large compared to the data things get unstable. In our case (other than for the output layer) the gradient is about 1000 times smaller than the respective weight, which is good. The last layer has a much higher standard deviation, about 10 times higher than the values before. This means that we would train the last layer 10 times faster. If we iterate 1,000 times things stabilize and the tails of the outer layer come in.
What is more important than the gradient to data ratio is the update to gradient ratio. We are adding this to keep track of the update to data ratio.
Andrej’s guideline is that the updates should be around 1/1000s of the value (the line in the graph plotted at -3 as a reference). The reason for the large update for the last layer is that we reduced the values through multiplication of the weights with 0.1 to be closer to an even distribution as discussed in the beginning.
If the learning rate is low the updates are too small and we train too slow. For example, if we change the learning rate from 0.1 to 0.001 the plot looks like this and all the updates are way too small.
Now let’s bring back batch normalization. The batch normalization is placed between the linear layer and the non-linearity. You do get similar results if you place the batch norm layer after the non-linearity. You can also place it at the end before the loss function. To make the softmax of the last layer less confident we are changing the gamma to make the values smaller and avoid the learning ‘hockeystick’. As we are now normalizing before every tanh layer the activation looks good and we don’t need to ‘balance the pencil’ any longer.
Now we are less brittle with respect to the initialization values. So if we change the gain to different factors e.g. .2 instead of 5/3 and the activations will remain unaffected. The learning rate is still affected by the linear weights, as the learning process now adjusts not just the weights but also the gain and bias for the normal distribution. It may be necessary to retune the learning rate.
Because we built our classes in a way that is Pytorch compatible you can just import pytorch’s nn and prepend ‘nn.’ to each class for the corresponding pytorch implementation.
Another thing to keep in mind is that when using the batch initializtion we also need to standardize during inference
Leave a Reply