보상은 어떻게 토큰의 생성 확률을 바꾸는가

앞 글에서는 LLM이 토큰을 하나씩 선택해 응답을 만들고, 완성된 응답에 보상을 받는 과정을 살펴봤습니다. 이제 정답 응답에 +1이라는 점수를 얻었다고 해보겠습니다. 이 숫자만으로는 어떤 모델 가중치를 어느 방향으로 바꿔야 할지 알 수 없습니다. 채점 결과와 모델의 계산을 연결해야 합니다.

먼저 기록된 토큰의 현재 확률을 다시 계산하고, 보상으로부터 이번 선택을 얼마나 더 선호할지 나타내는 어드밴티지를 구합니다. 두 값을 손실에서 연결하는 기본 경로를 이해한 뒤, 같은 데이터를 다시 학습할 때 사용하는 확률비와 clipping을 추가하겠습니다. 마지막에는 학습의 출발점과 너무 달라지지 않도록 하는 기준 정책의 역할을 구별합니다.

이미 선택한 토큰의 확률을 다시 계산하기

학습 데이터에는 어떤 문맥에서 어떤 토큰을 선택했는지가 기록되어 있습니다. 토큰 위치 t의 문맥을 sₜ, 그때 선택한 토큰을 xₜ라고 쓰겠습니다. θ는 현재 학습하는 모델의 가중치입니다.

현재 모델에 같은 문맥 sₜ를 넣으면 다음 토큰 후보들의 확률을 계산할 수 있습니다. 그중 기록에 있는 xₜ의 확률을 꺼낸 값이 pθ(xₜ | sₜ)입니다. 세로선은 “이 문맥이 주어졌을 때”라는 뜻입니다. 여기서 하는 일은 새 응답을 생성하는 것이 아니라, 기록된 선택을 현재 모델이 얼마나 그럴듯하게 보는지 다시 계산하는 것입니다.

프롬프트를 P라고 쓰겠습니다. 프롬프트 P에 대한 응답이 설명용 토큰 u → b → d로 기록되어 있다고 해보겠습니다. u를 평가할 때의 입력은 P, b를 평가할 때는 P + u, d를 평가할 때는 P + u + b입니다. 서로 다른 모델 세 개가 아니라 같은 모델의 서로 다른 예측 위치입니다.

같은 Transformer가 세 예측 위치에서 P, P와 u, P와 u와 b를 문맥으로 받는다. 문맥 표현이 LM head와 softmax를 통과한 뒤 기록된 선택 u, b, d의 확률 0.5, 0.2, 0.4를 각각 꺼낸다.

각 예측 위치에서 Transformer는 문맥을 반영한 히든 스테이트 hₜ를 계산합니다. LM head가 이를 어휘의 각 토큰에 대한 점수로 바꾸고, softmax가 확률로 변환합니다. 그중 실제로 선택했던 토큰의 확률을 꺼냅니다. 따라서 두 번째 위치에서는 다른 후보의 확률이 더 높더라도 pθ(b | P, u) = 0.2를 사용합니다. 그림의 ‘기타’는 다른 어휘 토큰을 합친 값이며, 숫자는 설명용입니다.

생성할 때는 토큰을 순서대로 뽑아야 합니다. 학습할 때는 기록된 전체 시퀀스가 있으므로 causal attention mask를 사용해 여러 위치를 한꺼번에 평가할 수 있습니다. b를 예측하는 위치는 b 자신이나 뒤의 d를 볼 수 없습니다. 학습 forward는 역전파에 필요한 계산을 남기지만, 과거에 토큰을 뽑았던 선택 자체를 미분하지는 않습니다.

예를 들어 선택한 토큰의 현재 확률이 0.5라면 자연로그를 취한 값은 약 −0.693입니다. 확률이 0.6이면 로그확률은 약 −0.511입니다. 확률이 커질수록 로그확률도 커집니다. 로그확률 자체가 정답 여부나 보상은 아닙니다. 가중치를 바꾸면 값이 어떻게 달라지는지 미분할 수 있는 모델 쪽의 계산 결과입니다.

보상과 비교 기준으로 어드밴티지 구하기

이번에는 평가 쪽 입력을 보겠습니다. 선택 이후에 받은 보상들을 모은 값을 return이라고 합니다. 마지막에만 보상 +1을 받고 앞뒤 보상을 동일하게 취급하는 간단한 예에서는, 그 응답 안의 선택들에 대해 관측한 return이 1입니다.

그런데 같은 1점이라도 원래 쉽게 얻던 결과인지, 예상보다 좋은 결과인지에 따라 학습 신호가 달라질 수 있습니다. 이를 비교하기 위해 기준값을 둡니다. 가장 간단한 형태는 다음과 같습니다.

어드밴티지 = 관측한 return − 비교 기준값
             1          −     0.4      = +0.6

어드밴티지(advantage)는 선택의 결과가 기준보다 얼마나 좋았는지를 나타내는 추정값입니다. 이 예의 +0.6은 기준보다 좋은 결과였다는 뜻입니다. 같은 기준에서 return이 0이었다면 −0.4가 되어, 그 선택을 덜 선호하게 만드는 신호가 됩니다.

비교 기준은 현재 문맥의 기대 보상을 예측하는 critic에서 얻을 수도 있고, 같은 질문에 생성한 여러 응답의 점수에서 얻을 수도 있습니다. 전자는 PPO에서 흔히 사용하는 방식이고, 후자는 GRPO의 중심 아이디어입니다. 구체적인 계산은 다음 글에서 다루겠습니다. 이 단계에서는 보상과 어드밴티지가 서로 다른 값이라는 점이 핵심입니다.

어드밴티지를 토큰 위치에 붙였다고 해서 그 토큰의 인과적 기여도를 정확히 밝혀냈다는 뜻은 아닙니다. 관측한 결과와 비교 기준으로 학습에 사용할 신호를 추정한 것입니다. RLHF Book — Policy Gradient Algorithms

확률을 바꾸는 방향을 손실로 표현하기

토큰 위치 t의 어드밴티지를 Aₜ, 이 선택에 대한 손실을 L이라고 쓰겠습니다. 두 입력을 연결하는 기본 정책 그래디언트의 손실은 다음과 같습니다. 이후 설명을 위한 단순한 형태이며, 아직 PPO의 확률비나 clipping을 넣은 식은 아닙니다.

L = - A_{t} \log p_{θ} (x_{t} | s_{t})

학습기는 이 L을 작게 만드는 방향으로 가중치를 바꿉니다. Aₜ가 양수라면 로그확률이 커질수록 L이 작아집니다. 따라서 선택한 토큰의 확률을 높이는 방향으로 신호가 전달됩니다. Aₜ가 음수라면 반대로 그 확률을 낮추는 방향입니다.

Aₜ = +0.6일 때를 계산해보겠습니다. 확률 0.5에서는 L이 약 0.416이고, 확률 0.6에서는 약 0.306입니다. 좋은 선택의 확률이 커진 쪽에서 손실이 작아집니다. 손실이라는 이름은 이처럼 어떤 가중치 변화를 선호하는지 나타내는 최적화 기준이라는 뜻입니다. 모든 손실이 오답 개수처럼 해석되거나 항상 양수여야 하는 것은 아닙니다.

로그확률을 사용하는 이유도 단지 숫자를 작게 만들기 위해서가 아닙니다. 정책이 생성한 경험의 기대 보상을 미분하면, 보상이나 어드밴티지로 가중한 로그확률의 기울기 형태로 표현할 수 있습니다. 이 식은 그 기울기를 자동 미분으로 구하기 위한 경로를 제공합니다. PPO §2.1

보상에서 구한 return 1에서 기준값 0.4를 빼 어드밴티지 +0.6을 얻는다. 기록된 문맥과 선택 토큰을 현재 모델에 넣어 로그확률을 계산한다. 고정한 어드밴티지와 로그확률로 손실을 계산하고 모델 파라미터에 역전파한다.

그림 2는 그림 1의 첫 토큰 u를 따라갑니다. 로그확률을 잠시 z라고 쓰면, A₁ = +0.6일 때 손실은 L₁ = −0.6z입니다. z가 0.01 커지면 손실은 0.006 작아집니다. 따라서 z에 대한 손실의 기울기는 −0.6입니다. 그림의 첫 역전파 블록은 이 계산을 나타냅니다.

여기서 −0.6은 로그확률 z에 대한 미분값입니다. 확률 p 자체에 대한 미분은 −0.6/p이므로 p = 0.5에서는 −1.2가 됩니다. 모델 가중치에 대한 기울기는 이 값들을 그대로 복사한 것이 아니라, 로그확률을 계산한 경로를 따라 연쇄적으로 미분한 결과입니다. 역전파는 softmax, LM head, Transformer로 이어지고, optimizer는 학습 가능한 가중치를 갱신합니다. 확률표의 숫자를 직접 고쳐 쓰는 것이 아닙니다. 바뀐 가중치로 다음 forward를 수행할 때 새로운 확률 분포가 계산됩니다.

Aₜ는 이 정책 업데이트에서 고정해 사용합니다. critic을 학습하는 경우에도, 정책 손실을 통해 Aₜ 자체를 마음대로 키우게 만드는 것이 아니라 critic의 예측 오차를 별도로 학습합니다.

또한 “확률을 높이는 방향”은 이 토큰 항이 주는 신호를 설명한 것입니다. 실제로는 여러 토큰과 응답의 손실을 모으고, 같은 가중치가 많은 문맥의 출력에 관여합니다. 따라서 전체 업데이트 뒤에 모든 양수 어드밴티지 토큰의 확률이 반드시 증가한다고 보장할 수는 없습니다.

같은 데이터를 다시 사용할 때의 확률비

새로운 응답을 만드는 데는 계산이 필요합니다. 어렵게 생성하고 채점한 데이터를 한 번의 가중치 갱신에만 쓰지 않고 여러 번 사용할 수 있다면 효율적일 수 있습니다. 다만 첫 업데이트가 끝난 순간부터 현재 모델은 그 데이터를 만들었던 모델과 달라집니다.

이를 비교하는 값이 importance ratio, 즉 확률비 rₜ입니다. 여기서는 생성 당시 정책을 old, 지금 학습하는 정책을 current라고 부르겠습니다. 앞에서 pθ라고 쓴 현재 모델의 확률을 아래에서는 pcurrent로 표시합니다.

r_{t} = \frac{p_{current} (x_{t} | s_{t})}{p_{old} (x_{t} | s_{t})}

생성 당시 확률이 0.5이고 현재 확률이 0.6이라면 rₜ는 1.2입니다. 같은 문맥에서 같은 토큰을 지금은 1.2배의 확률로 선택한다는 뜻입니다. 이 데이터를 다음 업데이트에서도 사용한다면 분모는 여전히 0.5입니다. 분모를 매번 현재 확률로 바꾸면 데이터가 만들어진 기준과의 차이를 잃게 됩니다.

이제 계산을 설명하기 위해 작게 만들 손실 대신 크게 만들 목적함수를 쓰겠습니다. 학습에 사용할 손실은 이 목적함수에 음수를 붙인 값입니다. PPO에서는 확률비와 고정된 어드밴티지를 곱한 rₜAₜ를 최대화하는 항에서 출발해, 곧 설명할 clipping을 적용합니다. 앞의 −Aₜ log p 식에 rₜ를 그대로 하나 더 곱하는 방식으로 설명하면 실제 PPO의 미분과 달라집니다.

왜 rₜAₜ에도 로그확률과 같은 학습 방향이 들어 있을까요? 분모와 Aₜ를 고정하면 rₜAₜ의 기울기는 rₜAₜ × 현재 로그확률의 기울기가 됩니다. 현재 정책이 old와 같아 rₜ = 1인 시작점에서는 앞서 본 기본 정책 그래디언트와 연결됩니다. 다만 rₜ가 1.2라는 것이 실제 확률이나 가중치 변화량도 정확히 1.2배라는 뜻은 아닙니다. 모델의 기울기와 학습률, 다른 샘플도 결과에 관여합니다.

여기서는 생성과 학습이 같은 확률 계산을 수행한다고 가정했습니다. 같은 가중치라도 엔진이나 양자화 때문에 계산한 확률이 달라질 수 있다는 문제는 9편에서 분리해 살펴보겠습니다.

충분히 바뀐 선택을 더 밀지 않는 clipping

확률비를 넣는 것만으로 큰 업데이트가 자동으로 막히지는 않습니다. Aₜ가 양수일 때 rₜAₜ만 계속 키우면, 같은 데이터에 대해 이미 확률이 높아졌어도 더 높일 유인이 남습니다. PPO의 clipped objective는 이 유인을 제한합니다.

먼저 rₜ를 정해진 구간으로 잘라 계산한 값도 만듭니다. 아래 그림에서는 설명을 위해 구간을 0.8부터 1.2로 정했습니다. 표기 순서는 clip(입력값, 하한, 상한)입니다. 잘라낸 확률비에도 Aₜ를 곱한 뒤, 원래 항 rₜAₜ와 비교해 더 작은 쪽을 최대화합니다.

원래 항:       rₜ × Aₜ
잘라 계산한 항: clip(rₜ, 0.8, 1.2) × Aₜ
목적함수 J:    두 항 중 작은 값
최소화할 손실: −J

생성 당시 확률 0.5를 고정한다. 왼쪽은 A=+1일 때 확률 0.5→0.6→0.7과 r=1→1.2→1.4를, 오른쪽은 A=−1일 때 확률 0.5→0.4→0.3과 r=1→0.8→0.6을 연결한다. clip한 확률비에 A를 곱한 계산과 목적함수 J의 그래프를 보여준다.

Aₜ = +1이고 rₜ가 1.4라면 두 항은 1.4와 1.2입니다. 작은 값인 1.2를 사용합니다. rₜ를 1.5로 더 키워도 J는 1.2이므로, 이 항만 놓고 보면 더 높일 이득이 없습니다. Aₜ가 양수인 그래프에서 rₜ > 1.2 구간이 평평해진 이유입니다.

Aₜ = −1일 때는 반대쪽이 평평해집니다. rₜ = 0.6을 자르면 clip(0.6, 0.8, 1.2) = 0.8로, 이 값은 여전히 양수입니다. 여기에 Aₜ = −1을 곱하면 −0.8이 됩니다. 원래 항은 0.6 × (−1) = −0.6이므로, J는 두 항 중 작은 값인 −0.8입니다. 그래프의 세로축은 확률비가 아니라 어드밴티지까지 곱한 목적함수 J이기 때문에 오른쪽 값들이 음수입니다.

음수인 목적함수도 최대화할 수 있습니다. 예를 들어 rₜ가 1에서 0.8로 낮아지면 J는 −1에서 −0.8로 커집니다. 하지만 rₜ를 0.6으로 더 낮춰도 J는 −0.8입니다. 좋지 않았던 선택의 확률이 이미 충분히 낮아졌으므로, 더 누르는 유인을 없애는 것입니다.

그런데 왜 양쪽 그래프를 모두 0.8과 1.2에서 평평하게 만들지는 않을까요? 그림은 clip 결과 자체가 아니라, 원래 항과 clip한 항 중 작은 값을 고른 J를 그린 것이기 때문입니다. 원하는 방향과 반대로 확률이 움직인 경우에는 원래 항이 선택됩니다.

Aₜ = +1인데 rₜ = 0.6이면, 좋은 선택의 확률이 오히려 낮아진 상황입니다. 원래 항은 0.6, clip한 항은 0.8이므로 J = min(0.6, 0.8) = 0.6입니다. 이 구간까지 평평하게 만들면 확률을 다시 높이는 신호도 사라지므로, 기울기를 남겨둡니다.

Aₜ = −1인데 rₜ = 1.4이면, 나쁜 선택의 확률이 오히려 높아진 상황입니다. 원래 항은 −1.4, clip한 항은 1.2 × (−1) = −1.2이므로 J = min(−1.4, −1.2) = −1.4입니다. 이때도 확률을 다시 낮추는 신호를 남깁니다. 이처럼 좋은 방향으로 충분히 움직였을 때 더 밀어주는 유인은 끊고, 나쁜 방향으로 움직였을 때 되돌리는 신호는 유지하는 것이 min의 역할입니다. PPO §3

Clipping은 확률비 자체를 0.8~1.2 안에 강제로 가두지 않습니다. 손실의 모양을 바꾸는 장치입니다. 다른 샘플과 보조 손실을 통해 가중치가 움직일 수 있으므로 실제 확률비는 구간 밖에 있을 수 있습니다. 또한 온라인으로 데이터를 새로 만든다고 해서 모든 RL 알고리즘이 반드시 같은 clipping을 사용하는 것은 아닙니다. 데이터와 현재 정책의 관계를 관리하는 한 가지 대표적인 방법입니다.

기준 정책과의 차이를 따로 관리하기

데이터를 만든 뒤 얼마나 바뀌었는지와, 학습의 기준으로 삼고 싶은 모델에서 얼마나 멀어졌는지는 다른 질문입니다. 후자를 위해 reference policy, 즉 기준 정책을 둘 수 있습니다. 예를 들어 RL 시작 시점의 모델을 고정해 두고 현재 정책과 비교하는 방식입니다.

같은 입력 P를 각 모델에 넣는다. 왼쪽은 current와 old가 선택 토큰 u에 준 확률 0.6과 0.5를 나누어 r=1.2를 구한다. 오른쪽은 current와 reference의 네 후보 확률분포를 KL로 비교한다. 두 current 상자는 같은 모델이다.

그림 4의 왼쪽은 선택한 토큰 u의 확률이 생성 당시 0.5에서 현재 0.6으로 달라진 예입니다. 오른쪽은 어휘가 네 토큰뿐이라고 가정해 두 분포를 그렸습니다. 양쪽의 current는 같은 현재 모델이고, 상자는 물리적인 서버 수가 아니라 비교에 필요한 모델의 역할을 나타냅니다.

같은 문맥에서 현재 정책과 기준 정책이 만드는 확률분포의 차이를 KL divergence로 측정해 규제할 수 있습니다. 별도의 손실로 추가하는 경우에는 정책 손실에 β × KL을 더해 최소화합니다. β는 보상에 따라 정책을 바꾸려는 목표와 기준을 유지하려는 목표 사이의 비중을 정합니다.

기준 정책의 KL을 보상 쪽의 penalty로 넣어 어드밴티지에 반영하는 학습 구성도 있습니다. 어디에 넣는지는 구별해야 하며, 모든 구성에서 두 위치에 동시에 넣는 것은 아닙니다. 기준 정책을 사용하지 않는 구성도 가능합니다. RLHF Book — Regularization

old와 reference가 처음에는 같은 가중치일 수 있습니다. 하지만 학습이 진행되면 새 데이터를 만드는 old는 달라질 수 있고, 이 예의 reference는 처음 정한 상태를 유지합니다. 또 old와 비교하기 위해 옛 모델을 매번 실행할 필요는 없습니다. 선택한 토큰의 생성 당시 로그확률을 데이터에 기록해 사용할 수 있습니다.

이제 보상에서 시작한 정보가 어디로 이어지는지 볼 수 있습니다. 보상과 비교 기준은 어드밴티지를 만들고, 현재 모델의 확률 계산은 손실에서 가중치로 이어지는 미분 경로를 제공합니다. 데이터 재사용을 위한 확률비와 clipping, 기준 정책을 유지하는 KL은 이 경로의 서로 다른 부분을 조절합니다.

다음 글에서는 지금까지 하나의 상자로 둔 비교 기준을 자세히 보겠습니다. 같은 질문의 응답들을 비교하는 방법부터 시작해, 각 문맥에서 미래 보상을 예측하는 critic으로 넘어가겠습니다.