Implementing an MLP Character-Level Language Model with Core Machine Learning Concepts (Makemore Part 2)

Andrej Karpathy

Introduction and Motivation [00:00:00]

The video continues the "makemore" series, moving from simple bigram models to a more advanced neural network. The previous bigram model used only a single preceding character, leading to poor name generation and an exponential increase in table size when considering longer contexts.

Bengio et al. 2003 (MLP Language Model) Paper Walkthrough [00:01:48]

The new approach is based on the influential Bengio et al. [2003] paper, "A Neural Probabilistic Language Model," which introduced using neural networks for language modeling.

• Addressing the curse of dimensionality: The paper proposes associating each word with a lower-dimensional feature vector, or "embedding," (e.g., 30 dimensions for a vocabulary of 17,000 words).
  • These embeddings are initialized randomly but are tuned during training via backpropagation.
  • Words with similar meanings or contexts are expected to cluster together in this embedding space, enabling generalization to unseen phrases.
• Neural Network Architecture:
  • The model takes N previous words (or characters in this video's implementation) as input.
  • Each input word index is looked up in an embedding table (C) to get its feature vector.
  • These embeddings are concatenated to form the input to a hidden layer.
  • The hidden layer applies a tanh non-linearity.
  • An output layer then projects the hidden state to a vector of logits, representing scores for each possible next word/character.
  • A softmax function converts these logits into a probability distribution over the next word/character.
  • The model is trained by maximizing the log-likelihood of the correct next word/character.

    

    • The C matrix and the weights and biases of the hidden and output layers are the parameters optimized by backpropagation.

(Re-)building Our Training Dataset [00:09:03]

The dataset consists of a list of names.

• Vocabulary and Mappings:
  • Create a mapping from characters (strings) to integers and vice-versa, including a special '.' token for start/end of words.
• Context Length (block_size):
  • Defined as the number of previous characters used to predict the next. For this video, it's initially set to 3.
• Generating (Context, Target) Pairs:
  • For each word in the dataset, a rolling window approach is used.
  • Example: For "emma" and block_size=3:
    • Context: ..., Target: e
    • Context: ..e, Target: m
    • Context: .em, Target: m
    • Context: emm, Target: a
    • Context: mma, Target: .
  • These pairs form the input X (context as integer indices) and target Y (next character's integer index).

    

  • Initially, a small subset (e.g., first 5 words) is used to quickly develop the code. The full dataset contains over 228,000 examples.

    

Implementing the Embedding Lookup Table [00:12:19]

• Embedding Matrix C:
  • A torch.randn tensor is initialized with dimensions (vocab_size, embedding_dim).
  • Initially, embedding_dim is set to 2 to allow for visualization. For 27 possible characters, C is a 27x2 matrix.
• Indexing C:
  • PyTorch allows direct indexing of C using integer tensors for efficiency.
  • C[X] where X is a (num_examples, block_size) tensor of integers, will directly return a (num_examples, block_size, embedding_dim) tensor of embeddings.
  • This is computationally equivalent to a one-hot encoding followed by matrix multiplication, but much faster.

Implementing the Hidden Layer + Internals of torch.Tensor: Storage, Views [00:18:35]

• Weights and Biases:
  • W1 (weights) is (input_dim, hidden_layer_size), and B1 (biases) is (hidden_layer_size).
  • The input_dim for the hidden layer is block_size * embedding_dim (e.g., 3 * 2 = 6 initially).
  • hidden_layer_size is a hyperparameter, initially 100.
• Reshaping Embeddings:
  • The (num_examples, block_size, embedding_dim) embeddings need to be reshaped to (num_examples, block_size * embedding_dim) for matrix multiplication with W1.
  • This is achieved efficiently using embeddings.view(-1, block_size * embedding_dim).
  • torch.Tensor Internals (Storage and Views):
    • view() creates a new logical interpretation of the same underlying physical memory (storage) without copying data. This is very efficient.
    • torch.cat() (concatenate) creates new memory to store the concatenated tensor, making it less efficient for reshaping operations when view is possible.

      

• Hidden State Calculation:
  • h = torch.tanh(embeddings_reshaped @ W1 + B1).
  • Biases (B1) are automatically broadcasted across the num_examples dimension.

    

Implementing the Output Layer [00:29:15]

• Weights and Biases:
  • W2 (weights) is (hidden_layer_size, vocab_size), and B2 (biases) is (vocab_size).
  • The vocab_size is 27 (for 26 letters + '.').
• Logits Calculation:
  • logits = h @ W2 + B2. These are the raw scores before probability normalization.

Implementing the Negative Log Likelihood Loss [00:29:53]

• Manual Loss Calculation Steps:
  1. counts = logits.exp(): Exponentiate logits to get "fake counts" (non-normalized probabilities).
  2. probs = counts / counts.sum(1, keepdim=True): Normalize counts across rows to get probabilities summing to 1.
  3. loss = -probs[range(num_examples), Y].log().mean(): Select the probability assigned to the true target character for each example, take its log, negate, and then average.

     

Summary of the Full Network [00:32:17]

All parameters (embedding table C, W1, B1, W2, B2) are defined. The initial loss (before training) is around 17, indicating random predictions.

Introducing F.cross_entropy and Why [00:32:49]

PyTorch's torch.nn.functional.cross_entropy function should be used in practice for classification loss.

• Benefits:
  1. Efficiency: It internally fuses multiple operations into optimized kernels, avoiding the creation of intermediate tensors and reducing memory overhead.
  2. Simplified Backward Pass: The analytical derivative of cross-entropy is often simpler and more efficient to compute than chaining individual derivatives.
  3. Numerical Stability: It handles extreme logit values gracefully.
  • Large positive logits can cause exp() to overflow (resulting in inf then NaN).
  • F.cross_entropy subtracts the maximum logit from all logits before exponentiation. This "shifting" doesn't change the final probabilities (due to softmax properties) but brings numbers into a safer range (max logit becomes 0, others become negative), preventing overflow.

    

Implementing the Training Loop, Overfitting One Batch [00:37:56]

• Training Steps:
  1. Set all parameter gradients to zero (p.grad.zero_()).
  2. Perform a forward pass to calculate the loss (loss = F.cross_entropy(logits, Y)).
  3. Perform a backward pass (loss.backward()) to compute gradients for all parameters.
  4. Update parameters (p.data -= learning_rate * p.grad).
• Overfitting a Small Batch:
  • Initially, the model is trained on a small subset (e.g., 32 examples).
  • With 3,400 parameters, the model quickly achieves a very low loss (e.g., 0.05), demonstrating its capacity to memorize the data.
  • The loss doesn't reach exactly zero if multiple target characters are associated with the same context in the dataset (e.g., ... can predict 'e', 'o', 'a', 's').

    

Training on the Full Dataset, Minibatches [00:41:25]

• Full Dataset:
  • The full dataset has over 228,000 examples. Processing all of them in one go (full-batch gradient descent) is computationally expensive and slow.
• Mini-batch Gradient Descent:
  • In practice, mini-batches of data are randomly sampled for each training step.
  • This provides an approximate gradient direction but allows for many more updates in the same amount of time, leading to faster overall convergence.
  • Example: Select 32 random indices (ix = torch.randint(...)) and use X[ix] and Y[ix] for training.

    

    

Finding a Good Initial Learning Rate [00:45:40]

• Learning Rate Tuning:
  • The learning rate (lr) is a crucial hyperparameter. Too small, and training is slow; too large, and the loss might diverge (explode).
  • A common strategy is to sweep through learning rates and plot their effect on the loss.
• Exponentially Spaced Learning Rates:
  • It's often more effective to search for lr on a logarithmic scale.
  • Create a range of learning rate exponents (e.g., lre = torch.linspace(-3, 0, 1000) for 10^-3 to 10^0).
  • The actual learning rates are then 10**lre.
  • Train the model for a few steps with each lr and record the loss.
• Visualization:
  • Plotting lre against loss often reveals a "valley" where the loss decreases steadily before becoming unstable. The bottom of this valley suggests a good learning rate (e.g., 10^-1 = 0.1).

    

• Training with Chosen LR:
  • Start with the identified optimal learning rate (e.g., 0.1) for many iterations.
  • Periodically, decay the learning rate (e.g., multiply by 0.1) in later stages of training to fine-tune the model and prevent oscillations around the minimum.

Splitting Up the Dataset into Train/Val/Test Splits and Why [00:53:20]

• The Problem of Overfitting:
  • A very powerful model might achieve a very low loss on the training data by simply memorizing it (overfitting), without learning generalizable patterns.
  • Such a model performs poorly on unseen data and might only generate data identical to its training set.
• Solution: Data Splits:
  • Training Set (e.g., 80%): Used exclusively for optimizing the model's parameters (weights, biases, embeddings) via gradient descent.
  • Development/Validation Set (e.g., 10%): Used for tuning hyperparameters (e.g., hidden layer size, embedding dimension, learning rate schedule, regularization strength). The best model configurations are chosen based on their performance on this set.
  • Test Set (e.g., 10%): Used only once at the very end to provide a final, unbiased estimate of the model's performance on completely unseen data. Repeated evaluation on the test set can lead to "overfitting" the test set itself.

    

• Implementation:
  • The full dataset of words is shuffled and then split into three partitions (X_train, Y_train), (X_dev, Y_dev), and (X_test, Y_test).
  • Initial training with a small model (100 hidden neurons, 2D embeddings) shows similar train and dev loss (e.g., 2.3 vs 2.38), indicating underfitting (model not powerful enough to fully learn patterns). This suggests increasing model capacity.

    

Experiment: Larger Hidden Layer [01:00:49]

• Hypothesis: Since the model was underfitting, increasing its capacity should improve performance.
• Change: The number of neurons in the hidden layer is increased from 100 to 300.
• Result: After training, the dev loss slightly improves (e.g., from 2.38 to 2.23). This small improvement suggests another bottleneck.

Visualizing the Character Embeddings [01:05:27]

• Current State: Before increasing embedding dimensions, the 2D embeddings (C matrix, 27x2) can be plotted.
• Observations:
  • The model learns meaningful relationships: Vowels (a, e, i, o, u) are clustered together, indicating the model treats them similarly.
  • Unique characters like 'q' and the padding character '.' are located far from others, reflecting their distinct roles.
  • This structure confirms that the embedding layer is learning useful representations.

    

Experiment: Larger Embedding Size [01:07:16]

• Hypothesis: The 2D embedding space might be a bottleneck, limiting the model's ability to represent complex character relationships.
• Change: The embedding dimension is increased from 2 to 10.
  • This also increases the input size to the hidden layer (e.g., 3 * 10 = 30).
  • The hidden layer size is adjusted (e.g., to 200).
• Result: After retraining, the dev loss significantly improves (e.g., from 2.23 to 2.19), confirming that a larger embedding space allows the model to learn better representations.

Summary of Our Final Code, Conclusion [01:11:46]

The final training loop involves 200,000 steps, with an initial learning rate of 0.1 for the first 100,000 steps, then decaying to 0.01 for the remaining 100,000 steps. The best validation loss achieved is 2.17.

• Exercises for the User:
  1. Tune hyperparameters (hidden layer neurons, embedding size, block_size, learning rate schedule, batch size) to beat a validation loss of 2.17.
  2. Experiment with network initialization to get a starting loss closer to uniform probability.
  3. Implement and test ideas from the Bengio et al. [2003] paper.

Sampling from the Model [01:13:24]

• The trained model can be used to generate new character sequences (names).
• Process:
  1. Start with an initial context (e.g., ...).
  2. Feed the context embeddings through the network to get logits.
  3. Apply F.softmax to get probabilities.
  4. Sample the next character's index from the probability distribution using torch.multinomial.
  5. Append the sampled character to the context, sliding the window, and repeat until the '.' token is generated.
• Results: The generated names (e.g., "joes", "ham") sound significantly more "name-like" compared to the bigram model, indicating successful learning.