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.

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 set
torch.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 = True
11897
# 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 loss
note 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>]
plot

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 word
carlah.
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).

ScenarioLogitsp[2] softmaxLoss
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)
5randn(4)variesunpredictable — likely >> 1.386
  • Scenario 1 is the ideal initialisation — all logits equal, softmax is uniform, every class (i.e. character) gets 1/N probability.
  • 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.randn weights: 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 just torch.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.131074905395508

Bug 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])
plot

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:

  1. 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.
  2. 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 tanh can only stay the same or decrease — it is a one-way squash.
    • Inspect Desmos graph above.

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>
plot

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 tanh neuron (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 tanh tails 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 hpreact so its distribution is concentrated near zero (the active region of tanh), instead of .
    • b1 = torch.randn(n_hidden, generator=g) * 0.01
    • W1 = 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])
plot
<matplotlib.image.AxesImage at 0x118b8c190>
plot