Goals
Rectifying poor network initialisation
Good initialisation of the network parameters (weights and biases) is crucial for better learning. This is especially true for deeper or more complex networks (e.g. 50 layers), which are especially unforgiving to poor initialisations.
Solution: scale parameter values. First manually via hard coding, later more systematically (Kaiming initialisation)
Batch normalisation (following notebooks)
- Understanding neuron activations and backward gradient flow behaviour during training.
Why this is important
- As an example: RNNs are very expressive (universal approximator, so in principle, can implement all algorithms).
- But they’re very difficult to optimise with commonly used first-order gradient based optimisation techniques.
This chapter covers:
- The behaviour of activations and gradients during training explain why training is difficult
- Also why NN variants since RNNs have taken different approaches
Rebuild the MLP from karpathy-03-mlp
# imports, build vocabulary, build_dataset function, create train/val/test data split.
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt # for making figures
%matplotlib inline
# import data (32,033 words, 228,146 training examples)
words = open('data/names.txt', 'r').read().splitlines()
# build the vocabulary of characters, and mappings to/from integers
chars = sorted(list(set(''.join(words))))
stoi = {s:i+1 for i,s in enumerate(chars)}
stoi['.'] = 0
itos = {i:s for s,i in stoi.items()}
vocab_size = len(itos)
# fn: build dataset (training examples X, and labels Y) for an INPUT list of names only
block_size = 3 # context length: how many characters do we take to predict the next one?
def build_dataset(words):
X, Y = [], [] # X: NN input training examples, Y: labels for each input in X
for w in words:
#print(w)
context = [0] * block_size
for ch in w + '.':
ix = stoi[ch]
X.append(context)
Y.append(ix)
#print(''.join(itos[i] for i in context), '--->', itos[ix])
context = context[1:] + [ix] # crop and append
X = torch.tensor(X)
Y = torch.tensor(Y)
print(X.shape, Y.shape)
return X, Y
# randomly shuffle words data set, and create train, val, test splits
import random
random.seed(42)
random.shuffle(words)
n1 = int(0.8*len(words)) # to index 80th percentile word (i.e. words[0] to words[n1])
n2 = int(0.9*len(words)) # to index the 90th percentile word (i.e. words[n1] to words[n2])
Xtr, Ytr = build_dataset(words[:n1]) # 80% test set (Xtr: training examples, Ytr: training labels)
Xdev, Ydev = build_dataset(words[n1:n2]) # 10% validation set
Xte, Yte = build_dataset(words[n2:]) # 10% test settorch.Size([182625, 3]) torch.Size([182625])
torch.Size([22655, 3]) torch.Size([22655])
torch.Size([22866, 3]) torch.Size([22866])See 01_build_mlp if definitions / dimensionality below is confusing
# MLP revisited (network parameters no longer hardcoded)
n_embd = 10 # the dimensionality of the character embedding vectors
n_hidden = 200 # the number of neurons in the hidden layer of the MLP
g = torch.Generator().manual_seed(2147483647) # for reproducibility
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)
b2 = torch.randn(vocab_size, generator=g)
parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters)) # number of parameters in total
for p in parameters:
p.requires_grad = True11897# o - same optimization as last time (200,000 iters; batch size 32 examples (per iter.))
max_steps = 200000
batch_size = 32
lossi = []
print('note high first iter loss')
for i in range(max_steps):
# minibatch construct
ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)
Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y
# forward pass
emb = C[Xb] # embed the characters into vectors
embcat = emb.view(emb.shape[0], -1) # concatenate the vectors (flattening emb dims)
hpreact = embcat @ W1 + b1 # hidden layer pre-activations
h = torch.tanh(hpreact) # hidden layer activations
logits = h @ W2 + b2 # output layer
loss = F.cross_entropy(logits, Yb) # loss function
# 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())
# break # to view ONLY zero'th iteration lossnote high first iter loss
0/ 200000: 27.8817
10000/ 200000: 2.8138
20000/ 200000: 2.5218
30000/ 200000: 2.7874
40000/ 200000: 2.0334
50000/ 200000: 2.6237
60000/ 200000: 2.3289
70000/ 200000: 2.0826
80000/ 200000: 2.2784
90000/ 200000: 2.2252
100000/ 200000: 2.0428
110000/ 200000: 2.3121
120000/ 200000: 2.0570
130000/ 200000: 2.4546
140000/ 200000: 2.2233
150000/ 200000: 2.1551
160000/ 200000: 2.0597
170000/ 200000: 1.7981
180000/ 200000: 2.0194
190000/ 200000: 1.7459# io - if break statement in cell above on 0th iteration
print('initial logits are quite extreme. not all near 0:\n', logits[0])
# this creates the fake overconfidence, and makes initial loss so high (~27.9)initial logits are quite extreme. not all near 0:
tensor([ 0.0747, 3.5932, 1.7715, 1.0445, 2.1256, 3.4390, 0.5381, -0.1884,
1.8437, 4.4816, 0.2604, 0.8402, 4.1507, 3.5982, 5.6979, 3.4520,
1.9421, 0.3094, 3.8519, 3.4718, 2.5219, 1.6419, 2.1227, 1.3424,
1.0001, 4.3635, 2.2235], grad_fn=<SelectBackward0>)plt.plot(lossi)[<matplotlib.lines.Line2D at 0x114fc2c10>]NB: @torch.no_grad() disables gradient tracking. If we’re only doing a forward pass, no need to waste memory / compute storing gradients (which are only needed in backward pass)
# fn: evaluate loss for specific data split. args: 'train', 'val', 'test'
@torch.no_grad() # this decorator disables gradient tracking
def split_loss(split):
x,y = {
'train': (Xtr, Ytr),
'val': (Xdev, Ydev),
'test': (Xte, Yte),
}[split]
# forward pass (calculate loss)
emb = C[x] # embed chars into vectors (N, block_size, n_embd)
embcat = emb.view(emb.shape[0], -1) # concatenate vectors into (N, block_size * n_embd)
hpreact = embcat @ W1 + b1 # hidden layer pre-activ^ns (N, n_hidden)
h = torch.tanh(hpreact) # hidden layer activations (N, n_hidden)
logits = h @ W2 + b2 # output layer (N, vocab_size)
loss = F.cross_entropy(logits, y) # loss function
print(split, loss.item())
split_loss('train')
split_loss('val')train 2.1267659664154053
val 2.1697638034820557# sample from the model
g = torch.Generator().manual_seed(2147483647 + 10)
for _ in range(20):
out = []
context = [0] * block_size # initialize with all '...'
while True:
# forward pass the neural net
emb = C[torch.tensor([context])] # embed context of that 1 token into a vector (1,block_size,n_embd)
h = torch.tanh(emb.view(1, -1) @ W1 + b1) # concat vectors -> hidden preacts -> hidden activations
logits = h @ W2 + b2 # output layer (log-counts)
probs = F.softmax(logits, dim=1) # exponentiate and normalise
# sample from the distribution
ix = torch.multinomial(probs, num_samples=1, generator=g).item()
# shift the context window and track the samples
context = context[1:] + [ix]
out.append(ix)
# if we sample the special '.' token, break
if ix == 0:
break
print(''.join(itos[i] for i in out)) # decode and print the generated wordcarlah.
amorie.
khirmin.
rey.
cassanden.
jazhubedah.
sart.
kaeli.
nellara.
chaiir.
kaleigh.
ham.
jore.
quint.
salin.
alianni.
wazthoniearyxi.
jace.
pirran.
eddeci.Two bugs in above MLP initialisation
We expect all characters to be initialised with roughly equal softmax probability (prob) of — i.e. a uniform distribution over 27 characters. Therefore, the expected loss (negative-log likelihood) for a single training example where the model assigns uniform prob distribution across all 27 characters is:
An initial loss of ~27.9 is far greater than this, which tells us something is badly miscalibrated at initialisation.
Bug 1: Parameter initialisation causing initial loss ~27.8
Why this is wrong
Expecting loss of ~3.3, getting ~27.9. The network has non-uniform initial logits — specifically, the “wrong” neuron is confidently high.
The network is very confidently wrong. Some characters are very confident, others are very not confident.
Why it’s occurring
When logits are randomly initialised (e.g. via a standard normal), rather than of all logits being equal, some logits will be large in magnitude. See break statement above.
Softmax exponentiates logits, so large values produce a very confident (but wrong) probability distribution. A model that is confidently wrong incurs a very high loss.
Consider a single 4-dimensional (i.e. 4 characters only) softmax example:
# io - Consider a 4-dimensional (i.e. 4 characters only) softmax example:
# logits = torch.tensor([0.0, 0.0, 0.0, 0.0]) # scenario 1
# logits = torch.tensor([0.0, 0.0, 5.0, 0.0]) # scenario 2
# logits = torch.tensor([0.0, 5.0, 0.0, 0.0]) # scenario 3
# logits = torch.tensor([-3.0, 5.0, 0.0, 2.0]) # scenario 4
logits = torch.randn(4) # * 100 # scenario 5
probs = torch.softmax(logits, dim=0)
loss = -probs[2].log()
print('logits:', logits, '\nprobs:', probs, '\nloss:', loss)logits: tensor([ 0.5636, 2.9025, -1.7178, 0.4539])
probs: tensor([0.0809, 0.8384, 0.0083, 0.0725])
loss: tensor(4.7965)We pick target = 2 arbitrarily — in real training it would be the index of the actual correct next character. loss = -log(p[target]) is simply: how much probability did the model assign to the correct answer? (e.g. if the vocabulary is [a, b, c, d] and the correct next character is c, then target = 2).
| Scenario | Logits | p[2] softmax | Loss |
|---|---|---|---|
| 1 | [0, 0, 0, 0] | 0.25 (uniform) | ≈ 1.386 ← expected |
| 2 | [0, 0, 5, 0] | ≈ 0.980 | ≈ 0.020 (lucky) |
| 3 | [0, 5, 0, 0] | ≈ 0.007 | ≈ 5.0 (confidently wrong) |
| 4 | [-3, 5, 0, 2] | ≈ 6.4e-3 | ≈ 5.1 (confidently wrong) |
| 5 | randn(4) | varies | unpredictable — likely >> 1.386 |
- Scenario 1 is the ideal initialisation — all logits equal, softmax is uniform, every class (i.e. character) gets
1/Nprobability. - Scenario 2 — the correct next character happens to have the large logit; loss is near zero. This is just luck, not learning.
- Scenarios 3 and 4 reflect what actually happens with
torch.randnweights: large logits cause softmax to concentrate (i.e. steal) probability mass on one character, and if that character isn’t the target, loss spikes.
Solution
We want logits to be approximately zero at initialisation. More precisely, softmax is invariant to constant shifts:
So what matters is that logits are equal, not that they are literally zero. All-zeros is a natural choice because it avoids introducing an arbitrary positive or negative bias.
Concretely:
- Scale down final layer weights at init (
W2 * 0.01), but do not zero it out — reasons covered later
- MLP init code:
W2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.01- And zero-initialise its bias (
b2 = 0), since it feeds directly into the logits and therefore directly into the loss
- MLP init code:
b2 = torch.randn(vocab_size, generator=g) * 0(or justtorch.zeros(vocab_size))
Outcomes of this change
- logits start near zero
- → uniform softmax
- → loss starts at
-ln(1/N)as expected (~3.296 for 27 characters, instead of a huge loss value) - → training begins learning immediately rather than wasting steps unwinding overconfidence (less hockey stick shape)
# o - scale down W2, zero-initialise b2
n_embd = 10 # the dimensionality of the character embedding vectors
n_hidden = 200 # the number of neurons in the hidden layer of the MLP
g = torch.Generator().manual_seed(2147483647) # for reproducibility
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) * 0.01
b2 = torch.randn(vocab_size, generator=g) * 0
parameters = [C, W1, b1, W2, b2]
# print(sum(p.nelement() for p in parameters)) # number of parameters in total
for p in parameters:
p.requires_grad = True
# rerun same optimization as last time (200,000 iters; batch size 32 examples (per iter.))
max_steps = 200000
batch_size = 32
lossi = []
print('note first iter loss now reasonable')
for i in range(max_steps):
# minibatch construct
ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)
Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y
# forward pass
emb = C[Xb] # embed the characters into vectors
embcat = emb.view(emb.shape[0], -1) # concatenate the vectors (flattening emb dims)
hpreact = embcat @ W1 + b1 # hidden layer pre-activations
h = torch.tanh(hpreact) # hidden layer activations
logits = h @ W2 + b2 # output layer
loss = F.cross_entropy(logits, Yb) # loss function
# 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())
# break # to view ONLY zero'th iteration loss
# fn: evaluate loss for specific data split. args: 'train', 'val', 'test'
@torch.no_grad() # this decorator disables gradient tracking
def split_loss(split):
x,y = {
'train': (Xtr, Ytr),
'val': (Xdev, Ydev),
'test': (Xte, Yte),
}[split]
# forward pass (calculate loss)
emb = C[x] # embed chars into vectors (N, block_size, n_embd)
embcat = emb.view(emb.shape[0], -1) # concatenate vectors into (N, block_size * n_embd)
hpreact = embcat @ W1 + b1 # hidden layer pre-activ^ns (N, n_hidden)
h = torch.tanh(hpreact) # hidden layer activations (N, n_hidden)
logits = h @ W2 + b2 # output layer (N, vocab_size)
loss = F.cross_entropy(logits, y) # loss function
print(split, loss.item())
split_loss('train')
split_loss('val')note first iter loss now reasonable
0/ 200000: 3.3221
10000/ 200000: 2.1900
20000/ 200000: 2.4196
30000/ 200000: 2.6067
40000/ 200000: 2.0601
50000/ 200000: 2.4988
60000/ 200000: 2.3902
70000/ 200000: 2.1344
80000/ 200000: 2.3369
90000/ 200000: 2.1299
100000/ 200000: 1.8329
110000/ 200000: 2.2053
120000/ 200000: 1.8540
130000/ 200000: 2.4566
140000/ 200000: 2.1879
150000/ 200000: 2.1118
160000/ 200000: 1.8956
170000/ 200000: 1.8644
180000/ 200000: 2.0326
190000/ 200000: 1.8417
train 2.0695888996124268
val 2.131074905395508Bug 2: tanh Saturation
tanh is a squashing function that maps arbitrary inputs smoothly into .
The histogram of hpreact (PRE-activations entering tanh) shows a broad distribution, roughly . This is a problem: most of the probability mass lands in the flat tails of tanh, producing the histogram h of saturated hidden layer activations clustered near .
# o - inspect hidden layer PREactivation values (`hpreact` feeding INTO tanh) and activation values (`h`)
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
# Hidden layer pre-activations (pre-tanh)
print('hpreact.shape:', hpreact.shape)
axes[0].hist(hpreact.view(-1).tolist(), 50)
axes[0].set_title('hidden layer PREactivations `hpreact` (pre-tanh)')
# Hidden layer activations (post-tanh)
print('h.shape:', h.shape)
axes[1].hist(h.view(-1).tolist(), 50)
axes[1].set_title('hidden layer activations `h` (post-tanh activations)')
plt.tight_layout()
plt.show()hpreact.shape: torch.Size([32, 200])
h.shape: torch.Size([32, 200])Why saturation kills backpropagation
Recall, from 04_backprop_train_a_neuron and 07_breaking_up_tanh, the local gradient of tanh is:
Possible resulting scenarios:
- Inactive
tanh: Upstream gradient destroyed: When , this term approaches zero, and thus, by the chain rule:- The upstream gradient is effectively zeroed out (flat tails of the
tanh) - Changes to the pre-activation inputs NO LONGER affect the loss, so the weights and biases feeding into this neuron learn nothing.
- Note the converse: Active
tanh: when , the local gradient is exactly , and the upstream gradient passes through unmodified.- The local gradient is always in , so gradient magnitude through a
tanhcan only stay the same or decrease — it is a one-way squash. - Inspect Desmos graph above.
- The local gradient is always in , so gradient magnitude through a
Visualising saturation (tanh activation grid)
Plot a grid (examples × hidden neurons): white if the tanh output , black otherwise. A high density of white cells indicates widespread saturation. See 04_backprop_train_a_neuron and 07_breaking_up_tanh.
plt.figure(figsize=(20,10))
plt.imshow(h.abs() > 0.99, cmap='gray', interpolation='nearest')<matplotlib.image.AxesImage at 0x11509f770>Dead neurons
A given white pixel is where h.abs() > 0.99; the absolute hidden layer activation was > 0.99. This is the flat, inactive part of the tanh hidden layer neuron.
- So that
tanhneuron (1 of 200 columns) was “inactive” for that particular example (1 of 32 rows)
The critical failure mode is a fully white column
- All 32 examples (rows) producing a saturated activation for the same neuron.
- No matter the input (preactivation, ), the hidden neuron always fires
- That neuron propagates zero gradient for every example,
- So the weights and biases feeding into that neuron never receive an update. It is dead.
Common causes:
- Unlucky initialisation: by chance, all examples land in the inactive
tanhtails at step zero.- Damage was done on first iteration, neurons will never learn over the training process, regardless of iteration count.
- Learning rate too high: a gradient step knocks a neuron’s weights into a region where no example ever activates it again.
Every tanh neuron needs some examples that don’t saturate it, so that gradient can flow through at least some of the time.
Same failure mode in other activations
- Sigmoid: same squashing shape, same saturation problem in the tails
- ReLU: no saturation for positive pre-activations (gradient = 1), but a negative pre-activation produces exactly zero output and zero gradient — the “dying ReLU” problem
Solution
To ensure no dead neurons:
- Scale down
hpreactso its distribution is concentrated near zero (the active region oftanh), instead of .
b1 = torch.randn(n_hidden, generator=g) * 0.01W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * 0.2- This reduces saturation and ensures more neurons receive meaningful gradient signal throughout training.
Since no column is fully white in the current run, we are not in a catastrophic state — but this is a poor starting point.
See also: vanishing-gradient-problem, stats.stackexchange discussion
# o - scale down hpreact by scaling down incoming weights `W1` and biases `b1` (and from before scaled down W2, zero-initialised b2)
n_embd = 10 # the dimensionality of the character embedding vectors
n_hidden = 200 # the number of neurons in the hidden layer of the MLP
g = torch.Generator().manual_seed(2147483647) # for reproducibility
C = torch.randn((vocab_size, n_embd), generator=g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * 0.2
b1 = torch.randn(n_hidden, generator=g) * 0.01
W2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.01
b2 = torch.randn(vocab_size, generator=g) * 0
parameters = [C, W1, b1, W2, b2]
# print(sum(p.nelement() for p in parameters)) # number of parameters in total
for p in parameters:
p.requires_grad = True
# rerun same optimization as last time (200,000 iters; batch size 32 examples (per iter.))
max_steps = 200000
batch_size = 32
lossi = []
print('note first iter loss now reasonable')
for i in range(max_steps):
# minibatch construct
ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)
Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y
# forward pass
emb = C[Xb] # embed the characters into vectors
embcat = emb.view(emb.shape[0], -1) # concatenate the vectors (flattening emb dims)
hpreact = embcat @ W1 + b1 # hidden layer pre-activations
h = torch.tanh(hpreact) # hidden layer activations
logits = h @ W2 + b2 # output layer
loss = F.cross_entropy(logits, Yb) # loss function
# 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())
# break # to view ONLY zero'th iteration loss
# fn: evaluate loss for specific data split. args: 'train', 'val', 'test'
@torch.no_grad() # this decorator disables gradient tracking
def split_loss(split):
x,y = {
'train': (Xtr, Ytr),
'val': (Xdev, Ydev),
'test': (Xte, Yte),
}[split]
# forward pass (calculate loss)
emb = C[x] # embed chars into vectors (N, block_size, n_embd)
embcat = emb.view(emb.shape[0], -1) # concatenate vectors into (N, block_size * n_embd)
hpreact = embcat @ W1 + b1 # hidden layer pre-activ^ns (N, n_hidden)
h = torch.tanh(hpreact) # hidden layer activations (N, n_hidden)
logits = h @ W2 + b2 # output layer (N, vocab_size)
loss = F.cross_entropy(logits, y) # loss function
print(split, loss.item())
split_loss('train')
split_loss('val')
# o - inspect hidden layer PREactivation values (`hpreact` feeding INTO tanh) and activation values (`h`)
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
# Hidden layer pre-activations (pre-tanh)
print('hpreact.shape:', hpreact.shape)
axes[0].hist(hpreact.view(-1).tolist(), 50)
axes[0].set_title('hidden layer PREactivations `hpreact` (pre-tanh)')
# Hidden layer activations (post-tanh)
print('h.shape:', h.shape)
axes[1].hist(h.view(-1).tolist(), 50)
axes[1].set_title('hidden layer activations `h` (post-tanh activations)')
plt.tight_layout()
plt.show()
plt.figure(figsize=(20,10))
plt.imshow(h.abs() > 0.99, cmap='gray', interpolation='nearest')note first iter loss now reasonable
0/ 200000: 3.3135
10000/ 200000: 2.1648
20000/ 200000: 2.3061
30000/ 200000: 2.4541
40000/ 200000: 1.9787
50000/ 200000: 2.2930
60000/ 200000: 2.4232
70000/ 200000: 2.0680
80000/ 200000: 2.3095
90000/ 200000: 2.1207
100000/ 200000: 1.8269
110000/ 200000: 2.2045
120000/ 200000: 1.9797
130000/ 200000: 2.3946
140000/ 200000: 2.1000
150000/ 200000: 2.1948
160000/ 200000: 1.8619
170000/ 200000: 1.7809
180000/ 200000: 1.9673
190000/ 200000: 1.8295
train 2.0355966091156006
val 2.1026782989501953
hpreact.shape: torch.Size([32, 200])
h.shape: torch.Size([32, 200])<matplotlib.image.AxesImage at 0x118b8c190>