Mathematical Foundations of Large Language Models: Training Objectives, Token Probability, and Loss Functions

• \( w_{
• \( T \): Total length of the sequence.

Example: Input: “Once upon a time, there was a ___.”
The model predicts:

• “princess”: 0.9 (nailed it!)
• “dragon”: 0.05
• “robot”: 0.05

The training process tweaks the model to assign higher probabilities to the right answer.

Key Differences Between MLMs and CLMs

Numerical Example: Causal Language Model

Imagine training a causal language model on the sentence: “AI is amazing.”

Step-by-Step:

1. Tokenize the Input:
   • Tokens: [“AI”, “is”, “amazing”].
2. Predictions:
   • \( P(\text{is} | \text{AI}) = 0.9 \)
   • \( P(\text{amazing} | \text{AI is}) = 0.8 \)
3. Loss: Use the negative log probabilities: \[ \mathcal{L}_{\text{CLM}} = - \left[ \log(0.9) + \log(0.8) \right] \] Approximate: \[ \mathcal{L}_{\text{CLM}} = - (-0.105 + -0.223) = 0.328 \]

The smaller the loss, the better the model is at predicting!

Fun Analogy

Training objectives are like exam prep:

• Masked Language Models (MLMs): A fill-in-the-blank test. The model studies the whole sentence, then guesses the missing piece.
• Causal Language Models (CLMs): A guessing game where the model predicts what comes next, like completing a story.

Mermaid.js Diagram: MLM vs. CLM Workflow

graph TD
    Input[Input Sentence] --> MLM[Masked Language Model]
    Input --> CLM[Causal Language Model]
    MLM --> MaskedTokens[Predict Missing Tokens]
    CLM --> NextToken[Predict Next Token]
    MaskedTokens --> MLM_Loss[Calculate MLM Loss]
    NextToken --> CLM_Loss[Calculate CLM Loss]

2. Token Probability Computation: Math Behind the Magic 🧮✨

How Token Probabilities Are Computed

The journey to token probabilities involves these steps:

1. Input Embedding: Convert tokens into vectors.
2. Contextual Understanding: Apply self-attention to capture relationships between tokens.
3. Logits Calculation: Use linear layers to get raw scores (logits) for each token.
4. Softmax: Transform logits into probabilities.

Let’s break it down! 🛠️

1. Input Embedding: Words Become Numbers

Tokens are converted into dense vectors using an embedding matrix \( E \):

Where:

• \( \mathbf{x} \): One-hot encoded token (e.g., “AI”).
• \( E \): Embedding matrix of size \( V \times d \) (\( V \): vocab size, \( d \): embedding dimension).

2. Contextual Understanding: Attention to Details

The self-attention mechanism learns relationships between tokens. For each token:

Where:

• \( Q \): Query matrix.
• \( K \): Key matrix.
• \( V \): Value matrix.

The attention mechanism ensures the model knows which tokens to focus on. For example:

• In “AI is amazing,” “is” relates most to “AI.”

3. Logits Calculation: Predict Raw Scores

Once contextualized, the hidden state \( \mathbf{h}_T \) (for the last token) is passed through a linear layer:

Where:

• \( W \): Weight matrix.
• \( b \): Bias vector.
• Logits are raw, unnormalized scores for all tokens in the vocabulary.

4. Softmax: Turning Logits into Probabilities

Softmax normalizes logits into probabilities:

import torch
import torch.nn as nn

# Example vocab size and embedding dim
vocab_size = 3
embedding_dim = 2

# Embedding layer
embedding = nn.Embedding(vocab_size, embedding_dim)
tokens = torch.tensor([0, 1])  # Example tokens

# Forward pass
embedded = embedding(tokens)

# Attention (simplified)
Q = K = V = embedded
attention_scores = torch.softmax(torch.matmul(Q, K.T) / (embedding_dim ** 0.5), dim=-1)
context = torch.matmul(attention_scores, V)

# Logits and softmax
linear = nn.Linear(embedding_dim, vocab_size)
logits = linear(context[-1])  # Take the last token's context
probs = torch.softmax(logits, dim=-1)

print("Logits:", logits)
print("Probabilities:", probs)

graph TD
    Token[Input Token] --> Embedding[Embedding Layer: Convert to Vectors]
    Embedding --> Attention[Self-Attention: Contextualize Tokens]
    Attention --> Logits[Linear Layer: Compute Logits]
    Logits --> Softmax[Apply Softmax]
    Softmax --> Probability[Token Probabilities]

import torch
import torch.nn as nn

# Example token predictions (logits) and true labels
logits = torch.tensor([[2.0, 1.0, 0.1], [0.5, 2.5, 0.2]])  # Shape: (batch_size=2, vocab_size=3)
true_labels = torch.tensor([0, 1])  # Correct tokens for each example

# Cross-Entropy Loss
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, true_labels)

print("Cross-Entropy Loss:", loss.item())

graph TD
    Input[Input Tokens] --> Embedding[Embedding Layer]
    Embedding --> Attention[Self-Attention: Contextualize Tokens]
    Attention --> Logits[Linear Layer: Compute Logits]
    Logits --> Softmax[Softmax: Convert Logits to Probabilities]
    Softmax --> LossFunction[Cross-Entropy Loss]
    LossFunction --> Backprop[Backpropagation to Update Model]