20.3. Implementation

Alpha Version: Work in progress.

This section details the implementation of components specific to GPT-1, namely its positional encoding mechanism and the language modeling loss function. It then walks through executing the pre-training scripts.

Section Contents

20.3.1. Embedding and Position Encoding
20.3.2. Loss Function
20.3.3. Running Pre-Training

20.3.1. Embedding and Position Encoding

The code below demonstrates how GPT-1’s embedding layers are implemented. Both the token embedding and positional embedding layers are instantiated using tf.keras.layers.Embedding().

By configuring embeddings_initializer and assigning explicit name attributes, both components are registered as trainable parameters within the model.

# GPT-1 Embedding Layers

class GPT1Model(tf.keras.Model):
    def __init__(self, vocab_size, d_model, max_seq_len, ...):
        super().__init__()
        _init = tf.keras.initializers.TruncatedNormal(stddev=0.02)

        # Token embedding (We)
        self.token_embedding = tf.keras.layers.Embedding(
            vocab_size,
            d_model,
            embeddings_initializer=_init,
            name="token_embedding",
        )

        # Positional embedding (Wp)
        self.position_embedding = tf.keras.layers.Embedding(
            max_seq_len,
            d_model,
            embeddings_initializer=_init,
            name="position_embedding",
        )

    def call(self, input_ids, training=False, mask=None):
        # ... snip ...
        seq_len = tf.shape(input_ids)[1]
        x = self.token_embedding(input_ids)
        positions = tf.range(seq_len)
        x += self.position_embedding(positions)
        # ... snip ...

20.3.2. Loss Function

This implementation leverages tf.keras.losses.SparseCategoricalCrossentropy() to calculate the negative log-likelihood loss for language modeling.

# ============================================================
# Loss function - causal language-model cross-entropy
# ============================================================

_lm_loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction="none"
)


def lm_loss_fn(targets, logits):
    """
    Language-model cross-entropy, ignoring padding positions (token-id == 0).

    targets : (B, T)     integer next-token ids;  0 = pad -> excluded
    logits  : (B, T, V)  raw predicted logits
    Returns : scalar mean loss
    """
    loss = _lm_loss_object(targets, logits)  # (B, T)
    mask = tf.cast(tf.not_equal(targets, 0), tf.float32)
    loss = loss * mask
    return tf.reduce_sum(loss) / (tf.reduce_sum(mask) + 1e-9)

20.3.3. Running Pre-Training

This subsection explains how to prepare the SNLI dataset and run the GPT-1 pre-training pipeline.

20.3.3.1. Prepare dataset

Run the following script to create the local ./data directory, download the SNLI source dataset, and execute the preprocessing pipeline required for pre-training:

$ python prepare_datasets.py

20.3.3.2. Pre-Training

To start training the model on the preprocessed SNLI corpus, execute the gpt1_pretrain.py script.

The trailing argument specifies the number of training epochs (e.g., 50 epochs).

$ python gpt1_pretrain.py 50

Depending on your hardware acceleration setup, this training phase may take several hours to complete.