Homepage

• Tech Blog
• Suibi
• Search

✕

From Math to Practices

A sharing session at Stanford XCS 236 (Deep Generative Models)

By Cheng Shen
Posted on October 18, 2024

A knowledge sharing session at Stanford XCS 236 (Deep Generative Models) where I shared how math concepts from Stanford Course XCS 236 are related to the Transformer and Large Language Models.

Agenda

0. Why this session?

1. Model architecture: Autoregressive model architecture in language modeling

1. Recap on conditional probability
2. Dependency assumption in n-gram model
3. Transformer
4. Dependency assumption in BERT (encoder-only model)
5. Dependency assumption in LLM (decoder-only model, example running with Mistral)

2. Model pretraining

1. Maximum likelihood objective
2. Transferable knowledge from pretraining

3. Supervised fine-tuning

1. Same maximum likelihood objective, but different data distribution
2. Multi-modal model, and its assumptions of likelihoods (example running with Pixtral)

4. Contrastive learning / RLHF

1. PPO
2. Direct Preference Optimization method and theory

0. Why this session?

• Course content from XCS236 Deep Generative Models (and others)
  • Maximum likelihood
  • Probability
  • Distribution, KL-divergence
  • ...
• boom!
• Fancy and exciting machine learning models
  • Transformer
  • GPT
  • RLHF
• How are mathematical concepts related to the latest structures and techniques?

1.1 Model architecture -- Recap on conditional probability

• Large Language Model is still a "language model"
• A language model estimates the distribution of words, or sequence of words.

• With the very straightforward dependency assumption, and leveraging conditional probability, we have $$p_\theta(x) = p_\theta(x_1) \cdot p_\theta(x_2|x_1) \cdot p_\theta(x_3|x_1, x_2)...$$

• For example, given input sentence I love machine learning, $$p_\theta(\text{``I love machine learning"}) = p_\theta(\text{``I"}) \cdot p_\theta(\text{``love"}|\text{``I"}) \cdot p_\theta(\text{``machine"}|\text{``I love"}) \cdot p_\theta(\text{``learning"}|\text{``I love machine"})$$

• $\theta$ is the parameters used to model the distribution
  • $\theta$ could be treated as the parameters in a neural network. But it could also be something else...
  • In n-gram model (markov assumption), $\theta$ could be as simple as a look-up table

1.2 n-gram model

• Without loss of generality, we take bigram model here as an example. But n could be any number
• Markove assumption for bigram model -- each token depends only one token in the past to make the prediction $$p_\theta(x) = p_\theta(x_1) \cdot p_\theta(x_2|x_1) \cdot p_\theta(x_3|x_2)...$$
• $\theta$ could be a python dictionary that looks like below

from typing import List, Tuple
from collections import defaultdict
import math
import random

class NGramModel:
    def __init__(self, n: int):
        if n < 1:
            raise ValueError("n must be at least 1")
        self.n = n
        self.prob_table = defaultdict(lambda: defaultdict(float))
        self.ngram_counts = defaultdict(int)
        self.context_counts = defaultdict(int)
        self.vocabulary = set()

    def train(self, training_corpus: List[List[str]]):
        # Count n-grams
        for sentence in training_corpus:
            sentence = ['<START>'] * (self.n - 1) + sentence + ['<END>']
            for i in range(len(sentence) - self.n + 1):
                ngram = tuple(sentence[i:i+self.n])
                context = ngram[:-1]
                word = ngram[-1]

                self.ngram_counts[ngram] += 1
                self.context_counts[context] += 1
                self.vocabulary.add(word)

        # Calculate probabilities
        for ngram, count in self.ngram_counts.items():
            context = ngram[:-1]
            word = ngram[-1]
            self.prob_table[context][word] = count / self.context_counts[context]

    def generate(self, start_words: List[str], num_words: int) -> List[str]:
        context = tuple(['<START>'] * (self.n - 1 - len(start_words)) + start_words)
        generated = list(context)

        for _ in range(num_words):
            if context not in self.prob_table:
                break

            word = self._sample_next_word(context)
            generated.append(word)

            if word == '<END>':
                break

            context = tuple(generated[-self.n+1:])

        return [word for word in generated if word not in ['<START>', '<END>']]

    def _sample_next_word(self, context: Tuple[str, ...]) -> str:
        choices, weights = zip(*self.prob_table[context].items())
        return random.choices(choices, weights=weights)[0]

    def probability(self, word: str, context: Tuple[str, ...]) -> float:
        if context not in self.prob_table or word not in self.prob_table[context]:
            return 0.0
        return self.prob_table[context][word]

    def perplexity(self, test_corpus: List[List[str]]) -> float:
        log_probability = 0
        total_words = 0

        for sentence in test_corpus:
            sentence = ['<START>'] * (self.n - 1) + sentence + ['<END>']
            for i in range(self.n - 1, len(sentence)):
                context = tuple(sentence[i-self.n+1:i])
                word = sentence[i]
                prob = self.probability(word, context)
                if prob > 0:
                    log_probability += math.log(prob)
                total_words += 1

        return math.exp(-log_probability / total_words)

import nltk
from nltk.corpus import brown
from typing import List

# Download the Brown corpus if you haven't already
nltk.download('brown')

def preprocess_brown_corpus() -> List[List[str]]:
    # Get all sentences from the Brown corpus
    sentences = brown.sents()

    # Lowercase all words and remove sentences with less than 2 words
    processed_sentences = [
        [word.lower() for word in sentence]
        for sentence in sentences
        if len(sentence) > 1
    ]

    return processed_sentences

# Preprocess the Brown corpus
training_data = preprocess_brown_corpus()

# Create and train the bigram model
bigram_model = NGramModel(n=2)
bigram_model.train(training_data)

# Test the model
start_words = ["the"]
generated_sequence = bigram_model.generate(start_words, 10)
print(f"Generated sequence: {' '.join(generated_sequence)}")

# Calculate perplexity on a subset of the corpus (e.g., last 1000 sentences)
test_data = training_data[-1000:]
perplexity = bigram_model.perplexity(test_data)
print(f"Perplexity on test data: {perplexity}")

# Example of probability lookup
print(f"Probability of 'dog' given 'the': {bigram_model.probability('dog', ('the',))}")
print(f"Probability of 'cat' given 'the': {bigram_model.probability('cat', ('the',))}")

# Create and train the tri-gram model
trigram_model = NGramModel(n=3)
trigram_model.train(training_data)

# Test the model
start_words = ["the"]
generated_sequence = trigram_model.generate(start_words, 10)
print(f"Generated sequence: {' '.join(generated_sequence)}")

# Calculate perplexity on a subset of the corpus (e.g., last 1000 sentences)
test_data = training_data[-1000:]
perplexity = trigram_model.perplexity(test_data)
print(f"Perplexity on test data: {perplexity}")

# Example of probability lookup, the trigram model here does not have any fall-back
# strategy. Therefore the probability could be 0
print(f"Probability of 'dog' given 'the cute': {trigram_model.probability('dog', ('the', 'cute',))}")
print(f"Probability of 'cat' given 'the cute': {trigram_model.probability('cat', ('the', 'cute',))}")

1.3 Transformer

• Much more powereful than n-gram model
• (Theoretically) no Markov assumption anymore!
• Recap bi-gram model $$p_\theta(x) = p_\theta(x_1) \cdot p_\theta(x_2|x_1) \cdot p_\theta(x_3|x_2)...$$ where $\theta$ is a probability look-up table as python dictionary
• Transformer takes all past inputs into consideration, so the given likelihood is a full expansion of conditional probability $$p_\theta(x) = p_\theta(x_1) \cdot p_\theta(x_2|x_1) \cdot p_\theta(x_3|x_1, x_2)...$$
• Wait!
  • Doesn't it mean $\theta$ needs to have enough parameters to cover an indefinite number of combinations $(x_1, x_2, ... x_n)$?
  • And I don't care about possibilities, I just want my LLM to generate next tokens, how does it relate to probability distribution $p_\theta(x)$?

tl;dr: $\theta$ is not the Transformer (NN) parameters. Transformer (NN) is not generating next token, but generating the distribution for next token.

• What happens is that:
  1. Let $(x_1, x_2, ..., x_n)$ be the input sequence, $f_{enc}$ is the Transformer encoder, $f_{dec}$ is the Transformer decoder
  2. The encoder produces "features" for each input token $(enc_1, enc_2, ..., enc_n) = f_{enc}(x_1, x_2, ..., x_n)$
  3. The decoder takes encoder outputs, and autoregressively produces distribution for next token $$\theta_i = f_{dec}((enc_1, enc_2, ..., enc_n), (y_1, y_2, ..., y_{i-1}))$$ And $\theta_i$ is a probabiliy distribution for the possible tokens on position $i$, something like $$p_{\theta_i}(y_i=\text{'I'}) = 0.03$$ $$p_{\theta_i}(y_i=\text{'cat'}) = 0.07$$ $$...$$
• Therefore
  • Transformer produces $\theta$ as inference result -- thus parameter space is reduced
  • Sampling techniques like temperature, beam-search also makes sense, because Transformer output is not a deterministic token, but a probability distribution $\theta$ that could be sampled from.

import torch
import torch.nn as nn
import torch.optim as optim
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(0), :]

class NaiveTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers):
        super(NaiveTransformer, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src = self.embedding(src) * math.sqrt(self.embedding.embedding_dim)
        tgt = self.embedding(tgt) * math.sqrt(self.embedding.embedding_dim)
        src = self.pos_encoder(src)
        tgt = self.pos_encoder(tgt)
        output = self.transformer(src, tgt)
        return self.fc_out(output)

# Demo data and training
vocab_size = 1000
d_model = 512
nhead = 8
num_encoder_layers = 3
num_decoder_layers = 3
batch_size = 32
seq_length = 20

# Create random input data
src = torch.randint(1, vocab_size, (seq_length, batch_size))
tgt = torch.randint(1, vocab_size, (seq_length, batch_size))

# Initialize the model
model = NaiveTransformer(vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    output = model(src, tgt[:-1])
    loss = criterion(output.view(-1, vocab_size), tgt[1:].reshape(-1))
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}")

# Inference
model.eval()
with torch.no_grad():
    input_seq = torch.randint(1, vocab_size, (seq_length, 1))
    output_seq = torch.zeros((seq_length, 1), dtype=torch.long)
    for i in range(seq_length):
        out = model(input_seq, output_seq)
        next_token = out[-1].argmax(dim=-1)
        output_seq[i] = next_token
    print("Input sequence:", input_seq.squeeze().tolist())
    print("Generated sequence:", output_seq.squeeze().tolist())

1.4 Independence Assumption in BERT

Proposed in https://arxiv.org/pdf/1810.04805

• (BERT) Bidirectional Encoder Representations from Transformers
• BERT is one big thing before GPT came out, and is still widely applied in many applications
• As suggested by its name "Bidirectional", BERT does not have autoregressive nature (and is not a generative model at all)
• Instead it was pre-trained with 2 tasks

Pretraining Task 1 -- Masked LM

• Train the model to predict the masked token given context before and after the masked token $$p(x_i | x_1, x_2, ..., x_{i-1}, MASKED, x_{i+1}, ..., x_{n})$$

Pretraining Task 2 -- Next Sentence Prediction

• Train the model to predicts whether two concatenated sentences are the true continuation. $$p(NSP | x_1, x_2, ..., x_n, SEP, y_1, y_2, ..., y_m)$$

• BERT is not trained for generation
• It has limited capability of language modeling as well
• Nevertheless, pretrained BERT model learned to understand the language, and has been used in many tasks (NER, question answering) effectively

!pip install datasets

import torch
from torch.utils.data import DataLoader, RandomSampler, TensorDataset
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import numpy as np
from datasets import load_dataset

# Load IMDB dataset
dataset = load_dataset("imdb", split="train[:5000]")  # Using only 5000 samples for demonstration

# Prepare data
texts = dataset["text"]
labels = dataset["label"]

# Split the data
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)

# Load pre-trained BERT tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Tokenize and encode sequences
MAX_LENGTH = 256

def encode_texts(texts):
    return tokenizer.batch_encode_plus(
        texts,
        max_length=MAX_LENGTH,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )

train_encodings = encode_texts(train_texts)
val_encodings = encode_texts(val_texts)

# Create torch datasets
train_dataset = TensorDataset(
    train_encodings['input_ids'],
    train_encodings['attention_mask'],
    torch.tensor(train_labels)
)
val_dataset = TensorDataset(
    val_encodings['input_ids'],
    val_encodings['attention_mask'],
    torch.tensor(val_labels)
)

# Create data loaders
train_loader = DataLoader(train_dataset, sampler=RandomSampler(train_dataset), batch_size=16)
val_loader = DataLoader(val_dataset, batch_size=32)

# Load pre-trained BERT model
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

# Set up optimizer
optimizer = AdamW(model.parameters(), lr=2e-5)

# Training loop
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

num_epochs = 3

for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch[0].to(device)
        attention_mask = batch[1].to(device)
        labels = batch[2].to(device)
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

    # Validation
    model.eval()
    val_preds = []
    val_true = []
    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch[0].to(device)
            attention_mask = batch[1].to(device)
            labels = batch[2].to(device)
            outputs = model(input_ids, attention_mask=attention_mask)
            preds = torch.argmax(outputs.logits, dim=1)
            val_preds.extend(preds.cpu().numpy())
            val_true.extend(labels.cpu().numpy())

    accuracy = accuracy_score(val_true, val_preds)
    print(f"Epoch {epoch+1}/{num_epochs}, Validation Accuracy: {accuracy:.4f}")

# Final evaluation
print("\nClassification Report:")
print(classification_report(val_true, val_preds, target_names=['Negative', 'Positive']))

# Example inference
def predict_sentiment(text):
    encoding = tokenizer.encode_plus(
        text,
        max_length=MAX_LENGTH,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)
        pred = torch.argmax(outputs.logits, dim=1)

    return "Positive" if pred.item() == 1 else "Negative"

# Test the model
test_text = "This movie was absolutely fantastic! I loved every minute of it."
print(f"\nSample text: '{test_text}'")
print(f"Predicted sentiment: {predict_sentiment(test_text)}")

1.5 Decoder only LLM (e.g. GPT, Mistral)

decoder-only architecture, deviating from traditional encoder-decoder transformer models

• Input data is fed directly into the decoder without being transformed by an encoder
• Uses masked self-attention during training

Pretraining Task -- Next Token Prediction

• Recap -- for encoder-decoder Transformer

  1. Let $(x_1, x_2, ..., x_n)$ be the input sequence, $f_{enc}$ is the Transformer encoder, $f_{dec}$ is the Transformer decoder
  2. The encoder produces "features" for each input token $(enc_1, enc_2, ..., enc_n) = f_{enc}(x_1, x_2, ..., x_n)$
  3. The decoder takes encoder outputs, and autoregressively produces distribution for next token $$\theta_i = f_{dec}((enc_1, enc_2, ..., enc_n), (y_1, y_2, ..., y_{i-1}))$$ And $\theta_i$ is a probabiliy distribution for the possible tokens on position $i$, something like $$p_{\theta_i}(y_i=\text{'I'}) = 0.03$$ $$p_{\theta_i}(y_i=\text{'cat'}) = 0.07$$ $$...$$

• For decoder-only Transformer

  1. Let $(x_1, x_2, ..., x_i)$ be the input prefix, $f_{dec}$ is the Transformer decoder.
  2. The decoder takes the whole prefix, and autoregressively produces distribution for next token $$\theta_i = f_{dec}((x_1, x_2, ..., x_i))$$ And $\theta_i$ is a probabiliy distribution for the possible tokens on position $i$, something like $$p_{\theta_i}(y_i=\text{'I'}) = 0.03$$ $$p_{\theta_i}(y_i=\text{'cat'}) = 0.07$$ $$...$$

• NOTE All the attention and secret sauces are in $f_{dec}$ used to produce $\theta_i$

• The task is autoregressive - the model predicts future tokens based on past tokens. It's trained to generate text in a left-to-right manner.

!pip install datasets

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login

# Authenticate with Hugging Face
# Replace 'your_huggingface_token' with your actual token, or add to colab secrets HF_TOKEN
# login(token=open('./hf_token', 'r').read().strip())

# Load the Llama-2 1B model
model_name = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

# Prepare input
# prompt = "Translate the following English text to French: 'Hello, how are you?'\nAnswer: "

prompt = "Translate the following English text to French: 'Hello, how are you?'"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# Generate
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50)

# Decode and print the result
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(result)

2.1 Model pretraining

Large Language Models (LLMs) like GPT (Generative Pre-trained Transformer) are initially trained on vast amounts of unlabeled text data. This process, known as pretraining, allows the model to learn general patterns and structures of language before being fine-tuned for specific tasks.

The primary goal of pretraining is to build a model that can effectively capture the underlying distribution of natural language. This is achieved through self-supervised learning, where the model learns to predict parts of the input data given other parts, without the need for explicit labels.

1. Maximum likelihood objective

If we denote our training corpus as a sequence of tokens $X=(x_1,x_2,...,x_n)$. The goal is to learn a probability distribution p(X) that maximizes the likelihood of observing this sequence.

For an autoregressive model like GPT, we can factorize this joint probability as:

$$p(X) = p(x_1)\cdot p(x_2|x_1)\cdot p(x_3|x_1,x_2)\cdot ...\cdot p(x_n|x_1,...,x_{n-1})$$

The log-likelihood objective for a single sequence is:

$$L(\theta) = \log p(X; \theta) = \sum_i \log p(x_i | x_1, ..., x_{i-1})$$

where θ represents the model parameters. The overall objective is to maximize this log-likelihood across all sequences in the training data:

$$\theta^* = \text{argmax} \sum_r L(\theta; X^r)$$

2.2 Transferable knowledge from pretraining

Interesting research from Allen-Zeyuan Zhu: https://physics.allen-zhu.com/home

Does LLM actually understand and learn the knowledge, or does it simply memorize the most probable next token?

Experiments on GPT-2 using Bio and Q&A dataset.

• Result 1 A model mixed-trained on both knowledge and its extraction QA tasks can effectively learn to extract knowledge.
• Result 2 A model pretrained to word-by-word memorize knowledge may never be fine-tuned to extract knowledge.
• Result 3 Adding multiplicity, permutations, or repeating full names, all help the model to better store knowledge during pretraining, making knowledge extraction easier later
  • Pretraining on a dataset where each person has 5 diverse biography entries (i.e., different wording, sentence shuffling) boosts the QA fine-tune accuracy from 9.7% to 96.6%
• Result 4 While bidirectional models like BERT is less sensitive to knowledge ordering, the MLM pretraining task does not necessarily promote knowledge storage for subsequent extrac- tion. Unless the knowledge is a standalone word or of independent words (like month, day, year), extracting knowledge after MLM pretraining might prove challenging, if not totally impossible.

3.1 Supervised fine-tuning

The primary goal of SFT is to adapt a pretrained model to perform specific tasks or to align its outputs with desired characteristics.

Here are some example goals of SFT:

1. Task-Specific Adaptation

   • Purpose: To optimize a general-purpose LLM for a particular task or domain.
   • Example: Fine-tuning BERT on a dataset of labeled movie reviews to classify sentiments accurately.

2. Alignment with Human Preferences

   • Purpose: To make model outputs more natural, helpful, and aligned with human expectations.
   • Example: InstructGPT - fine-tune the model to mimic these high-quality human responses.

3. Improving Factual Accuracy

   • Purpose: To reduce hallucinations and improve the model's factual reliability.

4. Multimodal Integration

   • Purpose: To enable models to process and generate content across different modalities.
   • Example: CLIP - understand relationships between images and text, enabling tasks like image search using natural language queries.

5. Stylistic Adaptation

   • Purpose: To tailor the model's output style for specific applications.

3.2 Multi-modal model

Consider a multi-modal model dealing with images (I) and text (T):

1. Joint Likelihood $$P(I, T) = P(I) * P(T|I)$$

2. Conditional Likelihood (for image captioning): $$P(T|I) = \prod_i P(t_i|t_1, ..., t_{i-1}, I)$$ Where tᵢ represents the i-th token in the text, conditioned on the previous tokens and the image.

3. Latent Space Modeling, $z$ represents the shared latent space: $$P(I, T) = ∫ P(I|z) ⋅ P(T|z) * P(z) dz$$

4. Fine-Tuning Objective, the negative log-likelihood is often used as the loss function during fine-tuning: $$L = -\log P(T|I; θ)$$

!pip install ftfy regex tqdm
!pip install git+https://github.com/openai/CLIP.git

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
from tqdm import tqdm
import clip

# Set up device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load CLIP model
clip_model, preprocess = clip.load("ViT-B/32", device=device)

# CIFAR-10 classes
cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

# Load CIFAR-10 dataset
train_data = datasets.CIFAR10(root='./data', train=True, download=True, transform=preprocess)
test_data = datasets.CIFAR10(root='./data', train=False, download=True, transform=preprocess)

# Use a subset of the data
train_subset = Subset(train_data, range(10000))  # 10,000 training samples
test_subset = Subset(test_data, range(1000))    # 1,000 test samples

# Create DataLoaders
train_loader = DataLoader(train_subset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_subset, batch_size=64, shuffle=False)

class EnhancedCLIPClassifier(nn.Module):
    def __init__(self, base_model, num_classes):
        super(EnhancedCLIPClassifier, self).__init__()
        self.base_model = base_model
        self.feature_adapter = nn.Sequential(
            nn.Linear(base_model.visual.output_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        with torch.no_grad():
            features = self.base_model.encode_image(x)
        return self.feature_adapter(features.float())

# Initialize the model
enhanced_classifier = EnhancedCLIPClassifier(clip_model, len(cifar10_classes)).to(device)

# Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(enhanced_classifier.feature_adapter.parameters(), lr=3e-4, weight_decay=0.01)

# Training function
def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    progress_bar = tqdm(loader, desc="Training")

    for images, labels in progress_bar:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

        progress_bar.set_postfix({'Loss': running_loss/len(loader), 'Acc': 100.*correct/total})

    return running_loss/len(loader), 100.*correct/total

# Evaluation function
def evaluate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    return running_loss/len(loader), 100.*correct/total

# Training loop
num_epochs = 3
for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(enhanced_classifier, train_loader, optimizer, criterion, device)
    val_loss, val_acc = evaluate(enhanced_classifier, test_loader, criterion, device)

    print(f'Epoch {epoch+1}/{num_epochs}:')
    print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')
    print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
    print('-' * 50)

print("Training completed!")

# Test the model on a few examples
enhanced_classifier.eval()
test_loader = DataLoader(test_subset, batch_size=5, shuffle=True)
images, labels = next(iter(test_loader))
images, labels = images.to(device), labels.to(device)

with torch.no_grad():
    outputs = enhanced_classifier(images)
    _, predicted = outputs.max(1)

# Print results
for i in range(5):
    print(f"True: {cifar10_classes[labels[i]]}, Predicted: {cifar10_classes[predicted[i]]}")

4.1 RLHF

Proximal Policy Optimization (PPO). For example GPT-4 (https://arxiv.org/pdf/2303.08774)

Supervised fine-tuning (instruction tuning) has its limitations:

• Data collection for generative modeling is expensive -- imagine asking vendors to provide data answering questions like "Please implement PPO algorithm for me"
• Misalignment between Maximum Likelihood optimization v.s. Human Preference

On the other hand:

• It is cheaper and easier to identify "which one is the better answer", for both human, and models
• Leveraging human preference so that training objective is aligned!

Instead of optimizing maximum likelihood, we optimize our policy to get better rewards. But how do we evaluate rewards $R$?

An example -- Bradley Terry model, where we train a reward model $r_\phi$ based on human preference data, to maximum the following likelihood $$p^*(y_1 \succ y_2 | x) = \frac{\exp(r^*(x, y_1))}{\exp(r^*(x, y_1)) + \exp(r^*(x, y_2))} = \sigma(r^*(x, y_2) - r^*(x, y_1))$$

Like usual, the Negative Log-Likelihood loss looks like this $$\mathcal{L}_R(r_\phi, \mathcal{D}) = -\mathbb{E}_{(x,y_1,y_2)\sim\mathcal{D}}\left[\log \sigma(r_\phi(x, y_1) - r_\phi(x, y_2))\right]$$

Then, with a well-trained reward model! We only need to optimize the large language model $\pi$ to maximize the rewards! $$\max_{\pi_\theta} \mathbb{E}_{x\sim\mathcal{D},y\sim\pi_\theta(y|x)}\left[r_\phi(x,y)\right]$$

with a little problem -- sampling operation $y ~ \pi_\theta(y|x)$ from LLM is not differentiable

Time for some RL!

About Reinforcement Learning

• Been used for action planning -- where actions are usually discrete, just like language model outputs
• Last time when it was really a buzz word: AlphaGo, AlphaZero, etc.
• Basic concepts: state, action, reward, and policy

$$\pi_\theta(s) \longrightarrow a $$

The probability of any "trajectory" under policy $\pi$ is given by: $$ P(\tau|\pi) = \rho_0(s_0) \prod_{t=0}^{T-1} P(s_{t+1}|s_t, a_t)\pi(a_t|s_t).$$ The expected value of rewards over policy $\pi$ is $$J(\pi) = \int_\tau P(\tau|\pi)R(\tau) = \mathbb{E}_{\tau\sim\pi} [R(\tau)]$$ And we want to find the optimal policy such that $$\pi^* = \arg\max_\pi J(\pi)$$

In language modeling

• Policy $\pi$ is the large language model
• Current state $s$ is the prefix/prompt -- $x$
• Action $a$ is the response -- $y$

Proximal Policy Optimization

• Used in GPT-4 Training
• Leverage KL-divergence, or gradient clipping to prevent the updated model $\pi$ from getting too far from the starting point
• Action sampling -> Estimation with reward model -> Backpropagate, and determine gradient scale/sign based on reward
• Objective function $$\max_{\pi_\theta} \mathbb{E}_{x\sim\mathcal{D},y\sim\pi_\theta(y|x)}\left[r_\phi(x,y)\right] - \beta\mathbb{D}_\text{KL}\left[\pi_\theta(y|x) \parallel \pi_\text{ref}(y|x)\right]$$

Some limitations

• Reward model needs to be trained separately, using human feedback
• Complicated, hard to tune and run
  • Multiple iterations of model training -> sampling -> reward model update -> model training
  • Within each round of model training: sampling -> grading -> model updates

!pip install gymnasium

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import gymnasium as gym

class PPO:
    def __init__(self, state_dim, action_dim, lr, gamma, epsilon, value_coef, entropy_coef):
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.Tanh(),
            nn.Linear(64, action_dim),
            nn.Softmax(dim=-1)
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )
        self.optimizer = optim.Adam(list(self.actor.parameters()) + list(self.critic.parameters()), lr=lr)
        self.gamma = gamma
        self.epsilon = epsilon
        self.value_coef = value_coef
        self.entropy_coef = entropy_coef

    def get_action(self, state):
        state = torch.FloatTensor(state)
        probs = self.actor(state)
        dist = Categorical(probs)
        action = dist.sample()
        return action.item(), dist.log_prob(action)

    def update(self, states, actions, old_log_probs, rewards, dones):
        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions)
        old_log_probs = torch.FloatTensor(old_log_probs)
        rewards = torch.FloatTensor(rewards)
        dones = torch.FloatTensor(dones)

        # Compute advantages
        values = self.critic(states).squeeze()
        returns = torch.zeros_like(rewards)
        advantages = torch.zeros_like(rewards)
        running_return = 0
        running_advantage = 0
        for t in reversed(range(len(rewards))):
            running_return = rewards[t] + self.gamma * running_return * (1 - dones[t])
            running_advantage = running_return - values[t].detach()
            returns[t] = running_return
            advantages[t] = running_advantage

        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # PPO update
        for _ in range(5):  # Multiple epochs
            probs = self.actor(states)
            dist = Categorical(probs)
            new_log_probs = dist.log_prob(actions)
            ratio = torch.exp(new_log_probs - old_log_probs)
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon) * advantages
            actor_loss = -torch.min(surr1, surr2).mean()

            value_pred = self.critic(states).squeeze()
            value_loss = nn.MSELoss()(value_pred, returns)

            entropy = dist.entropy().mean()

            loss = actor_loss + self.value_coef * value_loss - self.entropy_coef * entropy

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

# Example usage
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

ppo = PPO(state_dim, action_dim, lr=3e-4, gamma=0.99, epsilon=0.2, value_coef=0.5, entropy_coef=0.01)

for episode in range(1000):
    state, _ = env.reset()
    done = False
    episode_reward = 0
    states, actions, log_probs, rewards, dones = [], [], [], [], []

    while not done:
        action, log_prob = ppo.get_action(state)
        next_state, reward, done, _, _ = env.step(action)

        states.append(state)
        actions.append(action)
        log_probs.append(log_prob)
        rewards.append(reward)
        dones.append(done)

        state = next_state
        episode_reward += reward

    ppo.update(states, actions, log_probs, rewards, dones)

    if episode % 10 == 0:
        print(f"Episode {episode}, Reward: {episode_reward}")

env.close()

4.2 Direct Preference Optimization

Direct Preference Optimization (DPO, https://arxiv.org/pdf/2305.18290)

Starting with the objective we try to optimize in RLHF, rewrite the preference likelihood with policy (LLM) $\pi$. Assuming we are using Bradley-Terry model for the reward model

$$p^*(y_1 \succ y_2 | x) = \frac{1}{1 + \exp\left(\beta \log \frac{\pi^*(y_2|x)}{\pi_\text{ref}(y_2|x)} - \beta \log \frac{\pi^*(y_1|x)}{\pi_\text{ref}(y_1|x)}\right)}$$

Negative Log-Likelihood loss function could be written as $$\mathcal{L}_\text{DPO}(\pi_\theta; \pi_\text{ref}) = -\mathbb{E}_{(x,y_1,y_2)\sim\mathcal{D}}\left[\log \sigma \left(\beta \log \frac{\pi_\theta(y_1 | x)}{\pi_\text{ref}(y_1 | x)} - \beta \log \frac{\pi_\theta(y_2 | x)}{\pi_\text{ref}(y_2 | x)}\right)\right]$$

• No need to iterate between "sampling -> grading -> updating" during model training

!pip install datasets

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm
from huggingface_hub import login


# Load model and tokenizer
model_name = "meta-llama/Llama-3.2-1B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Freeze all parameters except for the last layer
for param in model.parameters():
    param.requires_grad = False
for param in model.lm_head.parameters():
    param.requires_grad = True

# Load and preprocess the dataset
# dataset URL: https://huggingface.co/datasets/Anthropic/hh-rlhf
# Note: using test set to fine-tune, just for demo purpose
dataset = load_dataset("Anthropic/hh-rlhf", split="test")

def preprocess_function(examples):
    tokenizer.pad_token = tokenizer.eos_token
    chosen = tokenizer(examples["chosen"], truncation=True, padding="max_length", max_length=512, return_tensors="pt")
    rejected = tokenizer(examples["rejected"], truncation=True, padding="max_length", max_length=512, return_tensors="pt")
    return {
        "chosen_input_ids": chosen["input_ids"],
        "chosen_attention_mask": chosen["attention_mask"],
        "rejected_input_ids": rejected["input_ids"],
        "rejected_attention_mask": rejected["attention_mask"],
    }

tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
dataloader = DataLoader(tokenized_dataset, batch_size=4, shuffle=True)

# DPO loss function
def dpo_loss(chosen_logits, rejected_logits, beta=0.1):
    chosen_rewards = chosen_logits.sum(dim=1)
    rejected_rewards = rejected_logits.sum(dim=1)
    loss = -torch.log(torch.sigmoid(beta * (chosen_rewards - rejected_rewards))).mean()
    return loss

# Training loop
optimizer = optim.Adam(model.lm_head.parameters(), lr=1e-5)
model.train()

num_epochs = 3
for epoch in range(num_epochs):
    total_loss = 0
    for batch in tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}"):
        chosen_input_ids = torch.stack(batch["chosen_input_ids"]).to(model.device)
        chosen_attention_mask = torch.stack(batch["chosen_attention_mask"]).to(model.device)
        rejected_input_ids = torch.stack(batch["rejected_input_ids"]).to(model.device)
        rejected_attention_mask = torch.stack(batch["rejected_attention_mask"]).to(model.device)

        chosen_outputs = model(input_ids=chosen_input_ids, attention_mask=chosen_attention_mask)
        rejected_outputs = model(input_ids=rejected_input_ids, attention_mask=rejected_attention_mask)

        loss = dpo_loss(chosen_outputs.logits, rejected_outputs.logits)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)
    print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")

# Save the fine-tuned model
model.save_pretrained("llama2-7b-dpo-tuned")
tokenizer.save_pretrained("llama2-7b-dpo-tuned")

print("DPO fine-tuning completed and model saved.")

# Example of using the fine-tuned model
input_text = "Translate the following English text to French: 'Hello, how are you?'"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)

with torch.no_grad():
    output = model.generate(input_ids, max_length=100, num_return_sequences=1)

generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(f"Generated text: {generated_text}")

Tags: techblog

Share: X (Twitter) Facebook LinkedIn

• Next Post

• Email me
• GitHub
• X (Twitter)
• LinkedIn

Cheng Shen  •  2026  •  scaufish.github.io  •  Edit page

Powered by Beautiful Jekyll