InstructGPT / RLHF (NeurIPS 2022)

 

Preliminaries

• Reinforcement Learning Basics
• PPO

 

 

Motivation

next token prediction으로 train하는 경우 사용자 intention과 align되지 않을 수 있으므로 다른 fine-tune 방법을 제안한다. 여기서는 RLHF를 사용한다. 

 

구체적으로는 reward model(RM)을 human labeled comparison dataset에 대해서 학습한다. 그 후 RM을 reward function으로 활용하여 PPO로 reward를 maximize하도록 supervised finetuning한다.

 

Figure 2: A diagram illustrating the three steps of the method.

 

Fig. 2를 참조하면, method는 세 단계로 구성된다:

(1) supervised fine-tuning(SFT)

(2) reward model(RM) training

(3) reinforcement learning via PPO on the reward model 

 

이를 통해 LLM이 human preference를 이용해 finetune할 경우 많은 task에서 성능이 향상됨을 보였다. 

 

Methods

Step 1: Collect demonstration data, and train a supervised policy

Labeler들이 input prompt distribution에 대해서 desired behavior를 정의하고, 이를 demonstration으로 제공했다. 그 후 이를 supervised learning 방식으로 train했다. 

 

처음 InstructGPT를 train하기 위해서 labeler들은 prompt를 직접 작성했다. 다음 세 종류의 prompt가 작성되었다. 

 

• Plain: 임의의 task를 떠올려서 작성하되 diversity를 갖도록 함.
• Few-shot: instruction과 그 instruction에 맞는 여러 query/response pair를 만들도록 함.
• User-based: waitlist application에 있는 use-case에 대한 prompt를 만들도록 함.

이를 바탕으로 세 종류의 dataset이 만들어졌는데, 

1. SFT dataset: labeler demonstration (13k)

2. RM dataset: labeler rankings of model outputs (33k)

3. PPO dataset:  RLHF finetuning (31k)

 

 

GPT-3를 finetune했다. 16 epoch train했다.

→ objective가 뭐임?

 

 

Step 2: Collect comparison data, and train a reward model

model output간 비교한 dataset을 이용하여 reward model이 human-preferred output을 predict하도록 train했다.

 

reward 모델로는 6B RM을 사용했다 → reward model이 LLM인가? 

reward model도 GPT-3 model, NLP dataset으로 finetune된 것

 

reward model의 loss function은 다음과 같다: 

 

Equation 1

전체 K개의 response 중 2개로 human preference를 확인하는데, 이때 human comparison dataset $D$ 중 prompt $x$에 대해 두 개의 output $y_w, y_l$이 있다고 했을 때 $y_w$가 preferred completion이다. $r(x,y)$는 reward model의 scalar output이다.

 → 앞에 negative 기호가 있으니까 차이를 키우는거다. 즉, preferred output에 대해서 reward function이 더 큰 값을 return하도록 train하는 방식이다. 

 

Step 3: Optimize a policy against the reward model using PPO

RM output을 scalar reward로 사용했다. 이를 이용해 PPO algorithm으로 train했다. 

 

random prompt를 보여주고 그에 대한 response를 받는다. 이 prompt와 response에 대해서 reward model이 reward를 부여한다. 

 

Equation 2.

 

→ 이 식이 PPO인 이유는 update 크기에 대한 soft constraint를 lagrangian의 형태로 넣어뒀기 때문이다. 

 

내 PPO 글에는 포함하지 않았는데, 원래 PPO paper를 보면 네 번째 section인 Adaptive KL Penalty Coefficient에서 TRPO를 soft constraint로 바꾼 형태의 KLPEN loss를 소개한다. 

 

Equation 3.

앞쪽 term은 eq 2에서 reward로 대체되었다. 이는 어차피 prompt x에 대한 output y가 pair 형태로 존재하므로 이에 대해서 probability를 계산할 필요가 없기 때문이다. 따라서 이미 얻은 x, y pair에 대해서 reward만 계산하면 아래의 estimated advantage function으로 weight된 prob값이 대체된다.

 

두 번째 term에 대해서는 원래 두 pdf간 차이가 커지지 않도록 하는 것이다. 

 

$$ β\sum π^{SFT} \log \frac{π^{SFT}}{π^{RL}_φ} $$

 

에서 summation은 어차피 (x,y) pair가 하나이므로 제거하고, $π^{SFT}$도 constant니까 무시하면 된다.

 

다만 eq 2식에서 KL divergence하고 순서가 다른데 왜 그렇게 됐는지는 모르겠다.

 

 

Experiments

 

Results

 

Discussion

 

 

References

[1] Ouyang, Long, et al. "Training language models to follow instructions with human feedback." Advances in neural information processing systems 35 (2022): 27730-27744.

[2] 

Footnotes