20.2. Implementation

Alpha Version: Work in progress.

This section describes the architectural components specific to GPT-1—namely, its position encoding mechanism and the language modeling loss function. It then walks through how to execute the pre-training scripts.

20.2.1. Embedding and Position Encoding

In the original Transformer architecture, the token embedding matrix is learnable, whereas the positional encoding relies on a fixed, predefined sinusoidal function.

In contrast, GPT-1 employs fully learnable weights for both its token embedding and its position encoding components.

Let $X_{pe}$ represent the final embedded input sequence fed into the first decoder unit. It is mathematically formulated as follows:

where:

• $X_{pe} \in \mathbb{R}^{N \times d_{\text{model}}}$ is the input matrix to the first decoder unit.
• $X \in \mathbb{R}^{N \times \vert{}\text{Vocabulary}\vert{}}$ represents the sequence of one-hot encoded input tokens.
• $W_{e} \in \mathbb{R}^{\vert{}\text{Vocabulary}\vert{} \times d_{\text{model}}}$ is the learnable token embedding matrix.
• $W_{p} \in \mathbb{R}^{N \times d_{\text{model}}}$ is the learnable position embedding matrix.

20.2.1.1. Implementation

The implementation details of GPT-1’s embedding layers are shown below. Both the token embedding and the positional embedding are instantiated using tf.keras.layers.Embedding().

By configuring the embeddings_initializer and assigning explicit name attributes, we define them as trainable parameters within the network.

# 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.2.2. Loss Function

This section describes the loss function optimized during the pre-training phase.

Following the original GPT-1 paper, the pre-training objective maximizes the log-likelihood of the language model.

To simplify the formal framework, let $h_{0}$ denote the combined input matrix $X_{pe}$ obtained from the token and position embeddings. Under this convention, the forward pass through the Transformer decoder blocks and the final output projection layer is defined as follows:

where:

• $Z \in \mathbb{R}^{N \times \vert{}\text{Vocabulary}\vert{}}$ is the matrix of logits produced by projecting the final decoder representations $h_{L}$ back onto the vocabulary space. (Note that GPT-1 binds or ties the weights of the embedding and output projection layers via $W_{e}^{T}$).
• $P \in \mathbb{R}^{N \times \vert{}\text{Vocabulary}\vert{}}$ is the matrix of predicted probability distributions over the vocabulary, obtained by applying the softmax function to each row of $Z$.

The $t$-th row of $P$ directly corresponds to the conditional probability distribution $P(X_{t} \mid X_{<t})$, which evaluates the likelihood of the next token at position $t$.

20.2.2.1. Implementation

Because TensorFlow optimizers are fundamentally designed to minimize a loss function rather than maximize an objective, our implementation minimizes the negative log-likelihood (NLL). This is mathematically equivalent to maximizing the log-likelihood formulation shown above.

In practice, this negative log-likelihood is calculated using tf.keras.losses.SparseCategoricalCrossentropy().

# ============================================================
# 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.2.3. Running Pre-Training

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

20.2.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.2.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.