Language Model Cross Entropy Loss Loss is calculated as the average negative log probs of the correct tokens. https://arxiv.org/pdf/1706.037

Attention Calculation. The attention calculation does two things for each destination token - it both moves information from source t

Evaluate the model on a test dataloader Args: model (Transformer): Transformer model test_dataloader (DataLoader): Test dataloade

Get the default device to use. Returns: torch.device: Device to use.

Mask the attention pattern. Each attention pattern is of shape (dest, src), and each element represents the attention that a destinat

Train loop Args: model: The Transformer model to train. train_dataloader: Dataloader for training data. test_dataloader:

Initialize the Embed layer.

Get the tokenized datasets.

Train.

Initialise the Transformer.

Initialize the positional encoding matrix.

Initialize the MultiHeadAttention module.

MLP Sub-Layer Initialisation.

Initialise the full layer.

Initialize the Unembed Layer.

Initialise the dataset. Args: num_samples (int): Number of samples d_vocab (int): Vocab size (e.g. if 3 then it'll ge

Initialize the dataset. Args: num_samples (int): Number of samples sequence_length (int): Sequence length

Get the number of samples in the dataset. Returns: int: Number of samples in the dataset

Forward pass. Args: tokens (BatchTokenIndices): Input tokens (indices rather than one-hot) Returns: BatchLog

Apply the positional encoding to the given input embedding. Args: embedding (ResidualStream): The input embedding with shape (bat

Attention layer forward pass. https://arxiv.org/pdf/1706.03762.pdf (p5)

Forward Pass through the MLP Sub-Layer. Args: residual_stream (ResidualStream): MLP input Returns: ResidualS

Forward pass. Args: residual_stream (ResidualStream): Residual stream Returns: ResidualStream: Updated resid

Forward Pass through the Embedding Layer. The original paper multiples the embedding by sqrt(d_model) during the forward pass, presum

Forward Pass through the Unembedding Layer. Args: residual_stream (ResidualStream): Residual stream Returns:

Forward Pass.

Forward pass Args: x (BatchTokenIndices): Inputs Returns: BatchLogits: 100% log probabilities of the specifi

Test that the model can attend to a specific (flagged) token.

Test a simple attention calculation

Verify that we don't have any repeated enum elements.

Test the MLP network learns to solve the regression tasks. Args: dataset_class (RegressionTaskDataset): Dataset to test.

Test that a model that always outputs the correct token index, gets 100% accuracy.

Test that a model that always outputs the incorrect token index, gets 100% accuracy.

Test that a valid PyTorch device is returned.

Test if Layer correctly adds attention and feed forward outputs to the residual stream.

Test that a model with just the embedding and unembedding can learn bigram statistics.

Test that relative positions can be obtained with just matrix multiplication. Check that the positional encoding of a token at position `initial_

Test loss with random logits.

Test loss with uniform logits.

Test that it masks correctly

Test that model parameters change after training.

Test that the encoding is numerically stable. Check our Positional Encoding doesn't create any NaN values.

Check the positional encoding is as expected for a specific position and dimension. This uses a scalar calculation (as opposed to the vectorized

Test that each token will have a unique positional encoding vector.

Test that the positional encoding is the same for each batch item.

Test that the train loop runs without error.

Check that a forward pass can be run.

Check that the Transformer is initialised correctly.

Check the transformer can learn to memorize a simple dataset.

Test loss where the logits are fully accurate (so it should be 0).

Tokenize code.