notes/

Kaiming Initialisation

9 min read

Goals

In the previous chapter we set the parameter scaling factors manually, for each layer. This is tedious and arbitrary for larger networks. Kaiming Initialisation is a method to automatically set the scaling factors for each layer, based on the number of input and output connections.

# 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

Bug: Non-Gaussian activations

Incorrect initialisation of weights can lead to non-Gaussian activations, which can cause problems during training. See the following example

# o - example
x = torch.randn(1000, 10)  # input data: 1000 training examples, each a 10-dim feature vector
w = torch.randn(10, 200)   # weight matrix: maps 10-dim input → 200-dim hidden layer (output)
y = x @ w                  # hidden layer preactivations: (1000, 200) — one 200-dim vector per example
y_w_scaled_x5 = x @ (w * 5)
y_w_scaled_x0_2 = x @ (w * 0.2)
y_w_scaled_root10 = x @ (w * 10**-0.5)
 
x_mean, x_std = x.mean().item(), x.std().item()
y_mean, y_std = y.mean().item(), y.std().item()
y_w_scaled_x5_mean, y_w_scaled_x5_std = y_w_scaled_x5.mean().item(), y_w_scaled_x5.std().item()
y_w_scaled_x0_2_mean, y_w_scaled_x0_2_std = y_w_scaled_x0_2.mean().item(), y_w_scaled_x0_2.std().item()
y_w_scaled_root10_mean, y_w_scaled_root10_std = y_w_scaled_root10.mean().item(), y_w_scaled_root10.std().item()
print('input, x                   ', x.mean(), x.std())
print('y (no scaling)             ', y.mean(), y.std(), "  std dev. NOT Gaussian")
print('y (w scaled by 5)          ', y_w_scaled_x5.mean(), y_w_scaled_x5.std(), " std dev. NOT Gaussian")
print('y (w scaled by 0.2)        ', y_w_scaled_x0_2.mean(), y_w_scaled_x0_2.std(), "  std dev. NOT Gaussian")
print('y (w scaled by 1/sqrt(10)) ', y_w_scaled_root10.mean(), y_w_scaled_root10.std(), "  std dev. IS Gaussian")
 
plt.figure(figsize=(20, 16))
plt.subplot(321)
plt.hist(x.view(-1).tolist(), 50, density=True)
plt.title(f'x is Gaussian: mean={x_mean:.4f}, std={x_std:.4f}')
 
plt.subplot(323)
plt.hist(y.view(-1).tolist(), 50, density=True , color='maroon')
plt.title(f'y is NOT Gaussian: mean={y_mean:.4f}, std={y_std:.4f} (should be ~1)')
plt.subplot(324)
plt.hist(y_w_scaled_x5.view(-1).tolist(), 50, density=True, color='maroon')
plt.title(f'y (w scaled by 5) NOT Gaussian: mean={y_w_scaled_x5_mean:.4f}, std={y_w_scaled_x5_std:.4f}')
plt.subplot(325)
plt.hist(y_w_scaled_x0_2.view(-1).tolist(), 50, density=True, color='maroon')
plt.title(f'y (w scaled by 0.2) NOT Gaussian: mean={y_w_scaled_x0_2_mean:.4f}, std={y_w_scaled_x0_2_std:.4f}')
plt.subplot(326)
plt.hist(y_w_scaled_root10.view(-1).tolist(), 50, density=True, color='green')
plt.title(f'y (w scaled by 1/sqrt(10)) IS Gaussian: mean={y_w_scaled_root10_mean:.4f}, std={y_w_scaled_root10_std:.4f}')
input, x                    tensor(-2.7266e-05) tensor(0.9987)
y (no scaling)              tensor(0.0085) tensor(3.1740)   std dev. NOT Gaussian
y (w scaled by 5)           tensor(0.0427) tensor(15.8701)  std dev. NOT Gaussian
y (w scaled by 0.2)         tensor(0.0017) tensor(0.6348)   std dev. NOT Gaussian
y (w scaled by 1/sqrt(10))  tensor(0.0027) tensor(1.0037)   std dev. IS Gaussian
Text(0.5, 1.0, 'y (w scaled by 1/sqrt(10)) IS Gaussian: mean=0.0027, std=1.0037')
plot

Bug: Hidden layer preactivations y = x @ w are NOT Gaussian

  • Std. dev. of y is ~3.2, ~15.8, or ~0.6, whereas it should be ~1

Kaiming initialisation

Scale the weight layer w by the fan-in, so y remains Gaussian

  • A good scaling factor is , where fan_in is the number of input connections to the layer (i.e. the number of columns in w)
  • In this case, the fan-in is 10, so scaling by , see green distribution above)

Why this works

  • The goal of Kaiming Initialisation is to maintain the variance of the activations across layers.
    • If the variance is too high, it can lead to exploding gradients.
    • If the variance is too low, it can lead to vanishing gradients.
    • By scaling the weights appropriately, we can ensure that the variance of the activations remains stable across layers, which helps with training deep networks.

Modern techniques mitigate poor initialisation issues

  • Residual connections (skip connections), and Normalisation layers (BatchNorm, LayerNorm, GroupNorm) help maintain the variance of the activations across layers, even with suboptimal initialisation.
  • Better optimisers than SGD (e.g., Adam, RMSprop) can also help mitigate the issues related to poor initialisation.

Reset MLP, and test Kaiming Init

# 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])

Apply Kaiming init to the weight layer W1

Scale the weight layer W1 by gain / sqrt(fan_in) (see docs)

For tanh activations,

  • gain is 5/3,
  • fan_in is the number of input connections to the layer
    • i.e. the number of columns in W1, (n_embed * block_size = 10 * 3 = 30)
# 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) * (5/3)/((n_embd * block_size)**0.5)
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
11897

Training (200k iterations)

Kaiming init means we don’t have to guess the right scaling factor for the weight layer W1 to maintain Gaussian activations across layers. Testing the activations will show that activations are now Gaussian with mean 0 and std 1.

# 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: 3.3179
  10000/ 200000: 2.1910
  20000/ 200000: 2.3270
  30000/ 200000: 2.5396
  40000/ 200000: 1.9468
  50000/ 200000: 2.3331
  60000/ 200000: 2.3852
  70000/ 200000: 2.1173
  80000/ 200000: 2.3159
  90000/ 200000: 2.2010
 100000/ 200000: 1.8591
 110000/ 200000: 2.0881
 120000/ 200000: 1.9389
 130000/ 200000: 2.3913
 140000/ 200000: 2.0949
 150000/ 200000: 2.1458
 160000/ 200000: 1.7824
 170000/ 200000: 1.7249
 180000/ 200000: 1.9751
 190000/ 200000: 1.8614

Loss: Not much better, but no guesswork for scaling W1!

Note, the training and validation losses are not much better, but we didn’t have to guess the right scaling factor for W1 to get Gaussian activations, which is a big win. In larger networks, it would be much harder to guess the right scaling factor for each layer, and Kaiming init helps with that.

# o - 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.0376641750335693
val 2.106989622116089
← Previous 05 Experiment 2: Increase embedding dims