3. REINFORCE — 샘플로 추정하고, 베이스라인을 빼다 — 생성모델의 강화학습

Chapter 3: REINFORCE — 샘플로 추정하고, 베이스라인을 빼다

의문

답이 셋뿐인 장난감 문제를 떠올려 보자. "대한민국의 수도는?"에 모델이 “서울입니다”(1점), “부산입니다”(0점), “서울이요”(0.6점) 가운데 하나로만 답한다. 평균 점수를 올리려고 샘플 하나씩 뽑아 걸으면 학습 곡선이 요동친다. 점수마다 똑같이 5를 더해 주면 요동은 더 심해진다. 모든 답에 똑같이 얹은 점수가 어느 답이 나은지를 바꿀 리 없고, 평균을 올리는 방향도 그대로인데 왜 그럴까? 그리고 그 요동은 어떻게 줄이는가?

평균 점수를 올리는 방향은 정책 그래디언트로 적힌다.

\nabla_\theta \textcolor{#d62728}{J} = \mathbb{E}_{\textcolor{#1c9c60}{y} \sim \textcolor{#1565c0}{\pi_\theta}}\big[\, \textcolor{#d9670b}{R}(\textcolor{#1c9c60}{y})\, \nabla_\theta \log \textcolor{#1565c0}{\pi_\theta}(\textcolor{#1c9c60}{y}) \,\big]

\small\begin{array}{ll} \textcolor{#d62728}{J} & \text{목적함수: 평균 점수} \\ \textcolor{#d9670b}{R}(\textcolor{#1c9c60}{y}) & \text{응답 }\textcolor{#1c9c60}{y}\text{의 점수} \\ \textcolor{#1565c0}{\pi_\theta} & \text{학습 중인 정책 (프롬프트를 받아 답을 고르는 언어모델)} \end{array}

정확한 식이지만 언어모델에서는 계산할 수 없다. 기대값을 구하려면 가능한 모든 응답을 더해야 하는데, 어휘가 1만 개인 모델이 20토큰짜리 답만 쓴다고 해도 가짓수가 $10000^{20} = 10^{80}$ 이다. 관측 가능한 우주의 원자 수로 흔히 꼽는 값(약 $10^{80}$ )과 맞먹고, 실제 답은 이보다 훨씬 길다. 그러니 할 수 있는 일은 하나뿐이다: 몇 개를 뽑아서 평균을 낸다.

몬테카를로 추정: 기대값 대신 샘플 평균

뽑은 몇 개의 평균이 정말 참 그래디언트를 대신할 수 있을까? 대신할 수 있다면, 한 걸음은 실제로 어떻게 돌아가고, 무엇이 비싼가?

샘플 평균으로 걷기

현재 모델로 응답 $\textcolor{#5f970c}{N}$ 개를 뽑고, 각자의 $\textcolor{#d9670b}{R} \cdot \nabla \log \textcolor{#1565c0}{\pi}$ 를 평균낸다.

\textcolor{#cc00ff}{\hat g} = \frac{1}{\textcolor{#5f970c}{N}} \sum_{i=1}^{\textcolor{#5f970c}{N}} \textcolor{#d9670b}{R}(\textcolor{#1c9c60}{y_i})\, \nabla_\theta \log \textcolor{#1565c0}{\pi_\theta}(\textcolor{#1c9c60}{y_i}), \qquad \textcolor{#1c9c60}{y_i} \sim \textcolor{#1565c0}{\pi_\theta}

\small\begin{array}{ll} \textcolor{#cc00ff}{\hat g} & \text{그래디언트 }\nabla_\theta \textcolor{#d62728}{J}\text{의 추정치 (샘플 평균)} \\ \textcolor{#5f970c}{N} & \text{뽑은 응답 수 (배치 크기)} \\ \textcolor{#d9670b}{R}(\textcolor{#1c9c60}{y_i}) & \text{}i\text{번째 응답의 점수} \\ \textcolor{#1c9c60}{y_i} \sim \textcolor{#1565c0}{\pi_\theta} & \text{지금의 정책이 뽑은 }i\text{번째 응답} \end{array}

각 항의 기대값이 $\nabla \textcolor{#d62728}{J}$ 이므로 평균의 기대값도 $\nabla \textcolor{#d62728}{J}$ 다. 이런 추정을 불편(unbiased, 치우침 없음)하다고 한다. 여러 번 뽑아 평균내면 참값으로 수렴한다는 뜻이다. 미니배치 확률적 경사하강법(SGD)이 전체 데이터의 그래디언트 대신 배치 평균을 쓰는 것과 같은 원리다.

실제 학습 루프

REINFORCE 한 걸음은 여섯 단계다.

단계	하는 일	비용
① 생성	현재 모델로 프롬프트마다 응답을 뽑는다 (샘플링 온도 1: 모델이 준 확률 그대로 뽑는다)	모델을 돌려 답을 뽑는 계산 — 가장 느린 단계
② 채점	각 응답에 보상 $\textcolor{#d9670b}{R_i}$ 를 매긴다 (검증기, 보상 모델, 사람)	채점기에 따라 다름
③ 로그확률	응답의 토큰 로그확률을 더해 $\log \textcolor{#1565c0}{\pi_\theta}(\textcolor{#1c9c60}{y_i})$ 를 구한다	순전파 1회
④ 대리 손실(기울기만 빌려 쓰는 손실)	$\textcolor{#d62728}{\mathcal{L}} = -\frac{1}{\textcolor{#5f970c}{N}}\sum_i \textcolor{#d9670b}{R_i}\log \textcolor{#1565c0}{\pi_\theta}(\textcolor{#1c9c60}{y_i})$	무시할 만함
⑤ 역전파·갱신	$\nabla \textcolor{#d62728}{\mathcal{L}}$ 로 파라미터를 한 걸음 옮긴다	역전파 1회
⑥ 폐기	뽑았던 응답을 버리고 ①로 돌아간다	—

샘플링 온도를 1로 두는 까닭은 위 식의 조건 $\textcolor{#1c9c60}{y_i} \sim \textcolor{#1565c0}{\pi_\theta}$ 에 있다. 샘플링 온도를 낮춰 뽑으면 샘플이 $\textcolor{#1565c0}{\pi_\theta}$ 가 아닌 더 뾰족한 분포에서 나오고, 그 샘플의 평균은 더 이상 $\nabla \textcolor{#d62728}{J}$ 를 겨누지 않는다. 코드로 쓰면 SFT(정답을 그대로 따라 쓰게 하는 지도 미세조정)에서 몇 줄 바뀔 뿐이다.

responses = policy.generate(prompts, do_sample=True)          # ① 생성
rewards = score(prompts, responses)                            # ② 채점, shape [N]
logp = policy.logprob(prompts, responses).sum(dim=-1)          # ③ 토큰 로그확률의 합, shape [N]
loss = -(rewards.detach() * logp).mean()                       # ④ 대리 손실
loss.backward(); optimizer.step(); optimizer.zero_grad()       # ⑤

⑥이 중요하다. 로그 미분 트릭은 "지금의 $\textcolor{#1565c0}{\pi_\theta}$ 에서 뽑은 샘플"에만 성립하므로, 파라미터를 한 번 갱신하면 방금 뽑은 응답은 남의 응답이 된다. 비싸게 생성한 샘플을 그래디언트 한 걸음에 쓰고 버린다. 이 낭비가 REINFORCE의 큰 약점이다.

왜 이렇게 흔들리는가

첫머리의 장난감 문제로 재 보자. 처음 소프트맥스 정책이 세 답 “서울입니다”(1점), “부산입니다”(0점), “서울이요”(0.6점)에 약 0.14, 0.61, 0.25의 확률을 준다. 이때 "서울입니다"의 로짓(소프트맥스를 거쳐 확률이 되기 전의 값)에 대한 참 그래디언트는 $+0.098$ 이다. 샘플 하나로 추정하면 뽑힌 답에 따라 값이 달라진다. 이 추정치의 표준편차를 계산해 보면:

설정	추정치 평균	표준편차	신호 대 잡음비	잡음보다 신호가 커지려면 필요한 샘플 수
보상 그대로	0.098	0.307	0.32	10개
보상에 +5	0.098	2.024	0.05	430개

(신호 대 잡음비는 평균 ÷ 표준편차다. 추정치가 가리키는 값이 흔들림보다 얼마나 큰지를 잰다. 마지막 열은 $(\text{표준편차}/\text{평균})^2$ , 즉 $\textcolor{#5f970c}{N}$ 개를 평균내 표준편차가 $1/\sqrt{\textcolor{#5f970c}{N}}$ 로 줄었을 때 평균과 같아지는 $\textcolor{#5f970c}{N}$ 이다.)

평균은 두 줄 모두 똑같다. 달라지는 것은 흔들림의 폭, 즉 분산이다. 보상에 5를 더했을 뿐인데 같은 정확도를 얻으려면 샘플이 43배 필요하다. 샘플 하나가 LLM의 응답 생성 한 번이라는 걸 생각하면, 분산은 곧 GPU 비용이다.

샘플 하나의 추정치가 실제로 어떤 값들인지 펼쳐 보면 까닭이 보인다.

왜 상수가 분산을 키우는가? 추정치는 $\textcolor{#d9670b}{R} \cdot \nabla \log \textcolor{#1565c0}{\pi}(\textcolor{#1c9c60}{y})$ 다. $\nabla \log \textcolor{#1565c0}{\pi}$ 는 어떤 답이 뽑혔느냐에 따라 방향이 크게 바뀌는 벡터이고, $\textcolor{#d9670b}{R}$ 은 그 벡터의 길이를 늘린다. 모든 답의 $\textcolor{#d9670b}{R}$ 이 5 근처라면, 매 샘플이 "뽑힌 답을 5배 세게 올려라"라고 외친다. 좋은 답이 뽑히든 나쁜 답이 뽑히든. 그림 아래쪽에서 "부산입니다"가 뽑힌 샘플이 "서울입니다"의 로짓을 −0.68만큼 내리라고 외치는 것이 그 외침이다. 이 외침들은 평균내면 상쇄되지만(상수 보상의 기대 그래디언트는 0이다 — 확률의 합이 늘 1이므로), 하나하나는 크다.

트릭의 뿌리 · 전체를 셀 수 없으면 일부를 뽑아 센다

응답 공간 전체의 합은 계산할 수 없었다. 그래서 전체 대신 몇 개를 뽑아 평균냈다. 정확한 답을 내려놓고, 그 대신 「뽑을수록 참값에 다가가는」 답을 얻는다. 치르는 값은 흔들림이다. 그 뒤의 일은 모두 흔들림과의 싸움이 된다. 더 많이 뽑거나, 흔들림 자체를 줄이는 요령을 찾는다.

1802년 라플라스(Pierre-Simon Laplace)는 프랑스의 인구를 알고 싶었다. 전국 인구조사를 한 적이 없었고, 한 사람씩 셀 여력도 없었다. 그는 교구 30곳만 골라 주민 2,037,615명을 세고, 같은 교구의 세 해치 세례 기록에서 한 해 평균 세례 수(71,866.333건)를 구했다. 세례 한 건에 주민 28.35명꼴이다. 세례 기록은 나라 전체의 것이 있었으므로, 전국의 세례 수에 이 비율을 곱해 인구를 어림했다. 일부에서 잰 비율이 전체에서도 같다고 믿고 건넌 것이다.

1908년, 더블린의 기네스 양조장에서 일하던 고셋(William Sealy Gosset)은 「스튜던트(Student)」라는 필명으로, 측정값이 몇 개뿐일 때 평균을 어떻게 믿을지를 다룬 논문을 냈다. 논문에 따르면 그는 식으로 답을 얻기 전에 먼저 뽑아서 알아보려 했다. 범죄자 3000명의 키와 왼손 가운뎃손가락 길이를 판지 3000장에 옮겨 적고, 잘 섞은 뒤 한 장씩 뽑아 적어 내려가, 넷씩 끊어 표본 750개를 만들었다. 표본마다 평균과 표준편차를 구해, 식이 내놓은 분포와 견주었다.

같은 역사에 경고도 하나 있다. 1901년 이탈리아의 라차리니(Mario Lazzarini)는 바늘을 3408번 던져 원주율을 355/113 = 3.1415929로 얻었다고 보고했다. 오차가 1000만분의 3보다 작다. 그런데 3408은 바로 355/113이 나올 수 있는 던진 횟수이고, 중간 결과도 지나치게 좋다. 그래서 좋은 값이 나온 순간 멈췄으리라는 의심을 받는다. 샘플 평균이 치우침 없는 것은 뽑은 것을 고르지 않고 다 쓸 때뿐이다.

알아보는 법: 합·적분·기대값 앞에 「계산할 수 없다」가 붙으면, 먼저 그 분포에서 뽑을 수 있는지 묻는다. 뽑을 수 있다면 문제는 「얼마나 흔들리나」로 바뀐다.

문제 1 — 몇 명에게 물어야 하나

동아리 회장이 MT를 바다로 갈지 산으로 갈지 정하려고 회원에게 한 명씩 묻는다. 바다를 고른 답은 +1, 산을 고른 답은 −1로 적고 그 평균으로 판단한다. 회원 전체로는 60%가 바다파다. (가) 아무나 한 명에게 물었을 때 그 답의 평균과 표준편차는? (나) 몇 명에게 물어 평균을 내야 평균이 흔들림(표준편차 ÷ √명수)만큼 커지는가? (다) 바다파가 55%라면?