Batch Norm

Normalizes the mean and variance for all instances of a feature in a batch. Commonly used in older architectures like CNNs by smoothing out the loss landscape and allowing for higher learning rates and faster convergence. Example implementation:

def batch_norm(features: np.array, i: int, eps=1e-5) -> np.array:
    x = features[:,i]
    mean = x.mean()
    variance = x.var()
    x = (x - mean) / np.sqrt(variance + eps)
    return x

In practice, we include a gamma and beta parameter which are learnable parameters of length F, where F is the number of features. These parameters that allow the network to undo the normalization if needed.

def batch_norm2(
        features: np.array,
        gamma: np.array,
        beta: np.array,
        eps=1e-5
    ):
        mean = features.mean(axis=0)
        variance = features.var(axis=0)
        normalized = (features - mean) / np.sqrt(variance + eps)
        return gamma * normalized + beta

See also: Layer norm

Byte Pair Encoding (BPE)

A tokenization algorithm which iteratively merges the most frequent pair of tokens. For example, let’s take the sentence snug_as_a_bug_in_a_rug as our corpus (underscores used to make things easier to see).

First, we break up the text into its individual characters:

corpus: ['s', 'n', 'u', 'g', '_', 'a', 's', '_', 'a', '_', 'b', 'u', 'g', '_', 'i', 'n', '_', 'a', '_', 'r', 'u', 'g']
vocab: ['s', 'n', 'u', 'g', '_', 'a', 'b', 'i', 'r']

After splitting by chars, we have 9 unique tokens in our vocabulary. Next, we merge the most common pair in our vocabulary, ('u','g') => 'ug' and retokenize:

corpus: ['s', 'n', 'ug', '_', 'a', 's', '_', 'a', '_', 'b', 'ug', '_', 'i', 'n', '_', 'a', '_', 'r', 'ug']
vocab: ['s', 'n', 'u', 'g', '_', 'a', 'b', 'i', 'r', 'ug']

Then, we find that the most common pair is ('_','a') which gives us the next merged token "_a":

corpus: ['s', 'n', 'ug', '_a', 's', '_a', '_', 'b', 'ug', '_', 'i', 'n', '_a', '_', 'r', 'ug']
vocab: ['s', 'n', 'u', 'g', '_', 'a', 'b', 'i', 'r', 'ug', '_a']

We can repeat this process until the vocabulary is the size we need.

See also: Wordpiece tokenization

Direct Preference Optimization

Direct preference optimization (DPO) is a contrastive learning method, often used for fine-tuning an LLM to align with human preferences.

Like other contrastive learning methods (for example those used to fine-tune encoders) DPO works with (anchor, positive, negative) triplets. In this case, the anchor is the prompt, the positive sample is the preferred completion, and the negative sample is the bad completion. By computing the probabilities the models assign to each of the responses, we can train the model to assign higher probabilities to the preferred response.

DPO is simpler, more stable, and more lightweight than more traditional methods like PPO and RLHF. It does not require a second model at training time like PPO and RLHF do. It doesn’t require a separately trained reward model at all. The model is updated by simple gradient descent rather than generating during training.

However, since DPO is technically off-policy from an RL perspective, since the samples it is scoring are not actually generated from the model.

DPO pairs technically do not need to be created based on human preferences. Any reward model will do. For example:

  • If we’re training a chess model, we can use a chess engine to score potential moves
  • If we’re training a code generation model, we can use the compiler’s feedback to create positive and negative samples

Distillation

Distillation is the task of training a smaller “student” model to behave like a larger “teacher” model.

Classically, this has meant directly using the teacher’s output distribution, and minimizing KL(p_teacher || p_student).

Less formally, the larger teacher model can be used as a pseudo-labeler of real inputs, and a smaller student can learn from those labels.

Entropy

When training LLMs, we want some level of the model’s uncertainty when generating tokens. To do this, we multiply the probability by the negative log of the probability.

import math
def calc_per_token_entropy(probabilities):
    entropy = 0
    for probability in probabilities:
        entropy -= probability * math.log(probability)
    return entropy

The negative log of the probability can be thought of as the “surprise” from that token: high probability tokens have very little surprise to them, and low probability tokens have surprise approaching infinity. Then, multiplying by the probability itself helps to balance the total by preventing it from exploding.

There is a connection to perplexity, which is $ e^{H(x)} $, where $H(x)$ is the entropy of the distribution.

Group-Relative Preference Optimization (GRPO)

GRPO is an RL algorithm for LLM post-training which attempts to determine how good a rollout is in order to apply a reward to it.

GRPO has some advantages over more traditional PPO. In PPO, we use a second model (the “critic”) to score rollouts, and in practice this model may on the order of the same size as the model we’re training. This is both slow and resource intensive.

In GRPO, we get rid of the critic entirely.

prompt = 'write a python function to determine if a number is prime'
group_size = 8 # usually between 4 and 16
responses = [
    generate_response(model, prompt, temp=1.1)
    for _ in range(group_size)
]

rewards = [
    calculate_reward(response) for response in responses
]

mean_reward = np.mean(rewards)
std_reward = np.std(rewards)
advantages = [(reward-mean_reward)/std_reward for reward in rewards]

# policy gradient update
log_probs = [model.compute_log_prob(r) for r in responses]
pg_loss = 0
for log_prob, advantage in zip(log_probs, advantages):
    pg_loss += -log_prob * advantage

# KL penalty prevents model from changing too quickly
ref_log_probs = [ref_model.compute_log_prob(r) for r in responses]
kl = sum(r - p for p, r in zip(log_probs, ref_log_probs)) / group_size

loss = pg_loss / group_size + beta * kl

KL-Divergence

A measurement of the “distance” between two distributions.

from math import log

def kl_divergence(p:list[float], q:list[float]) -> float:
  total = 0
  for a,b in zip(p, q):
    total += a * log(a/b) if p != 0 else 0
  return total

Note that KL-divergence is not symmetrical.

Where it shows up:

  • In RLHF, it is used to measure the change between old policy and the updated policy, and prevent it from drifting too far using a KL penalty.
  • In knowledge distillation, we directly train on minimizing KL(p_teacher||p_student) on the distribution of logits.

Layer Norm

This could probably be called “token norm” and avoid a lot of confusion since in practice this is often used in LLMs.


# hypothetical input embeddings for "ai is cool"
sequence = np.array([
    [0.3, 0.1, 0.4, 0.1], # "ai"
    [0.5, 0.9, 0.2, 0.6], # "is"
    [0.5, 0.3, 0.5, 0.8]  # "fun"
])

def layer_norm(sequence, eps=1e-5):
    """
    Example for learning, not vectorized.

    Here, `sequence` has dimensions (seq_length, emb_dim)
    and we are ignoring the batch dimension.
    """
    means = []
    variances = []
    for emb in sequence:
        means.append(emb.mean())
        variances.append(emb.var())
    result = []
    for emb,mean,variance in zip(sequence, means, variances):
        emb -= mean
        emb /= np.sqrt(variance + eps)
        result.append(emb)
    return np.array(result)

Today, many more modern LLMs use RMSNorm instead, which skips subtracting the mean.

See also: Batch norm

Offline RL

Models are trained from a fixed dataset of preferences rather than the rollouts they would naturally generate. This is much faster, but can be less useful if the training data distribution shifts far from the the model’s actual predictions.

Examples of offline RL algorithms include DPO, IPO, and KTO.

Online RL

Models generate rollouts, get scored, and update from that signal. Considered stronger than offline RL because the model is getting direct feedback on its own predictions. However, the scoring itself can be expensive, and gradient updates must wait for the scores to come back.

Examples of online RL algorithms include PPO, GRPO, and REINFORCE.

Parallelism

Data parallel

Replicate the same model over all devices.

  1. For each batch, split the data between devices.
  2. Each copy computes a forward pass, a loss, and a backwards pass to compute local gradients.
  3. Each device reports its gradients.
  4. Each copy is updated with the same average gradient.

Model parallel

This is used when a model is too large to fit on a single device.

The model can be split between multiple devices, with each device computing its part, handing it off to the next device, then waiting for the next batch.

This is very slow. Pipeline parallelism is an improvement to Model Parallelism.

Pipeline parallel

The idea here is to reduce wait time between devices.

Rather than a single device computing and then waiting as in model parallelism, a device finishes its forward pass on a “micro-batch” and then immediately starts working on the next “micro-batch”.

Similarly, backwards passes flow in reverse.

Tensor parallel

This is used when even a single layer cannot fit on a device.

The individual tensors (weights and activations) can be split across multiple devices to enable a huge matrix multiplications.

Reinforcement Learning from Human Feedback (RLHF)

RLHF has three steps:

  1. Train a reward model on human preference pairs. This is very similar to a standard pairwise Learn-to-Rank model where we train on (input, positive, negative) triplets. You can simply add a linear layer to the top of the LLM to create R(x,y) which generates the ranking scores.

2a. Calculate the expected reward $ E[R(x,y)] $

2b. Penalize large KL-divergence with the old policy

Rejection Sampling/Rejection Fine-Tuning

The idea of rejection sampling is to generate multiple responses or “roll-outs” to an input. Typically we will increase sampling temperature Responses which do not meet some criterion are “rejected” or thrown away. This leaves the responses which do pass the criterion. If we fine-tune on these examples, this is called rejection fine-tuning.

RFT has some similarity with RL techniques like GRPO, which also involve using multiple roll-outs.

On the one hand, RFT is very simple and intuitive. It avoids a lot of the RL complexity such as entropy collapse and policy divergence. Additionally, because GRPO relies on the relative advantages within the group, RFT can outperform GRPO in situations where all members of the group have perfect or close to perfect rewards.

  • If all members have perfect scores, the advantage will be zero for each completion.
  • If all members have close to perfect scores, GRPO will amplify noise within the group.

On the other hand, RFT can be far less sample efficient than RL, because it only learns from positive examples (since the negatives are rejected). For hard problems, your acceptance rate for the problem may be very low.

  • If your acceptance rate is 1 in 16, you will throw away 94% of your work to get the 6% to learn from.
  • Meanwhile with GRPO, if 1 out of 16 samples in the group is correct, your algorithm can learn from all 16.
  • However, if none of the samples in the group get a reward, the relative advantage for all members collapses to zero and you get nothing. In this case, it would be better to keep generating until we get a correct completion.

For this reason, we can use RFT until the acceptance rate is high enough for GRPO.

Speculative Decoding

The idea of speculative decoding is to use a small “draft” model to speed up inference for the big “target” model. This is possible because the target model can predict probabilities for multiple sequences in parallel. For example, a small draft model might predict:

draft_tokens = ['take', 'me', 'out', 'to', 'the', 'ballgame']

Now, the target model can simply predict the probabilities for each of these in parallel:

probabilities = model([
  {'token': 'take'},
  {'token': 'me', 'prefix':['take']},
  {'token': 'out', 'prefix':['take', 'me']},
  {'token': 'to', 'prefix':['take', 'me', 'out']},
  {'token': 'the', 'prefix':['take', 'me', 'out', 'to']},
  {'token': 'ballgame', 'prefix':['take', 'me', 'out', 'to', 'the']},
])

Because we’re not just using greedy decoding, we need to accept tokens from the draft model probabilistically. In other words, we accept a token with probability p_accept(token) = min(1, target_prob(token) / draft_prob(token)).

  • So if the target model thinks the token is a lot less likely than the draft model does, the ratio will be small and p_accept will be small for that token.
  • On the other hand, if the target model agrees with the draft model the ratio will be close to 1 and p_accept will be close to 1.
  • In the case where the target model thinks the token is more likely than the draft model did, it’s somewhat of a happy accident that the token was actually selected by the draft model but the acceptance probability will be 1.

If the target model rejects a token that the draft model proposes, none of its subsequent tokens will be considered. The draft model will rerun with all of the accepted tokens and propose new tokens.

This is how it is possible to speed up the target model’s inference without losing quality.

Wordpiece Tokenization

An algorithm for training a tokenizer with a slightly more advanced rule for merging tokens. For example, consider the following pair frequencies in a corpus abracadabra:

Counter({
    ('a', 'b'): 4,
    ('b', 'a'): 3,
    ('a', 'c'): 2,
    ('d', 'a'): 2,
    ('c', 'a'): 1,
    ('c', 'd'): 1,
    ('a', 'd'): 1,
})

Under byte-pair encoding tokenization, we would automatically merge ('a','b') => 'ab'. However, wordpiece tokenization also takes the frequencies of the tokens into account individually as well:

Counter({
    'a':7,
    'b':4,
    'c':2,
    'd':2
})

In wordpiece tokenization, we merge the tokens when the pair maximizes score(x,y) = freq(xy) / (freq(x) * freq(y)). In other words, we penalize pairs if they are made up of common tokens. Let’s compute the scores:

scores = {
    ('a', 'b'): 4÷(7*4) = 0.14,
    ('b', 'a'): 3÷(4*7) = 0.11,
    ('a', 'c'): 2÷(7*2) = 0.14,
    ('d', 'a'): 2÷(2*7) = 0.14,
    ('c', 'a'): 1÷(2*7) = 0.04,
    ('c', 'd'): 1÷(2*2) = 0.25,
    ('a', 'd'): 1÷(7*2) = 0.04,
}

In this case, ('c','d') is actually the winning pair because it is composed of rare tokens.