Reinforcement learning

Reinforcement learning is a way of training where a model takes actions, a reward scores each action, and the training raises the probability of actions with high rewards. In this chapter, the model is a language model, and an action is an answer to a question. For example, a model can solve 1,000 math problems, receive the reward 1 for each correct final answer and 0 for each wrong one, and become better at math from these rewards alone.

Fine-tuning (see Chapter 11) needs a correct answer for each example. Reinforcement learning only needs a way to score answers. For example, a program can check the final number of a solution, but it cannot write the steps that lead to it. The model can then learn steps that no example contains, and Chapter 13 shows models that are trained in this way.

12.1Policy gradient

The model that is trained is the policy. Let $p(a)$ denote the probability that the policy gives the answer $a$. Our task is to change the policy so that its answers get higher rewards on average.

The policy gradient method¹ trains the policy in rounds. In each round, the policy gives several answers to the same question, and a reward scores each answer. The advantage $A(a)$ of an answer $a$ is its reward minus the average reward of the round². After this, the parameters move in the direction that increases the following sum over the answers $a_1, a_2, \ldots, a_m$ of the round:

$$A(a_1) \cdot \ln p(a_1) + A(a_2) \cdot \ln p(a_2) + \cdots + A(a_m) \cdot \ln p(a_m)$$

This means that answers with a positive advantage get a higher probability, and answers with a negative advantage get a lower probability.

As an example, consider the question “What is 17 + 25?” and a policy that only chooses among the answers 42, 32 and 43. The policy has one parameter $z(a)$ for each answer, and $p(a)$ is proportional to $e^{z(a)}$. Initially, each parameter is 0, so each answer has probability 1/3. Suppose that the policy gives four answers in the first round, where the average reward is 0.5:

For this policy, the derivative of $\ln p(a)$ with respect to $z(b)$ is $1 - p(b)$ when $b = a$, and $-p(b)$ otherwise. Using this, we can calculate the change of each parameter. With the learning rate 1, $z(42)$ grows by $2 \cdot 0.5 \cdot 2/3 + 2 \cdot 0.5 \cdot 1/3 = 1$. A similar calculation shows that $z(32)$ and $z(43)$ fall by 0.5:

After one round, the probability of the correct answer has grown from 1/3 to about 0.69.

Note that the average reward is subtracted because rewards can be all positive, for example scores between 0 and 10. Without the subtraction, each answer would push its own probability up, also an answer that is worse than the others. When every answer of a round gets the same reward, every advantage is 0, and the round changes nothing.

In practice, writing the answers usually takes most of the training time, because the policy writes each answer one token (a piece of a word) at a time.

12.2Reward models

A program can score an answer when its result can be checked: a number, or code that must pass tests. However, most requests have no such program, for example a request to write a polite reply to an upset customer.

The idea is to ask people to compare two answers to the same question and to train a reward model that gives each answer a score. Let $s(x)$ denote the score of the answer $x$. The reward model predicts that people prefer the answer $x$ to the answer $y$ with the following probability:

$$\frac{1}{1 + e^{-(s(x) - s(y))}}$$

For example, when $s(x) = 2$ and $s(y) = 0$, the probability is $1/(1 + e^{-2}) \approx 0.88$. The following picture shows the probability for each difference of the scores:

When the scores are equal, the probability is 0.5, and a larger difference moves it nearer to 0 or 1.

The training changes the reward model so that the comparisons that people made get high probabilities, and after this, the score $s(a)$ can be used as the reward of the answer $a$. This combination is called reinforcement learning from human feedback³.

12.3Reward hacking

Reward hacking is the case where the policy gets high rewards with answers that are not good. For example, if the people who compared the answers often preferred longer answers, the reward model gives higher scores to longer answers. The policy then learns to write long answers, also when a short answer is better. The reward goes up, and the quality does not.

For this reason, the reward usually includes a penalty for moving away from the starting model, i.e., the model before reinforcement learning: $$s(a) - \beta \cdot \ln \frac{p(a)}{p_0(a)}$$ where $p_0(a)$ is the probability of $a$ in the starting model, and $\beta$ sets the strength of the penalty. For example, when $\beta = 0.1$, an answer whose probability has grown from 0.1 to 0.4 loses $0.1 \cdot \ln 4 \approx 0.14$ of its reward.

A larger $\beta$ keeps the policy nearer to the starting model, so the policy forgets less (see Section 11.1). On the other hand, a larger $\beta$ also limits how much the policy can improve.