index00
Visual01
Projects02
Writing03
Back to Writing

unscrambling the alphabet soup

an unserious guide to post training, eval, and rl environments.

Technology

If intelligence is a cake, the bulk of the cake is unsupervised learning, the icing on the cake is supervised learning, and the cherry on the cake is reinforcement learning

Yann LeCun 2016

Foundation Models

If you have ever attempted to communicate with a foundation model that isn't "instruction-tuned", you will soon come to realize the true statistical language nature of transformers.

Here is an example of Qwen 2.5 (base)'s output - equivalent to that of GPT-3.5-Turbo - from the prompt User: Can you tell me a joke? \n Assistant:.

User: Can you tell me a joke?
Assistant: 2.5
User: How does the moon cut his hair? With a moon-scissor!
Assistant: 3.7

In the above examples, one can see the resemblance of a conversation between a user and an AI-assistant, albeit either one or both of them have experienced a stroke.

Nature of LLMs

The foundation models are trained on large datasets of text using a technique called self-supervised learning. This is actually the breakthrough that made GPT, or Generative Pre-trained Transformer, possible.

Up until around 2018, state of the art natural language models still leaned heavily on quality labeled data for specific tasks. The GPT paper found that through unsupervised pre-training of transformers (now known as self-supervised learning), the model could capture the underlying semantic structure of the language, which then allowed the model to generalize to a wide range of different tasks.

But how do we get the model to actually generate the response that we want? Instead of just sampling whatever next token is "probable"?

Instruction Fine-Tuning

This is where post training comes in. Once LLMs have a grasp of how language works, it is time to make them actually helpful.

Instruction Fine-Tuning (IFT), or often known as Supervised Fine-Tuning (SFT), is a technique where we teach the model how to follow prompt instructions.

Given a prompt, we craft a quality response by hand (or today we often use a teacher model to generate proper responses). We then train the model to imitate responses with a small set of quality examples. An example of a prompt and response pair looks like this:

User: What is the capital of France?
Assistant: Paris

or maybe following our previous example, we have:

User: Can you tell me a joke?
Assistant: What do you call a fish with no eyes? A fsh.

We can then use this response pair to teach the model the way we want it to behave.

Let's see how SFT actually works.

Training Objective

SFT uses the same objective as that of pre-training. Given a prompt xx and a response y=(y1,,yT)y = (y_1, \dots, y_T) that we want the model to produce, we want the model πθ\pi_\theta with weights θ\theta to predict each token of that response in turn. Starting with the first:

πθ(y1x)\pi_\theta(y_1 \mid x)

We calculate this probability by passing the prompt into the model, which hands us back a vector of raw scores, or logits, z(1)=fθ(x)RVz^{(1)} = f_\theta(x) \in \mathbb{R}^V where VV is the vocabulary size (one score for every token in the vocabulary). The softmax function then turns these scores into probabilities.

πθ(y1=ix)=ezi(1)j=1Vezj(1),i{1,,V}\pi_\theta(y_1 = i \mid x) = \frac{e^{z^{(1)}_i}}{\sum_{j=1}^{V} e^{z^{(1)}_j}}, \qquad i \in \{1, \dots, V\}

And to get the probability of the entire response, we need to do the exact same thing for every token in the response. The catch is that the context grows by one token each time, so the model hands us back a different logit vector at every step. For the second token we feed it the prompt plus the token we just conditioned on, giving z(2)=fθ(x,y1)z^{(2)} = f_\theta(x, y_1):

πθ(y2=ix,y1)=ezi(2)j=1Vezj(2)\pi_\theta(y_2 = i \mid x, y_1) = \frac{e^{z^{(2)}_i}}{\sum_{j=1}^{V} e^{z^{(2)}_j}}

And we can see that by the chain rule of probability, that the probability of the entire response is just the product of the probabilities of each token.

πθ(yx)=s=1Tπθ(ysx,y<s)\pi_\theta(y \mid x) = \prod_{s=1}^{T} \pi_\theta(y_s \mid x, y_{<s})

where y<s=(y1,,ys1)y_{<s} = (y_1, \dots, y_{s-1}) is everything generated so far, and y<1y_{<1} is empty.

autoregressive rollout
User: Can you tell me a joke? Assistant:
step 1 of 6 — logits from the 1st forward pass
Whatz 3.844.0%
Az 3.224.2%
Whyz 3.019.8%
Surez 2.512.0%
this step0.440π(What | context)
running product4.40e-1chain rule, tokens 1..n
running log-sum-0.82negate this and you have the loss
Illustrative logits over the first six tokens of the response.

Loss Function

Now to train the model, we need to define a loss function for the model to optimize this objective.

We can define a loss function as the negative log probability of the response given the prompt.

Note: we take the log because it turns that product into a sum, which is both numerically stable (multiplying hundreds of numbers between 0 and 1 underflows to zero in floating point) and easier to differentiate. We take the negative because optimizers minimize, which maximize the likelihood of the response.

This gives us the loss function:

LSFT(θ)=E(x,y)D[s=1Tlogπθ(ysx,y<s)]\mathcal{L}_{\text{SFT}}(\theta) = -\mathbb{E}_{(x, y) \sim \mathcal{D}} \left[ \sum_{s=1}^{T} \log \pi_\theta(y_s \mid x, y_{<s}) \right]

Where D\mathcal{D} is our dataset of prompt-response pairs.

The one difference between SFT and pre-training is that the sum in the loss function runs over the response tokens only. Since the prompt is something we condition on, not something we asked the model to generate.

loss masking
prompt — x
User:Canyoutellmeajoke?Assistant:
response — y
Whatdoyoucallafishwithnoeyes?Afsh.
predicted — enters the lossconditioned on — masked out14 of 25 tokens scored

Gradient

Now that we have the loss function, we need to determine which way to nudge the parameters of the model to minimize that loss.

I won't do the full gradient calculation with respect to transformer weights here, but differentiating the loss with respect to the logits can give us some insights into how the model is learning.

First we can expand the token loss:

Ls=logπθ(ysx,y<s)=logezysj=1Vezj=zys+logj=1Vezj\begin{aligned} \mathcal{L}_s &= -\log \pi_\theta(y_s \mid x, y_{<s}) \\[4pt] &= -\log \frac{e^{z_{y_s}}}{\sum_{j=1}^{V} e^{z_j}} \\[4pt] &= -z_{y_s} + \log \sum_{j=1}^{V} e^{z_j} \end{aligned}

Then we can take the gradient by the chain rule:

Lszi=zyszi+zilogj=1Vezj=1[i=ys]  +  ezij=1Vezj=pi1[i=ys]\begin{aligned} \frac{\partial \mathcal{L}_s}{\partial z_i} &= -\frac{\partial z_{y_s}}{\partial z_i} + \frac{\partial}{\partial z_i} \log \sum_{j=1}^{V} e^{z_j} \\[6pt] &= -\mathbf{1}[i = y_s] \;+\; \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}} \\[4pt] &= p_i - \mathbf{1}[i = y_s] \end{aligned}

And to minimize the loss we would need to nudge the model parameters in the negative direction of the gradient.

We can see that this does teach the model to output the correct response, but to the detriment of other possible responses. Alternative tokens in the vocab are penalized harder for being more likely to be output, even though they could also be valid correct responses.

To address this, we will need to look at reinforcement learning next.

Reinforcement Learning

Further reading

About04