GRPO: How Language Models Learn to Reason

Do you remember the times when we used to make LLMs count the occurrence of a specific letter in a word, like "How many r's in strawberry?"

Back then, LLMs used to get it wrong a lot of times, but nowadays they don't. Well, one of the factors behind it is the emergence of reasoning capabilities (the main reason behind it was the tokenization issue). It allows the model to reason through the problem.

When this happened, the models could think like:

"I have to count the number of r's, so first let me break down the word into individual letters: S T R A W B E R R Y. Now let's count the number of r's sequentially: S is not an r, T is not an r, R is r, so the count becomes 1 ......"

This emergent property helped LLMs solve complex problems by breaking them down and thinking step by step.

In today's time, we use GRPO. It stands for Group Relative Policy Optimization.

Here is the abstract overview of what it does:

Step 1 — It generates a few model responses.

Step 2 — It scores every model response.

Step 3 — It compares every response by the model relatively in the group and assigns a score.

Step 4 — It updates the model parameters using the advantages, so that the good responses become more likely.

After having a quick overview, let's begin with the deep explanation.

Sampling

Given an input, the model generates G number of outputs. The output of the model is represented by oᵢ.

Advantage Calculation

Aᵢ = (rᵢ − mean(r₁, r₂, …, rG)) / std(r₁, r₂, …, rG)

After getting the G number of outputs, we simply grade each of the outputs by a reward function or another model, and use an advantage function to calculate the Advantage value for each of the outputs.

We simply take every output, subtract the mean, and divide it by the standard deviation. We calculate the mean and standard deviation using each of the outputs in a group.

After standardization, we can tell which responses are better than an average response: if Aᵢ > 0, then the response is actually better than the average response; if Aᵢ < 0, then the response is worse than the average response.

So we know which responses are better; now we have to update the model so that it produces better responses.

Policy/Model Update

Before we begin this section, I would like to tell you that a policy is something that generates responses or takes action, so in our case, the policy is the language model.

J_GRPO(θ) =

[ (1/G) × Σᵢ₌₁ᴳ min( rᵢ(θ)Aᵢ, clip(rᵢ(θ), 1−ε, 1+ε)Aᵢ ) ]

− βD_KL(π_θ ∥ π_ref)

The equation above can be broken down into its sub-pieces to make it more interpretable.

1 — The Probability Ratio

rᵢ(θ) = πθ(oᵢ | q) / πθₒₗd(oᵢ | q)

The equation above represents the probability ratio. Here, π represents the language model's policy, and θ represents the parameters of the model.

π_θ(oᵢ | q)

represents the probability that the current model assigns to generating oᵢ given prompt q.

The probability ratio is the probability assigned by the new model divided by the probability assigned by the old model. The difference between the old and the new models will be cleared in the example at the end.

If the probability ratio rᵢ(θ) > 1, then the model assigns a higher probability to the response oᵢ by the new model. If rᵢ(θ) < 1, then the model assigns a lower probability to the response oᵢ by the new model.

2 — The Clip Function

clip(πθ(oᵢ | q) / πθₒₗd(oᵢ | q), 1−ε, 1+ε)

The clip function prevents the model from changing itself too much; it doesn't allow the probability ratio to go beyond the ε range. The deviation is capped at the 1−ε and 1+ε range.

3 — KL Divergence

Eₒᵢ~πθ [ Σₜ₌₁ᵀⁱ log( πθ(oᵢ,ₜ | q, oᵢ,<ₜ) / π_ref(oᵢ,ₜ | q, oᵢ,<ₜ) ) ]

The above equation is the KL divergence equation. It tells us how far off the probability distribution of the new model πθ(oᵢ,ₜ | q, oᵢ,<ₜ) is from the reference model π_ref(oᵢ,ₜ | q, oᵢ,<ₜ).

The reference model is the model that we took right after SFT (Supervised Fine-Tuning) and before RL fine-tuning.

So the equation calculates how different the new model's probabilities are from the reference model for those responses (oᵢ), and takes the expected value of that difference.

Walkthrough

Let's say we take the batch size of 5, and each of the prompts in a batch contains 8 outputs of those responses, so the group size is 8, and the number of epochs is 3.

Step 1

We generate 5 × 8 = 40 responses.

Step 2

Then we calculate the Advantage of each of the prompt outputs, so 40 advantages each. We calculate Advantage for every example in the group independently.

Step 3

The minimum is between the unclipped surrogate objective (rᵢAᵢ) and the clipped surrogate objective:

clip(rᵢ, 1−ε, 1+ε)Aᵢ

Then we calculate the KL divergence function. Here, β is a hyperparameter that defines how strongly you want to penalize the model for deviating from the reference model.

Finally, after having everything we need, we calculate the objective function and do a single backward pass. And we do it three times for a single batch because remember our number of epochs is 3.

NOTE: At the first epoch of every batch, πθₒₗd = πθ, meaning both models are the same because we haven't done any backward pass yet.