makemore_MLP¶

  • Inspired by Andrej Karpathy's "Building makemore Part 2: MLP"

  • Useful links for practice

    • PyTorch internals reference
    "...The talk is in two parts: in the first part, I'm going to first introduce you to the conceptual universe of a tensor library ... The second part grapples with the actual nitty gritty details involved with actually coding in PyTorch..." - Edward Z. Yang



Exercises¶

  1. Tune the hyperparameters of the training to beat karpathy's best validation loss of 2.2.
  2. I was not careful with the intialization of the network in this video.
    1. What is the loss you'd get if the predicted probabilities at initialization were perfectly uniform? What loss do we achieve?
    2. Can you tune the initialization to get a starting loss that is much more similar to (1)?
  3. Read the "A Neural Probabilistic Language Model" paper (Bengio et al. 2003), implement and try any idea from the paper. Did it work?



image

1. Tune Hyperparameters (Val. loss < 2.2)¶

The goal is to find an "optimal" set of hyperparameters that yield a validation loss that is less than 2.2 (val_loss < 2.2). We'll achieve this by iterating over all combinations of two parameters:

  • learn_rate
  • batch_size

and selecting the parameter combination with the lowest validation loss. Other hyperparameters that can be considered for optimization but weren't in this solution include:

  • block_size
  • emb_size
  • hidden_size
In [1]:
import numpy as np
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt # for making figures
%matplotlib inline

import random
from sklearn.model_selection import ParameterGrid
from tqdm import tqdm
In [2]:
words = open('../data/names.txt', 'r').read().splitlines()
In [3]:
chars = sorted(list(set("".join(words))))
stoi = {s:idx + 1 for idx, s in enumerate(chars)}
stoi["."] = 0
itos = {idx: s for s, idx in stoi.items()}
In [4]:
def build_dataset(words, block_size):  
    X, Y = [], []
    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
In [5]:
# Hyperparameters
block_size = 3   # context length: how many characters do we take to predict the next one?
emb_size = 10
in_size = block_size * emb_size
hidden_size = 200
out_size = 27

n_iters = 50000

params = {
    "batch_size": [16, 32, 64, 128, 256],
    "learn_rate": [0.1, 0.05, 0.01]
}


param_grid = ParameterGrid(params)

result_dict = {}
for p in param_grid:
    #print(p)
    result_dict[tuple(p.items())] = [p["batch_size"], p["learn_rate"]]
result_dict
Out[5]:
{(('batch_size', 16), ('learn_rate', 0.1)): [16, 0.1],
 (('batch_size', 16), ('learn_rate', 0.05)): [16, 0.05],
 (('batch_size', 16), ('learn_rate', 0.01)): [16, 0.01],
 (('batch_size', 32), ('learn_rate', 0.1)): [32, 0.1],
 (('batch_size', 32), ('learn_rate', 0.05)): [32, 0.05],
 (('batch_size', 32), ('learn_rate', 0.01)): [32, 0.01],
 (('batch_size', 64), ('learn_rate', 0.1)): [64, 0.1],
 (('batch_size', 64), ('learn_rate', 0.05)): [64, 0.05],
 (('batch_size', 64), ('learn_rate', 0.01)): [64, 0.01],
 (('batch_size', 128), ('learn_rate', 0.1)): [128, 0.1],
 (('batch_size', 128), ('learn_rate', 0.05)): [128, 0.05],
 (('batch_size', 128), ('learn_rate', 0.01)): [128, 0.01],
 (('batch_size', 256), ('learn_rate', 0.1)): [256, 0.1],
 (('batch_size', 256), ('learn_rate', 0.05)): [256, 0.05],
 (('batch_size', 256), ('learn_rate', 0.01)): [256, 0.01]}
In [6]:
random.seed(42)
random.shuffle(words)
n1 = int(0.8*len(words))
n2 = int(0.9*len(words))

Xtr, Ytr = build_dataset(words[:n1], block_size)
Xdev, Ydev = build_dataset(words[n1:n2], block_size)
Xte, Yte = build_dataset(words[n2:], block_size)

torch.Size([182625, 3]) torch.Size([182625])
torch.Size([22655, 3]) torch.Size([22655])
torch.Size([22866, 3]) torch.Size([22866])
In [7]:
param_dict = {}
for pset in tqdm(param_grid):    
    g = torch.Generator().manual_seed(2147483647)
    C = torch.randn((out_size, emb_size), generator = g)
    W1 = torch.randn((in_size, hidden_size), generator = g)
    b1 = torch.randn(hidden_size, generator = g)
    W2 = torch.randn((hidden_size, out_size), generator = g)
    b2 = torch.randn(out_size, generator = g)

    parameters = [C, W1, b1, W2, b2]
    # print(f"number of params = {sum(p.nelement() for p in parameters)}")

    for p in parameters:
        p.requires_grad = True

    for i in range(n_iters):

        # minibatch construct
        ix = torch.randint(0, Xtr.shape[0], (pset["batch_size"], ))

        # forward pass
        emb = C[Xtr[ix]]                               
        h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)     
        logits = h @ W2 + b2                           
        loss = F.cross_entropy(logits, Ytr[ix])
        #print(loss.item())

        # backward pass
        for p in parameters:
            p.grad = None
        loss.backward()

        # update parameters
        lr = pset["learn_rate"] if i < n_iters/2 else pset["learn_rate"]/10
        for p in parameters:
            p.data += -lr * p.grad

        # store training loss
    #param_dict[tuple(pset.items())] = np.round(loss.item(), decimals=4) 
    #print(f"parameters: {pset}, training loss: {loss.item():.4f}, validation loss: {val_loss.item():.4f}")
    
    with torch.no_grad():
        emb = C[Xdev]
        h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)
        logits = h @ W2 + b2
        val_loss = F.cross_entropy(logits, Ydev)
        
    param_dict[tuple(pset.items())] = [np.round(loss.item(), decimals=4), np.round(val_loss.item(), decimals=4)]
    print(f"parameters: {pset}, training loss: {loss.item():.4f}, validation loss: {val_loss.item():.4f}")

  7%|▋         | 1/15 [00:18<04:14, 18.16s/it]
parameters: {'batch_size': 16, 'learn_rate': 0.1}, training loss: 2.5563, validation loss: 2.2570

 13%|█▎        | 2/15 [00:35<03:52, 17.91s/it]
parameters: {'batch_size': 16, 'learn_rate': 0.05}, training loss: 1.9905, validation loss: 2.3013

 20%|██        | 3/15 [00:52<03:32, 17.68s/it]
parameters: {'batch_size': 16, 'learn_rate': 0.01}, training loss: 2.6036, validation loss: 2.5070

 27%|██▋       | 4/15 [01:13<03:24, 18.63s/it]
parameters: {'batch_size': 32, 'learn_rate': 0.1}, training loss: 2.2959, validation loss: 2.2351

 33%|███▎      | 5/15 [01:35<03:16, 19.63s/it]
parameters: {'batch_size': 32, 'learn_rate': 0.05}, training loss: 2.1301, validation loss: 2.2860

 40%|████      | 6/15 [01:56<03:00, 20.06s/it]
parameters: {'batch_size': 32, 'learn_rate': 0.01}, training loss: 2.6998, validation loss: 2.5056

 47%|████▋     | 7/15 [02:20<02:50, 21.31s/it]
parameters: {'batch_size': 64, 'learn_rate': 0.1}, training loss: 2.0842, validation loss: 2.2304

 53%|█████▎    | 8/15 [02:45<02:35, 22.23s/it]
parameters: {'batch_size': 64, 'learn_rate': 0.05}, training loss: 2.3601, validation loss: 2.2766

 60%|██████    | 9/15 [03:09<02:17, 22.89s/it]
parameters: {'batch_size': 64, 'learn_rate': 0.01}, training loss: 2.6763, validation loss: 2.5050

 67%|██████▋   | 10/15 [03:40<02:06, 25.29s/it]
parameters: {'batch_size': 128, 'learn_rate': 0.1}, training loss: 2.2146, validation loss: 2.2326

 73%|███████▎  | 11/15 [04:11<01:47, 26.94s/it]
parameters: {'batch_size': 128, 'learn_rate': 0.05}, training loss: 2.1858, validation loss: 2.2696

 80%|████████  | 12/15 [04:42<01:24, 28.15s/it]
parameters: {'batch_size': 128, 'learn_rate': 0.01}, training loss: 2.5317, validation loss: 2.5254

 87%|████████▋ | 13/15 [05:25<01:05, 32.70s/it]
parameters: {'batch_size': 256, 'learn_rate': 0.1}, training loss: 2.2779, validation loss: 2.2428

 93%|█████████▎| 14/15 [06:09<00:35, 35.98s/it]
parameters: {'batch_size': 256, 'learn_rate': 0.05}, training loss: 2.2796, validation loss: 2.2738

100%|██████████| 15/15 [06:56<00:00, 27.75s/it]
parameters: {'batch_size': 256, 'learn_rate': 0.01}, training loss: 2.6329, validation loss: 2.5172





















In [8]:






# Extract the parameters and validation losses from param_dict
params = list(param_dict.keys())
val_losses = [param_dict[p][1] for p in params]

# Create a 2D array to store the validation losses
val_loss_array = np.array(val_losses).reshape(len(set(p[0][1] for p in params)), len(set(p[1][1] for p in params)))

# Get the unique batch sizes and learning rates
batch_sizes = sorted(set(p[0][1] for p in params))
learn_rates = sorted(set(p[1][1] for p in params))

# Create a heatmap
plt.imshow(val_loss_array, cmap='viridis', interpolation='nearest')
#plt.imshow(val_loss_array, cmap='gray_r', interpolation='nearest')
plt.xlabel('Learning Rate')
plt.ylabel('Batch Size')
plt.title('Validation Loss')
plt.colorbar()

# Set the tick labels and values
plt.xticks(range(len(learn_rates)), learn_rates)
plt.yticks(range(len(batch_sizes)), batch_sizes)
plt.show()
























No description has been provided for this image



















In [9]:






param_dict





















Out[9]:




{(('batch_size', 16), ('learn_rate', 0.1)): [2.5563, 2.257],
 (('batch_size', 16), ('learn_rate', 0.05)): [1.9905, 2.3013],
 (('batch_size', 16), ('learn_rate', 0.01)): [2.6036, 2.507],
 (('batch_size', 32), ('learn_rate', 0.1)): [2.2959, 2.2351],
 (('batch_size', 32), ('learn_rate', 0.05)): [2.1301, 2.286],
 (('batch_size', 32), ('learn_rate', 0.01)): [2.6998, 2.5056],
 (('batch_size', 64), ('learn_rate', 0.1)): [2.0842, 2.2304],
 (('batch_size', 64), ('learn_rate', 0.05)): [2.3601, 2.2766],
 (('batch_size', 64), ('learn_rate', 0.01)): [2.6763, 2.505],
 (('batch_size', 128), ('learn_rate', 0.1)): [2.2146, 2.2326],
 (('batch_size', 128), ('learn_rate', 0.05)): [2.1858, 2.2696],
 (('batch_size', 128), ('learn_rate', 0.01)): [2.5317, 2.5254],
 (('batch_size', 256), ('learn_rate', 0.1)): [2.2779, 2.2428],
 (('batch_size', 256), ('learn_rate', 0.05)): [2.2796, 2.2738],
 (('batch_size', 256), ('learn_rate', 0.01)): [2.6329, 2.5172]}




















In [10]:






min_key = min(param_dict, key=lambda k: param_dict[k][1]) # retrieve 2nd element in dict value (val loss)
min_value = param_dict[min_key]
print(f"The smallest validation loss value is {min_value[1]} and the parameters are {min_key}")

























The smallest validation loss value is 2.2304 and the parameters are (('batch_size', 64), ('learn_rate', 0.1))





















In [11]:






min_key[0][1], min_key[1][1]





















Out[11]:




(64, 0.1)




















In [12]:






g = torch.Generator().manual_seed(2147483647)
C = torch.randn((out_size, emb_size), generator = g)
W1 = torch.randn((in_size, hidden_size), generator = g)
b1 = torch.randn(hidden_size, generator = g)
W2 = torch.randn((hidden_size, out_size), generator = g)
b2 = torch.randn(out_size, generator = g)
parameters = [C, W1, b1, W2, b2]
print(f"number of params = {sum(p.nelement() for p in parameters)}")

























number of params = 11897





















In [13]:






learn_rate = 0.1
batch_size = 256
n_iters = 200000





















In [14]:






%%time
# Essential to always enable gradients !!!
for p in parameters:
    p.requires_grad = True

# Training loop !!!
for i in range(n_iters):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ))

    # forward pass
    emb = C[Xtr[ix]]                               
    h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)     
    logits = h @ W2 + b2                           
    loss = F.cross_entropy(logits, Ytr[ix])
    # add regularization
    #loss += 0.01 * (W1 ** 2).mean() + 0.01 * (W2 ** 2).mean()
    #print(loss.item())

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update parameters
    if i > 0.5 * n_iters and i < 0.75 * n_iters:
        lr = learn_rate / 10
    elif i >= 0.75 * n_iters:
        lr = learn_rate / 100
    else:
        lr = learn_rate
#     r = 10E-5 # n_iters = 300000 & lr_init = 0.1
#     lr = 0.1 / (1 + r*i)   
    for p in parameters:
        p.data += -lr * p.grad
        
    if i % 10000 == 0:
        print(f"mini batch {i} | loss {loss.item():.4f} at lr {lr}")
        
with torch.no_grad():
    emb = C[Xdev]
    h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)
    logits = h @ W2 + b2
    val_loss = F.cross_entropy(logits, Ydev)
    print(f"val_loss: {val_loss.item()}")

























mini batch 0 | loss 26.6556 at lr 0.1
mini batch 10000 | loss 2.3171 at lr 0.1
mini batch 20000 | loss 2.2889 at lr 0.1
mini batch 30000 | loss 2.2101 at lr 0.1
mini batch 40000 | loss 2.2581 at lr 0.1
mini batch 50000 | loss 2.2627 at lr 0.1
mini batch 60000 | loss 2.3664 at lr 0.1
mini batch 70000 | loss 2.0849 at lr 0.1
mini batch 80000 | loss 2.3078 at lr 0.1
mini batch 90000 | loss 2.1750 at lr 0.1
mini batch 100000 | loss 2.1651 at lr 0.1
mini batch 110000 | loss 2.0781 at lr 0.01
mini batch 120000 | loss 2.1338 at lr 0.01
mini batch 130000 | loss 2.1463 at lr 0.01
mini batch 140000 | loss 2.1469 at lr 0.01
mini batch 150000 | loss 2.0903 at lr 0.001
mini batch 160000 | loss 2.1117 at lr 0.001
mini batch 170000 | loss 2.1610 at lr 0.001
mini batch 180000 | loss 2.1123 at lr 0.001
mini batch 190000 | loss 2.1595 at lr 0.001
val_loss: 2.160660982131958
CPU times: user 3min 29s, sys: 23.3 s, total: 3min 53s
Wall time: 3min 3s





















In [15]:






# Evaluation - Training Loss

emb_tr = C[Xtr]                               
h_tr  = torch.tanh(emb_tr.view(-1, in_size) @ W1 + b1)   
logits_tr = h_tr @ W2 + b2                        
loss_tr = F.cross_entropy(logits_tr, Ytr)
print(f"The training loss: {loss_tr.item():.4f}")

# Evaluation - Validation Loss
emb_dev = C[Xdev]                                
h_dev = torch.tanh(emb_dev.view(-1, in_size) @ W1 + b1)    
logits_dev = h_dev @ W2 + b2                        
loss_dev = F.cross_entropy(logits_dev, Ydev)
print(f"The validation loss: {loss_dev.item():.4f}")

























The training loss: 2.1297
The validation loss: 2.1607

























Comment: The validation loss of the model based on batch_size = 256 and learn_rate = 0.01 is 2.152 which beats 2.2.


























2a. Initialization: Perfectly Uniform Predicted Probabilities¶


We have 27 possible characters (26 letters of the alphabet and 1 special token, .). If every character was equally likely, then we'd expect all the probabilities to be roughly 4% (1/27.0 ~= 3.7%), in essence the predicted probabilities at initialization are perfectly uniform.


$$\\{logits}_{ij} = \frac{e^{C_{ij}}}{\sum_{j} e^{C_{ij}}} = \frac{C}{27C} = \frac{1}{27}$$


















In [16]:






block_size = 3
emb_size = 10
in_size = block_size * emb_size
hidden_size = 200
out_size = 27





















In [17]:






C = torch.full((out_size, emb_size), fill_value=1.0)
W1 = torch.full((in_size, hidden_size), fill_value=1.0)
b1 = torch.full((1,hidden_size), fill_value=1.0)
W2 = torch.full((hidden_size, out_size), fill_value=1.0)
b2 = torch.full((1,out_size), fill_value=1.0)





















In [18]:






# Forward pass
emb = C[Xtr]
h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)
logits = h @ W2 + b2
print(torch.softmax(logits, dim=1))
loss = F.cross_entropy(logits, Ytr)

print(f"The training loss: {loss.item():.4f}\n")

emb_dev = C[Xdev]                                
h_dev = torch.tanh(emb_dev.view(-1, in_size) @ W1 + b1)    
logits_dev = h_dev @ W2 + b2                        
loss_dev = F.cross_entropy(logits_dev, Ydev)
print(f"The validation loss: {loss_dev.item():.4f}")

























tensor([[0.0370, 0.0370, 0.0370,  ..., 0.0370, 0.0370, 0.0370],
        [0.0370, 0.0370, 0.0370,  ..., 0.0370, 0.0370, 0.0370],
        [0.0370, 0.0370, 0.0370,  ..., 0.0370, 0.0370, 0.0370],
        ...,
        [0.0370, 0.0370, 0.0370,  ..., 0.0370, 0.0370, 0.0370],
        [0.0370, 0.0370, 0.0370,  ..., 0.0370, 0.0370, 0.0370],
        [0.0370, 0.0370, 0.0370,  ..., 0.0370, 0.0370, 0.0370]])
The training loss: 3.2958

The validation loss: 3.2958

























Comment: The loss, if the predicted probabilities at initialization were perfectly uniform, is 3.296


























2b. Tune Initialization (Similar Val. loss to that of No. 1)¶


We initialize the weights and biases to be as close to zero as possible. I calibrated the variances with sqrt (2/n) based on a stanford cs231n lecture on weight initialization.




















In [19]:






np.sqrt(2/len(words))





















Out[19]:




0.00790162092698779




















In [20]:






g = torch.Generator().manual_seed(2147483647)
C = torch.randn((out_size, emb_size), generator = g) 
W1 = torch.randn((in_size, hidden_size), generator = g) * np.sqrt(2/len(words))
b1 = torch.randn(hidden_size, generator = g) * 0.0
W2 = torch.randn((hidden_size, out_size), generator = g) * np.sqrt(2/len(words))
b2 = torch.randn(out_size, generator = g) * 0.0
parameters = [C, W1, b1, W2, b2]
print(f"number of params = {sum(p.nelement() for p in parameters)}")

























number of params = 11897





















In [21]:






%%time
n_iters = 50000
# Essential to always enable gradients !!!
for p in parameters:
    p.requires_grad = True

# Training loop !!!
for i in range(n_iters):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ))

    # forward pass
    emb = C[Xtr[ix]]                               
    h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)     
    logits = h @ W2 + b2                           
    loss = F.cross_entropy(logits, Ytr[ix])
    # add regularization
    #loss += 0.01 * (W1 ** 2).mean() + 0.01 * (W2 ** 2).mean()
    #print(loss.item())

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update parameters
    if i > 0.5 * n_iters:
        lr = learn_rate / 10
    elif i >= 0.75 * n_iters:
        lr = learn_rate / 100
    else:
        lr = learn_rate
#     r = 10E-5 # n_iters = 300000 & lr_init = 0.1
#     lr = 0.1 / (1 + r*i)   
    for p in parameters:
        p.data += -lr * p.grad
        
    if i % 10000 == 0:
        print(f"mini batch {i} | loss {loss.item():.4f} at lr {lr}")
        
with torch.no_grad():
    emb = C[Xdev]
    h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)
    logits = h @ W2 + b2
    val_loss = F.cross_entropy(logits, Ydev)
    print(f"\nval_loss: {val_loss.item():.4f}\n")

























mini batch 0 | loss 3.2959 at lr 0.1
mini batch 10000 | loss 2.2430 at lr 0.1
mini batch 20000 | loss 2.2546 at lr 0.1
mini batch 30000 | loss 2.1073 at lr 0.01
mini batch 40000 | loss 2.2433 at lr 0.01

val_loss: 2.1384

CPU times: user 54.1 s, sys: 6.18 s, total: 1min
Wall time: 47.3 s

























Comment: After tuning the initialization, the validation loss after training is 2.138 which is similar to the validation loss in no. 1 of 2.152. The starting loss here of 3.296 is better than that of no. 1 which is 25.74.


























3. Implement New Idea from Bengio 2003 Neural Probabilistic LM Paper¶






















Let's add an additional layer, the word features input to output weights $W$ (a $|V| \times(n-1) m$ matrix). We'll denote it as $W_{0}$.


$$
\begin{equation*}
y=b_{2}+W_{0} x+W_{2} \tanh (b_{1}+W_{1} x) \tag{1}
\end{equation*}
$$


$$
parameters: \theta=(b_{1}, b_{2}, W_{0}, W_{1}, W_{2}, C)
$$

Also, we'll add a ridge regularization penalty and only apply it to the weights ($W_{0}, W_{1}, W_{2}$) and not the biases as per the paper. Finally, we'll implement SGD with momentum since our gradient trajectory is noisy during our mini-batch gradient descent training loop, and this limits our use of large learning rates.


















In [22]:






block_size = 3
emb_size = 10
in_size = block_size * emb_size
hidden_size = 200
out_size = 27





















In [23]:






g = torch.Generator().manual_seed(2147483647)
C = torch.randn((out_size, emb_size), generator = g)
W1 = torch.randn((in_size, hidden_size), generator = g)
b1 = torch.randn(hidden_size, generator = g)
W2 = torch.randn((hidden_size, out_size), generator = g)
b2 = torch.randn(out_size, generator = g)
W0 = torch.randn((in_size, out_size), generator=g)
parameters = [C, W1, b1, W2, b2, W0]
param_grads = [0.0 for _ in range(len(parameters))]
print(f"number of params = {sum(p.nelement() for p in parameters)}")

























number of params = 12707





















In [24]:






learn_rate = 0.1
batch_size = 256
n_iters = 200000
beta = 0.9 # average over last 10 values.
reg = 0.01





















In [25]:






%%time
# Essential to always enable gradients !!!
for p in parameters:
    p.requires_grad = True

# Training loop !!!
for i in range(n_iters):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ))

    # forward pass
    emb = C[Xtr[ix]]                               
    h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)     
    logits = emb.view(-1, in_size) @ W0 + h @ W2 + b2                           
    loss = F.cross_entropy(logits, Ytr[ix])
    # add regularization
    loss += reg * ( (W0 ** 2).mean() + (W1 ** 2).mean() +  (W2 ** 2).mean())
    #print(loss.item())

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update parameters
    if i > 0.5 * n_iters and i < 0.75 * n_iters:
        lr = learn_rate / 10
    elif i >= 0.75 * n_iters:
        lr = learn_rate / 100
    else:
        lr = learn_rate
        
    for idx, p in enumerate(parameters):
        param_grads[idx] = beta * param_grads[idx] + (1 - beta) * p.grad
        p.data = p.data - lr * param_grads[idx]

        
    if i % 10000 == 0:
        print(f"mini batch {i} | loss {loss.item():.4f} at lr {lr}")
        
with torch.no_grad():
    emb = C[Xdev]
    h = torch.tanh(emb.view(-1, in_size) @ W1 + b1)
    logits = emb.view(-1, in_size) @ W0 + h @ W2 + b2
    val_loss = F.cross_entropy(logits, Ydev)
    print(f"val_loss: {val_loss.item()}")

























mini batch 0 | loss 29.5452 at lr 0.1
mini batch 10000 | loss 2.2582 at lr 0.1
mini batch 20000 | loss 2.2082 at lr 0.1
mini batch 30000 | loss 2.3550 at lr 0.1
mini batch 40000 | loss 2.4241 at lr 0.1
mini batch 50000 | loss 2.2602 at lr 0.1
mini batch 60000 | loss 2.1515 at lr 0.1
mini batch 70000 | loss 2.2701 at lr 0.1
mini batch 80000 | loss 2.2038 at lr 0.1
mini batch 90000 | loss 2.1660 at lr 0.1
mini batch 100000 | loss 2.1656 at lr 0.1
mini batch 110000 | loss 2.2317 at lr 0.01
mini batch 120000 | loss 2.2038 at lr 0.01
mini batch 130000 | loss 2.3513 at lr 0.01
mini batch 140000 | loss 2.1198 at lr 0.01
mini batch 150000 | loss 2.1437 at lr 0.001
mini batch 160000 | loss 2.2227 at lr 0.001
mini batch 170000 | loss 2.2649 at lr 0.001
mini batch 180000 | loss 2.0970 at lr 0.001
mini batch 190000 | loss 2.0675 at lr 0.001
val_loss: 2.155996799468994
CPU times: user 4min 53s, sys: 23.7 s, total: 5min 17s
Wall time: 4min 31s





















In [26]:






# Evaluation - Training Loss

emb_tr = C[Xtr]                               
h_tr  = torch.tanh(emb_tr.view(-1, in_size) @ W1 + b1)   
logits_tr = emb_tr.view(-1, in_size) @ W0 + h_tr @ W2 + b2                      
loss_tr = F.cross_entropy(logits_tr, Ytr)
print(f"The training loss: {loss_tr.item():.4f}")

# Evaluation - Validation Loss
emb_dev = C[Xdev]                                
h_dev = torch.tanh(emb_dev.view(-1, in_size) @ W1 + b1)    
logits_dev = emb_dev.view(-1, in_size) @ W0 + h_dev @ W2 + b2                       
loss_dev = F.cross_entropy(logits_dev, Ydev)
print(f"The validation loss: {loss_dev.item():.4f}")

























The training loss: 2.1310
The validation loss: 2.1560

























Comment: The validation loss is 2.156 with the addition of a new layer $W_{0}$, ridge regularization penalty, and stochastic gradient descent momentum. It is roughly the same as the validation loss of 2.1524 in no. 1.
























Sampling from the model¶


















In [27]:






# 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:
        emb = C[torch.tensor([context])] # (1,block_size,d)
        h = torch.tanh(emb.view(1, -1) @ W1 + b1)
        logits = emb.view(-1, in_size) @ W0 + h @ W2 + b2
        probs = F.softmax(logits, dim=1)
        ix = torch.multinomial(probs, num_samples=1, generator=g).item()
        context = context[1:] + [ix]
        out.append(ix)
        if ix == 0:
            break
    print(''.join(itos[i] for i in out))

























carmah.
ambrilli.
kimri.
rehty.
skaf.
ske.
rahnen.
den.
rha.
kaqui.
ner.
kia.
chaiir.
kaleigph.
bmandin.
quint.
sulin.
alian.
quinathon.
jarynix.