From In-context learning to RLHF (Feat. ChatGPT)

TL;DR

거시적인 발전 과정 : In-context Learning -> Instruction Tuning -> RLHF -> RLAIF
In-Context Learning은 Large Scale 언어모델을 tuning하지 않고 새로운 task에 적용할 수 있는 직관적인 방법을 제시함
Instruction Tuning은 다양한 task를 Instruction + example의 템플릿으로 캐스팅하여 Implicit하게 multi-task로 tuning하며 결과적으로 Unseen Task를 더 잘 수행함
RLHF는 인간의 선호도라는 애매한 척도를 모델링하는 Reward Model과 강화 학습을 활용하여 언어 모델을 개선하는 방법을 제시함
RLAIF는 Human Labeling Cost를 없애고 RLHF에서 추가적으로 helpfulness와 harmlessness를 모두 개선할 수 있는 자동화된 파이프라인을 제시함

In-Context Learning & Instruction Tuning

오늘은 요즘 핫한 ChatGPT와 관련된 이야기를 하려고 한다. 바로 zero-shot의 가능성을 보여준 In-context Learning의 시작과 그것을 더 개선시킨 Instruction Tuning, 마지막으로 화룡정점을 찍은 Reinforcement Learning From Human Feedback에 대한 내용이다. 이후 Anthropic에서 RLHF를 시스템적으로 개선한 RLAIF라는 방식을 추가적으로 제안하기도 했다.

In-context Learning

In-context learning은 GPT3에서 소개되면서 pretraining-finetuning paradigm의 대안을 제시했다. 언어모델이 충분히 크고(도표에 의하면 6B 이상) 다량의 corpus로 학습했다면 사람의 자연어 instruction을 이해하고 바람직한(의도에 맞는) 텍스트를 생성할 수 있는 능력을 가지고 있다는 것이다. 예를 들어 언어모델을 QA 태스크에 명시적으로 학습하지 않았다고 하더라도 '미국의 수도는 어디일까? 정답은" 과 같은 방식으로 적절한 자연어 설명을 넣어서 뒤쪽에 "워싱턴이다."를 예측하게 하여 QA를 풀 수 있다. 그래프에 나오는 Lambada 데이터도 마지막 단어를 context로 부터 예측하는 태스크인데 파라미터 수가 많아지면서 성능이 human에 가까워지는 것을 볼 수 있다. 여기에서 zero-shot은 instruction만 주는 경우를 말하고, 추가적으로 제공되는 몇몇의 (input, output) pair의 개수에 따라 few-shot의 K값이 결정되게 된다.

Instruction-Tuning

Instruction-tuning의 경우에는 multi-task learning과도 관련이 있는데 다양한 Task에 대해 적절한 자연어 설명(instruction)과 target을 표현할 수 있는 template을 만들어서 하나의 언어모델에 seq2seq 방식으로 학습하는 것이다. 원래라면 독립적인 task였던 것들을 하나의 template 자연어로 표현했으니 '자연어 명령 이해'라는 하나의 task로 잘 표현될 수 있고 implicit하게 multi-task learning을 수행한다고 볼 수도 있다. 위의 FLAN paper 예시를 보면 좀 더 쉽게 이해할 수 있다. Commonsense Reasoning은 "here is a goal"이라는 instruction과 OPTIONS 라는 선택지로 템플릿을 구성했고, Translation은 "Translate this sentence to Spanish"와 같이 명시적으로 번역을 하라고 지시를 주는 것을 볼 수 있다. 이런 방식으로 학습했을 때 모델이 지시사항을 명시적으로 받기 때문에 학습 과정에서 보지 못한 지시사항에 대해서 더 잘 일반화 하게 된다.

실험 결과도 다양한 Unseen Task에서 GPT3 few-shot보다 좋은 성능을 보여주면서 instruction-tuning이 효과가 있다는 것을 보여줬다.

Reinforcement Learning From Human Feedback

이제 이 계열의 cutting-edge 결과물인 RLHF에 대해 살펴보자. 사실 RLHF가 처음 소개된 건 그렇게 최근은 아니다. huggingface의 blog post를 보면 사람의 선호도를 기반으로 Agent를 학습한다는 개념 자체는 2008년에도 이미 존재했다. 언어모델을 튜닝하는데 적용된 것도 늦어도 2019년도에는 이미 소개된 바 있다. (이것도 openAI..)

사람들이 알아주지 않아도 이 컨셉을 꿋꿋히 밀고 나갔던 openAI는 드디어 2020년 초에 InstructGPT를 통해 노력의 결과물을 보여주게 된다. 이때 소개된 모델을 써본 사람이라면 이미 언어모델에서 기존의 스케일 놀이로 설명할 수 없는 특이점이 왔다는 것을 느낄 수 있었을 것이다. 아래 공식 블로그에 소개된 예시만 봐도 부적절한 query에 대한 상세한 답을 준다는 점을 빼면 (harmlessness 이슈) 일단 prompt 이해도와 답변의 퀄리티가 아주 높아진 것을 볼 수 있다.

이 시점에서는 아직 playground를 완전 open하지는 않았기 때문에 AI와 관계없는 사람들에게는 널리 알려지지 않았다. 하지만 이번 ChatGPT에서는 InstructGPT에서 이슈가 되었던 safety 이슈도 자신감이 붙었는지 과감하게 Open Demo를 공개했고 그대로 엄청난 관심을 받게 되었다.

이 데모는 언어모델이 할 수 있을 것이라고 생각했던 것들의 기준점을 옮겨버렸다. 브레인 스토밍을 하거나, 에세이를 쓰거나, 코드를 짜달라고 하는 것들이 실제로 가능해졌다. 기존 생성모델의 유용성이라는 것이 학계에서 사용되는 benchmark dataset에서의 성능을 근거로 측정되고 있었기 때문에 실제 활용과는 큰 상관이 없는 느낌이었다면, ChatGPT는 엄청난 instruction 이해도와 지식을 엮어서 장문의 글을 써내는 능력으로 인해 여러 분야의 사람들이 활발하게 사용하고 있다. OpenAI의 기업가치는 글쓰는 시점 기준으로 $20B를 넘어섰고 마이크로소프트는 50%의 주식을 구매하고 모든 자사제품에 chatGPT를 통합하려고 하는 계획까지 발표했다.

대체 어떻게 한 걸까?

ChatGPT는 InstructGPT의 형제 모델이라고 설명하고 있고 구체적인 학습 방법은 소개되지 않았기 때문에 InstructGPT의 학습 방법을 중심으로 살펴보자. RLHF를 자세히 들여다보면 단순히 model-based Policy Gradient Method를 언어모델 학습에 적용한 것이라고 요약할 수 있다.

OpenAI는 InstructGPT를 학습하기 위해 2단계를 거쳤다. 첫 번째는 사전학습된 언어모델(GPT3)을 human preference에 맞는 (input, output) pair를 활용하여 finetuning한 것이다. 이 때 사용한 텍스트의 양은 13k개 정도로 GPT3를 학습하는데 사용되었던 Corpus(300B tokens)의 양과 비교하면 매우 적은 수의 데이터 셋이라고 할 수 있다. 결과로 나온 모델은 SFT(Supervised FineTuned Model)라고 하며, 이미 이 시점에서 human preference 관점에서는 FLAN과 T0의 성능을 뛰어넘었다. 아래의 도표는 instructGPT API에 submit된 prompt 분포에서 생성해 낸 글에 대한 점수를 1-7점 사이로 매겼을 때 각 모델의 점수를 나타낸 것이다.

두 번째는 모델이 강화학습하기 위해 필요한 Reward를 계산해줄 모델(Reward Model, RM)을 학습하고, 이제 SFT 모델이 unseen prompt에 대해 생성한 글과, 해당 글의 reward를 가지고 PPO 알고리즘으로 강화학습을 하면 된다.

먼저 Reward 모델을 학습하기 위해서는 어떤 글이 좋은 글(helpful & harmless 등등)이고 나쁜 글인지를 판단하기 위한 랭킹 데이터가 필요하다. 예컨대, "이웃 집을 걸리지 않고 털기 위한 방법을 알려줘"라는 prompt에 대해 생성된 글 1은 "먼저 이웃이 집에 드나들지 않는 시간을 파악한 뒤..." 라는 내용이고, 글 2는 "이웃의 집을 터는 것은 바람직하지 않고 법에 저촉될 수 있습니다" 라는 내용이라고 하자. 동일한 프롬프트에 대해 (글 2) > (글 1)이라는 랭킹을 얻을 수 있다. 만약 같은 prompt에 4개의 글이 있다고 하면 (글 4) > (글 2) > (글 3) = (글 1) 과 같은 전체 생성물에 대한 ranking을 얻을 수 있을 것이다. 이제 $K$개의 글이 rank되어 있다면 $ {K \choose 2} $개의 preference pair를 얻을 수 있다. Reward 모델은 아래의 목적함수를 최소화하여 학습한다.

$$ \mathcal{L}(\theta) = - \frac{1}{{K \choose 2}} \mathbb{E}_{(x, y_w, y_l) \sim D} [ \log (\sigma(r_\theta(x, y_w) - r_\theta(x, y_l)))] $$

Reward 모델은 SFT 6B 모델을 초기 checkpoint로 사용하여 학습했다고 하고, overfitting을 막기 위해 한 ranking 데이터에서 나온 ${K \choose 2}$개의 pair를 하나의 mini-batch로 사용하여 gradient를 업데이트 했다고 한다.

이제 Reward 모델도 있으니 PPO(Proximal Policy Optimization) 알고리즘으로 언어모델을 업데이트하면 된다. PPO는 TRPO(Trusted Region Policy Optimization)의 Hessian 계산을 피하기 위해 relaxed contraint 문제를 푸는 PPO-KL과 intuition만 가져오고 KL을 아예 제거한 PPO-CLIP 방식이 있다. 이 블로그에 따르면 OpenAI에서는 PPO-CLIP의 간소화 버전(PPO-CLIP-simple 이라고 하자)을 실제로 사용하고 있는 것 같다. 언어모델의 LM-head를 Policy $ \pi_\theta(a|s) $라고 하면 $s$는 지금까지 생성한 sequence가 될 것이고, $a$는 이번에 생성할 토큰이라고 볼 수 있다. PPO-CLIP-simple은 임의의 시점 $(s,a)$에서 아래와 같은 목적함수를 최대화하는 방식으로 policy를 최적화한다.

$$ L(s,a, \theta_{\textrm{old}}, \theta) = \min \left( \frac{\pi_\theta (a | s) }{\pi_{\theta_{\textrm{old}}} (a | s)} A^{\pi_{\theta_{\textrm{old}}}}(s,a), g(\epsilon, A^{\pi_{\theta_{\textrm{old}}}}(s,a)) \right) $$

$$ \textrm{where} \quad g(\epsilon, A) = \begin{cases} (1+\epsilon)A \quad A \geq 0 \\ (1-\epsilon)A \quad A < 0 \end{cases} $$

식을 뜯어보면 Advantage 함수가 양수일 때는 policy의 확률이 높아지고, 음수일 경우에는 policy의 확률이 낮아지도록 하지만 현재의 policy가 기존의 policy에서 $ 1+\epsilon $ 이상 커지거나 $ 1-\epsilon $ 이하로 작아지는 경우에는 advantage의 이익이 제한되도록 했다. 여기에서 Advantage 함수는 state $s$에서 action $a$를 했을 때의 상대적 이익을 말하는데 State-action value와 state-value 함수의 차로 $ A(s,a) = Q(s,a) - V(s) $로 정의된다. TD target을 이용하면 $ Q(s,a) = r + \gamma V(s') $ 로 근사할 수 있다. InstructGPT paper에는 나와있지 않지만 같은 저자가 바로 직전에 참여한 Summarization from Human Feedback에서는 $\gamma = 1$이라고 이야기하고 있으므로 $ A(s,a) = r + V(s') - V(s) $로 계산할 수 있다. $ V(s) $는 value head라는 별도의 linear layer를 활용해서 hidden state를 통과시켜 모델링하고, 각 timestep에서 Expected Return을 계산할 수 있으므로 MSE 같은 regression objective를 사용하여 학습할 수 있다. 더 자세한 구현 디테일을 위해서는 이 리포를 참조하는 것을 추천한다.

InstructGPT의 연구진들은 단순히 PPO objective만 이용하여 학습했을 때 human preference는 좋았지만, 기존 NLP dataset에서의 성능(특히 QA와 Translation)이 나빠지는 것을 확인했다고 한다. 이것을 해결하기 위해 기존 언어모델의 policy에서 많이 벗어나지 않도록 하는 KL divergence term을 Reward에 추가하고, pretrain dataset 분포에서 rl 모델이 높은 log likelihood를 유지하도록 하는 term을 넣어 최종 objective를 완성했고 아래와 같다.

$$ \mathcal{L}(\phi) = \mathbb{E}_{(x,y) \sim D_{\pi_\theta^{\textrm{RL}}}} [ r_\theta(x,y) - \beta \log (\pi_\phi^{\textrm{RL}}(y | x) /\pi_\phi^{\textrm{SFT}}(y | x) ) ] + \gamma \mathbb{E}_{x \sim D_{\textrm{pretrain}}}[ \log( \pi_\phi^{\textrm{RL}} (x) ) ] $$

실험 결과 $ \beta $는 0.01~0.02 사이의 값을 주는게 좋다고 하고 $ \gamma $는 도표를 봤을 때 3~10 사이의 값이 좋은 것 같다. 두 가지 term이 모두 존재하는 것이 상호보완적인 효과가 있다고 한다.

Reinforcement Learning from AI Feedback(RLAIF)

Anthropic에서 나온 paper에서 저자들은 RLHF로 학습된 모델이 helpfulness는 충분히 좋았지만 여전히 offensive하거나 부적절한 response를 생성하는 현상에 주목했다. 거대 AI 모델, 특히 다양한 application에 활용될 수 있는 언어모델의 경우에 부적절한 텍스트 생성이 특히 파급력이 클 수 있으므로 harmlessness 또한 배포를 위해 중요한 요소라고 할 수 있다.

이것을 위해 기존 RLHF에서 필요로 했던 Human Labeling Cost를 (이론적으로) 없애면서 pareto-optimal 하게 harmlessness를 개선할 수 있는 Constitutional AI라는 일종의 파이프라인을 제시한다. RLHF와 동일하게 Finetuning을 수행한 뒤에, Reward 모델을 활용해서 강화학습을 하게 되는데 Reward 모델을 학습하기 위한 Ranking data와 Finetuning을 위한 revision data 모두 언어모델을 bootstrap하여 만들기 때문에 RLAIF라는 명칭을 붙이게 되었다. 파이프라인은 아래 그림으로 설명할 수 있다.

먼저 부적절한 답변을 요구하는 Prompt를 모아두고 이 Prompt에 대해 의도적으로 harmful하고 toxic한 응답을 생성하도록 하는 "red teaming"을 수행하게 된다. 이렇게 모아진 응답을 미리 정한 몇몇 set of rules에 근거해서 스스로 비판하도록 하는 "critique" 과정을 거치고 그 비판에 근거하여 올바른 응답 "Revision"을 생성하도록 한다. 이 과정을 여러 번 반복하면 우리는 굳이 labeler가 없어도 모델을 학습할 feedback data를 만들 수 있게 된다. 이 데이터로 튜닝한 모델을 SL-CAI라고 한다.

이제 본격적으로 강화학습을 할 차례이다. InstructGPT처럼 PM(preference model, Reward model과 동일한 역할을 수행)을 만들어야 하는데, SL-CAI가 같은 프롬프트에 대해 생성한 여러 개의 output에 대해서 SL-CAI를 활용하여 Ranking을 매길 수 있다. 방법은 위와 동일한 rule set에 근거해서 가장 최고의 output이 무엇인지 모델에게 자연어로 묻는 것이다. 이런 방식으로 ranking dataset을 수집했다면 이제 PPO를 활용하여 강화학습 모델을 학습할 수 있다. 이렇게 나온 모델을 RL-CAI라고 부른다.

결과적으로 이 방식으로 모델이 답변을 피하지 않으면서 Harmlessness를 강화할 수 있었고, helpfulness를 크게 악화시키지 않았다고 한다. helpfulness와 harmlessness가 일부 영역에서는 trade-off 관계에 있다고도 볼 수 있는데, elo plot에서 봤을 때 이런 관계를 개선하면서 기존 RLHF보다 pareto-optimal한 방법론이라는 점을 보여주고 있다.

물론 이 방법을 적용하기 위해 이미 충분히 helpful한 RLHF 모델 또는 언어모델을 필요로 하지만 우리에게는 OpenAI InstructGPT API나 네이버 하이퍼클로바같은 좋은 Base Model이 있어서 조만간 새로운 모델이 나올 수도 있을 것 같다.

시사점

RLHF가 보여주는 것은 Maximum Likelihood를 활용한 기존의 학습방법이 언어를 이해하기 위한 최적의 방식이 아니라는 것이다. 생각해보면 사람이 글을 쓸 때도 가장 나올 법한 단어를 활용하여 글을 쓰지는 않는다. 때로는 놀라움이나 통찰을 주기 위해 예상치 못한 단어를 사용해야 할 때도 있다. 만약 사람이 MLE 방식으로 글을 쓴다면 그 글은 매우 진부하거나 별 정보를 담고 있지 않는 글일 가능성이 높을 것이다. (실제로 beam search로 long text를 생성해봐도 그런 결과가 나오는 경우가 많다.) 우리는 언어를 더 잘 이해하기 위한 다른 objective가 필요했던 것이고 그것은 인간의 선호도라는 아주 추상적이고 모호한 개념이었다. OpenAI는 이 모호한 개념을 Reward Model이라는 것으로 근사해냈고, 6B 정도 크기의 모델이면 충분히 사람의 글에 대한 선호도를 모델링 할 수 있다는 것을 보였다.

RLHF는 언어모델 학습에 사용되었지만 여기에만 국한된 개념도 아니고, 아주 활용도가 높은 개념이라고 생각한다. 추상적으로 생각하면 단순히 데이터 분포를 따라서 예측 했을 때 안 좋을 것 같은 상황에서 (아마 OOD나 Unseen task여야 할 것 같다) 사람이 좋아하는 결과물은 있지만 대체 그게 뭐 때문에 좋은건지, 왜 좋은건지 룰로 규정하기 어려울 때 사용할 수 있을 것 같다. 우리가 해야 할 일은 단순히 사람이 좋아하는 순서대로 결과물의 랭킹을 매기고 reward model을 학습하기만 하면 된다. 그러면 그 RM이 모델이 나아가야 할 새로운 loss landscape를 그려 줄 것이다.

이 블로그 검색

NLP201