I’m writing this post mostly as a reminder to myself. Even though I learned this concept a long time ago, I still find myself coming back to these fundamentals whenever a new training technique comes along.

The Gap

Most tutorials illustrate next-token prediction like the picture on the left: text in, one token out. That’s fine as a mental model, but it doesn’t match what we see in the actual architecture and its outputs.

The actual input and output of an LLM are both sequences. Take the input “I love”. To keep things simple, one word is one token. The model returns two rows of logits:

• Row 1 sees only “I” and predicts what comes after “I”.
• Row 2 sees “I love” and predicts what comes after “I love”.

The causal mask makes this work. Position $t$ can attend only to positions $\le t$, so every row makes a prediction without looking at the answer.

Also note what the raw output is: logits, not tokens. Softmax turns each row into a probability distribution. Choosing a token, by argmax or by sampling, is a separate step that happens outside the model.

1
2

logits = model(ids)          # ids: [1, 2]  ->  logits: [1, 2, V]
probs = logits.softmax(-1)   # one distribution per position

Let’s demonstrate with actual code.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

from transformers import AutoModelForCausalLM, AutoTokenizer

# load tokenizer and model from huggingface
model_name = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# the input text
input_text = "The capital of France is Paris, and the capital of Japan is"
input_ids = tokenizer(text=input_text, return_tensors="pt").input_ids
token_to_string = [tokenizer.decode(token_id) for token_id in input_ids[0]]

print(f"The input: {input_text}")
print(f"input tokens: {input_ids}")
print(f"input tokens shape: {input_ids.shape}")
print("string to token mapping")
for t, (i, s) in enumerate(zip(input_ids[0].tolist(), token_to_string)):
    print(f"    [{t:>2}] {s:<16} {f'id={i}'}")
print()

# the output
logits = model(input_ids=input_ids).logits.float()
print("The output (one forward pass):")
print(f"input tokens shape: {input_ids.shape}")
print(
    f"output logits shape: {logits.shape} <- [batch, T={input_ids.shape[1]}, vocab V={logits.shape[-1]}]"
)
print(
    f"{input_ids.shape[1]} tokens in, {logits.shape[1]} tokens out. Each row is a score for every one of the {tokenizer.vocab_size} vocab entries: the model's guess for the token that comes AFTER that position."
)
print()

# decode every row
probs = logits[0].softmax(-1)
p_top, pred = probs.max(-1)
print("Every row = next-token prediction for its prefix")
head = f"  {'t':>3}  {'model sees input[0..t]':<{len(input_text)}}  {'predicts':<14} {'p':>5}   actual input[t+1]"
hits = 0
for t in range(input_ids.shape[1]):
    seen = tokenizer.decode(input_ids[0, : t + 1])
    guess = tokenizer.decode(pred[t])
    line = (
        f"  {t:>3}  {seen:<{len(input_text)}}  {guess:<14} {p_top[t].item():>5.2f}   "
    )
    if t < input_ids.shape[1] - 1:
        ok = pred[t].item() == input_ids[0, t + 1].item()
        hits += ok
        line += ("✓ " if ok else "✗ ") + token_to_string[t + 1]
    if t == input_ids.shape[1] - 1:
        line = line + "   <- the 'one next word'"
    print(line)
print(
    f"\n  {hits}/{input_ids.shape[1] - 1} rows guessed the token that really comes next in the input."
)
print(
    "The 'one next word' everybody expects is only the LAST row. Training grades ALL rows at once (labels = input shifted left by one)."
)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

The input: The capital of France is Paris, and the capital of Japan is
input tokens: tensor([[1784, 8961, 1307, 5498, 1395, 6993, 1044, 1321, 1278, 8961, 1307, 6777,
         1395]])
input tokens shape: torch.Size([1, 13])
string to token mapping
    [ 0] The              id=1784
    [ 1]  capital         id=8961
    [ 2]  of              id=1307
    [ 3]  France          id=5498
    [ 4]  is              id=1395
    [ 5]  Paris           id=6993
    [ 6] ,                id=1044
    [ 7]  and             id=1321
    [ 8]  the             id=1278
    [ 9]  capital         id=8961
    [10]  of              id=1307
    [11]  Japan           id=6777
    [12]  is              id=1395

The output (one forward pass):
input tokens shape: torch.Size([1, 13])
output logits shape: torch.Size([1, 13, 131072]) <- [batch, T=13, vocab V=131072]
13 tokens in, 13 tokens out. Each row is a score for every one of the 131072 vocab entries: the model's guess for the token that comes AFTER that position.

Every row = next-token prediction for its prefix
    0  The                                                          ,               0.23   ✗  capital
    1  The capital                                                   of             0.35   ✓  of
    2  The capital of                                                the            0.20   ✗  France
    3  The capital of France                                         is             0.98   ✓  is
    4  The capital of France is                                      Paris          0.86   ✓  Paris
    5  The capital of France is Paris                               ."              0.54   ✗ ,
    6  The capital of France is Paris,                               which          0.26   ✗  and
    7  The capital of France is Paris, and                           the            0.49   ✓  the
    8  The capital of France is Paris, and the                       capital        0.33   ✓  capital
    9  The capital of France is Paris, and the capital               of             0.99   ✓  of
   10  The capital of France is Paris, and the capital of            Germany        0.17   ✗  Japan
   11  The capital of France is Paris, and the capital of Japan      is             1.00   ✓  is
   12  The capital of France is Paris, and the capital of Japan is   Tokyo          0.99      <- the 'one next word'

  7/12 rows guessed the token that really comes next in the input.
The 'one next word' everybody expects is only the LAST row. Training grades ALL rows at once (labels = input shifted left by one).

Training Uses Every Row

During training the targets are simply the input shifted by one position. Each row is scored against the token that really comes next in the data:

So one forward pass over $T$ tokens gives $T$ training signals, all computed in parallel. Every position is a training example. This is why transformers train so efficiently.

We can see this actual code at HuggingFace’s transformers library:

1
2
3
4
5
6
7
8
9
10
11
12

# ForCausalLMLoss at src/transformers/loss/loss_utils.py

    if shift_labels is None:
        # Shift so that tokens < n predict n
        labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
        shift_labels = labels[..., 1:].contiguous()

    # Flatten the tokens
    logits = logits.view(-1, vocab_size)
    shift_labels = shift_labels.view(-1)
    ...
    loss = fixed_cross_entropy(logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)

On an older version of transfomers library, before the loss was moved into a shared function, the lines are even more readable. The following snippet is taken from modeling_llama.py of version 4.44.2:

1
2
3
4
5
6

# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss_fct = CrossEntropyLoss()
...
loss = loss_fct(shift_logits, shift_labels)

Everything Else is a Variation

Every other stage, be it SFT, knowlege distillation or RL, does the same thing: feed a sequence in, get one distribution per position back, apply cross entropy. The only change is which positions count, what the target is, and how much each position weighs.

If you think this post is helpful please share on the comment section.

Thanks.

From PPO to GDPO — Understanding the Evolution of Reinforcement Learning Algorithms

Introduction: Why Reinforcement Learning?

The Two Stages of LLM Training

Training a Large Language Model is fundamentally a two-stage process:

Stage 1: Pre-training (Knowledge Acquisition)

• The model learns from trillions of tokens of text
• It learns to predict the next token: $P(x_t | x_{<t})$
• Result: A “document completor” that understands language structure

Stage 2: Post-training (Behavioral Alignment)

• The model learns to be helpful, harmless, and honest
• It learns to follow instructions and provide useful responses
• Result: A helpful AI assistant

The Problem: A pre-trained model asked “How do I bake a cake?” might respond with “How do I bake a pie?” — continuing a pattern rather than answering the question.

Byte-Pair Encoding

What is Byte-Pair Encoding

Byte-pair encoding (BPE) is a compression algorithm that can be used for text tokenization.

LLaMA 2 Open Foundation and Fine-Tuned Chat Models

📖 Abstract

In this work, we develop and release Llama 2, a collection of pretrained and fine-tuned large language models (LLMs) ranging in scale from 7 billion to 70 billion parameters. Our fine-tuned LLMs, called LLaMA 2-Chat are optimized for dialogue use cases. Our models outperform open-source chat models on most benchmarks we tested, and based on our human evaluations for helpfulness and safety, may be a suitable substitute for closedsource models. We provide a detailed description of our approach to fine-tuning and safety improvements of LLaMA 2-Chat in order to enable the community to build on our work and contribute to the responsible development of LLMs.

LLaMA Open and Efficient Foundation Language Models

📖 Abstract

We introduce LLaMA, a collection of foundation language models ranging from 7B to 65B parameters. We train our models on trillions of tokens, and show that it is possible to train state-of-the-art models using publicly available datasets exclusively, without resorting to proprietary and inaccessible datasets. In particular, LLaMA-13B outperforms GPT-3 (175B) on most benchmarks, and LLaMA65B is competitive with the best models, Chinchilla-70B and PaLM-540B. We release all our models to the research community.

Recall that from Off-Policy Gradient Methods, the objective function is

$$
\begin{aligned}
J^{\theta’}(\theta) &= \hat{\mathbb{E}} _{(s_t, a_t) \sim \pi _{\theta’} (\tau)} \bigg[ \frac{\pi _{\theta}(a_t | s_t)}{\pi _{\theta’}(a_t | s_t)} A^{\theta’}(s_t, a_t) \bigg]
\end{aligned}
$$

We learned that from Importance Sampling, the two distributions shouldn’t differ too much, therefore in PPO, they added a constraint term:

$$
J ^{\theta’} _{PPO} (\theta) = J ^{\theta’} (\theta) - \beta KL(\theta, \theta’)
$$

Where $\beta$ is a weight, this constraint term ensures that $\theta$ and $\theta’$ are similar.

Recall that from Policy Gradient Methods, the calculation of gradient is

$$
\begin{aligned}
\nabla \hat{R} _{\theta} &= \hat{ \mathbb{E}} _{\tau \sim \pi _{\theta}(\tau)} \bigg[ R(\tau) \nabla \log \pi _{\theta} (\tau) \bigg]
\end{aligned}
$$

where

• $\tau$ is a trajectory
• $\pi_{\theta}$ is a stochastic policy
• $\nabla_ {\pi _{\theta}}$ is the gradient operator of the policy $\pi$ with respect to $\theta$
• $\hat{\mathbb{E}} _{\tau \sim \pi _{\theta}(\tau)}[\dots]$ is the empirical average expectation over a finite batch of samples in an algorithm that alternates between sampling and optimization

In this setting, we have to:

• Use $\pi _{\theta}$ to collect data. When the parameter $\theta$ is updated, we must sample the training data again.
• This is called on-policy training

What Are Policy Gradient Methods

Policy gradient methods are a class of reinforcement learning algorithms used for optimizing parametrized policies by gradient descent to maximize the expected return, which is the cumulative reward received over a long-term horizon. These methods are based on updating the policy parameters such that the probabilities of actions leading to higher returns are increased. In comparison, the probabilities of actions leading to lower returns are decreased.

Iteratively optimizing the policy using this approach makes it possible to arrive at an optimal policy that maximizes the expected return. This technique has proven effective in solving complex decision-making problems in various fields, including robotics, game playing, and natural language processing.

Importance sampling is a technique commonly used in statistics and machine learning to estimate properties of a distribution that may be difficult to compute directly.

What is Out-of-Distribution (OOD) Detection and Why is it Important?

Most machine learning models are trained based on the closed-world assumption, meaning that the test data is assumed to have the same distribution as the training data. When we train a model, we usually form our training and testing data by randomly splitting the data we have into two sets, in which both datasets have the same label distribution. However, in the real-world scenario, this assumption doesn’t necessarily hold. For example, a model trained to classify cats and dogs may receive an image of a dolphin. Even worse, it might highly confident classify such inputs into in-distribution classes.

Detecting if the received model inputs have the same distribution as the training data is called out-of-distribution detection.