Large Language Models from Scratch

GitHub

Taking a look under the hood of LLMs, I implement a variety of neural network architecture commonly used in language modeling from scratch. Many of these implementations follow from Andrej Karpathy's Neural Networks: Zero to Hero lecture series. I then dive deeper into representation learning, focusing on topics in model internals and mechanistic interpretability.

micrograd

micrograd includes an implementation of backpropagation and stochastic gradient descent on scalar-valued neural networks. The core class of the implementation is the Value class, which are used to define neurons of the neural network. Value is defined such that basic algebraic operations can be performed to create new Value, while keeping track of these operations in the form of an expression graph. This bookkeeping serves an important purpose for training these neural networks; the record of all the operations which went into a Value is the foundation of the backpropagation algorithm, allowing for the exact analytic calculation of gradients simply through the chain rule from calculus. Computing these gradients exactly rather than estimating them numerically with finite difference methods is key for accuracy and efficiency during training.

A single neuron consits of a data vector $x = (x_{1} x_{2})$ , a weight vector $w = (w_{1} w_{2})$ , a bias $b$ , and an activation function $tanh$ . We can see an example of an expression graph for the following neuron:

tanh (x \cdot w + b)

Figure 1: Single Neuron Expression Graph

micrograd culminates in a small multi-layer perception (MLP) model built and trained entirely in pure python.

makemore

Bigram Model

We begin with language modeling with the bigram model. This model is not built on a neural network architecture but rather serves as a conceptual introduction to language modeling. The goal of makemore is to generate text. It is an autoregressive character level model trained on bigrams. All this means is that the model learns to predict the next character, given the previous character. It looks at pairs of characters (bigrams) and learns a conditional probability distribution in form of a correlation matrix by finding the distribution which maximizes the likelihood of the training data.

We then implement this same bigram model in the form of a neural network. This is an extremely simple network, consisting only of a single linear layer followed by a softmax layer.

Multi-Layer Perception

We then begin to scale up with an implementation of a MLP following Bengio et al., 2003. We embedd characters into 10-dimensional vectors, extend the context window––enabling the model to look further back at the data proceeding each character, and deepen the neural network. We also include the addition of batch normalization layers following Ioffe et al., 2015, resulting in a network of 6 linear layers, each followed by a BatchNorm Layer and $tanh$ activation.

Figure 3: MLP Architecture (Bengio et al.)

Convolutional Neural Network

We then implement features from Google DeepMind's WaveNet CNN (van den Oord et al., 2016) into our MLP. In particular, we include the FlattenConsecutive layer, improving the embeddings of the training data such that information is compressed over time rather than in a single layer.

Figure 4: WaveNet Architecture (van den Oord et al.)

It is worth noting that this implementation is not a CNN but rather draws on this concept from the CNN WaveNet.

Transformer

makemore culminates in an implementation of a decoder-only transformer model, following the famous paper Vaswani et al., 2017 and similar to a mini version of GPT-2.

We implement a self-attention mechanism manually from torch.tensor primitives. The attention mechanism looks to solve the problem that when predicting the next token, certain tokens in this token's context have more predictive influence than others. Attention is a communication mechanism that allows us to take a data-dependent sum of the nodes pointing to a node of a directed graph. This enables the most influential tokens to talk to the prediction of the next token more than the less significant ones.

Figure 5: Encoder-Decoder Transformer Architecture (Vaswani et al.)

This implementation also includes residual connections from He et al., 2015, layer normalization from Ba, Kiros, and Hinton, 2016, and dropout to prevent overfitting Srivastava, 2014