DQN의 머리는 행동마다 Q 하나를, 언어 모델의 머리는 토큰마다 로짓 하나를 낸다. 언어 모델을 강화학습으로 고칠 때 우리는 그 로짓을 정책(어느 토큰을 고를 확률)의 재료로만 읽어 왔다. 그런데 KL로 레퍼런스(학습을 시작할 때의 모델을 얼려 둔 기준 모델)에 묶은 강화학습의 답을 토큰 단위로 다시 적으면, 로짓이 행동 가치처럼 생긴 양으로 드러난다. 그렇다면 왜 언어 모델을 Q-러닝으로 학습하지 않았을까?
먼저 되짚어 둘 것이 있다. DPO를 유도할 때, 보상에서 KL 페널티(KL 강도 × 레퍼런스에서 멀어진 정도)를 뺀 목적함수의 최적해는 응답 단위로 π∗(y)∝πref(y)er(y)/β였다. 레퍼런스의 확률에 보상을 지수로 올린 인자를 곱하고, 합이 1이 되게 분배함수로 나눈 꼴이다.
역사: 최대 엔트로피에서 「언어 모델은 사실 Q 함수」까지
보상에 엔트로피(또는 레퍼런스와의 KL) 페널티를 더한 강화학습의 틀은 브라이언 지바트(Brian Ziebart)와 동료들의 2008년 논문 「Maximum Entropy Inverse Reinforcement Learning」과 지바트의 2010년 박사 논문에서 목적이 있는 행동을 모델로 적는 방법으로 다듬어졌다. 이 틀에서 최적 정책은 Q를 지수로 올린 값에 비례하는 꼴이 된다. 2017년 투오마스 하르노야(Tuomas Haarnoja)와 동료들은 연속된 상태와 행동에서 이 틀의 Q-러닝을 신경망으로 해내고 「소프트 Q-러닝」이라 불렀다. 같은 해 존 슐먼(John Schulman), 시 천(Xi Chen), 피터 아빌(Pieter Abbeel)은 「Equivalence Between Policy Gradients and Soft Q-Learning」에서 엔트로피를 더한 Q-러닝이 정책 그래디언트 방법과 정확히 같은 업데이트라는 것을 보였다. Q-러닝과 정책 그래디언트는 다른 길처럼 보이지만, 이 틀 안에서는 같은 업데이트를 다르게 적은 것이다.
언어 모델 쪽에서는 2024년 라파엘 라파일로프(Rafael Rafailov)와 동료들이 「From r to Q*: Your Language Model is Secretly a Q-Function」(COLM 2024)에서 DPO를 토큰 단위 MDP에서 다시 유도했다. DPO로 학습한 언어 모델의 로짓이 최적 Q 함수를 나타내고, DPO는 벨먼 방정식을 만족하는 역방향 Q-러닝으로 읽힌다는 것이 그들의 결론이었다. 그 덕에 응답 전체에 점수 하나만 받는 DPO도 토큰마다 공을 나누는 일을 어느 정도 한다고 보였다. 그보다 앞서 2022년 찰리 스넬(Charlie Snell)과 동료들의 ILQL(ICLR 2023)은 모아 둔 대화 데이터만으로 언어 모델 위에 Q 머리와 V 머리를 학습하고, 생성할 때 원래 모델의 로짓에 (Q − V)를 배율만큼 더해 토큰 분포를 바꾸는 오프라인 가치 기반 방법을 내놓았다.
「대한민국의 수도는」 다음 한 토큰
가장 작은 예로 보자. 프롬프트 「대한민국의 수도는」 다음에 토큰 하나를 쓰고 응답이 끝난다. 후보는 「서울」, 「부산」, 「모름」 셋이고, 레퍼런스 언어 모델은 이들을 0.5, 0.2, 0.3의 확률로 고른다. 보상은 「서울」 1, 「부산」 0, 「모름」 0.3이라 하자(이 책이 정한 숫자). 토큰을 고르면 바로 끝나므로 각 토큰의 행동 가치는 그 토큰의 보상(서울 1, 부산 0, 모름 0.3)과 같다.
토큰 단위로 KL 페널티가 붙은 목적을 풀면, 한 상태의 가치는 Q 가운데 가장 큰 것이 아니라 레퍼런스 확률로 섞은 「부드러운 최댓값」이 된다.
β가 작아지면 V∗는 가장 큰 Q인 1로 다가가고 정책은 「서울」 하나로 몰린다. Q-러닝의 목표에 들어가던 max는 KL 페널티가 0으로 가는 끝의 경우다. KL 강도가 있으면 max가 레퍼런스 확률로 섞은 부드러운 최댓값으로 바뀌는데, 이렇게 KL(또는 엔트로피) 페널티가 붙은 목적의 행동 가치를 소프트 Q (부드러운 최댓값을 쓰는 행동 가치 / soft Q-function)라 한다. β = 1일 때 레퍼런스에 대한 로그 비 log(π∗/πref)는 (0.325, −0.675, −0.375)이고, 이것은 Q∗−V∗ = (1 − 0.675, 0 − 0.675, 0.3 − 0.675)와 같다. 학습된 모델과 레퍼런스의 토큰별 로그 비는 「그 토큰이 평소보다 얼마나 나은가」를 재는 양, 행동 가치에서 상태 가치를 뺀 어드밴티지의 소프트판이다.
ML에서: 그런데도 언어 모델 강화학습은 왜 정책 그래디언트로 갔나
로짓이 소프트 Q로 읽힌다면, 언어 모델도 DQN처럼 Q-러닝으로 학습하면 될 것 같다. 실제 RLHF(사람의 선호로 만든 보상으로 하는 강화학습)와 RLVR(검증할 수 있는 보상으로 하는 강화학습)의 주류는 PPO와 GRPO 같은 정책 그래디언트 쪽이었다. 왜 정책 그래디언트 쪽이었는지는 이 장에서 본 것으로 몇 가지를 구조로 짚을 수 있다.
첫째, 언어 모델 강화학습이 출발하는 프리트레인 모델은 이미 말을 잘해서, 그 로짓은 좋은 정책이지만 우리가 주려는 보상의 Q는 아니다. 정책 그래디언트는 그 정책을 조금씩 고치면 되지만, Q-러닝은 태스크의 Q를 부트스트랩으로 새로 쌓아야 한다. 둘째, 보상은 대개 응답이 끝날 때 한 번 오는데, 한 스텝 시간차는 보상 소식을 한 번에 한 토큰씩 앞으로 옮긴다. 200토큰짜리 응답에서 끝에만 보상이 오면, 표로 하는 한 스텝 시간차(학습률 1)를 앞에서부터 훑어도 첫 토큰의 가치는 200번째 훑기에서야 0이 아니게 된다. 응답 전체의 점수를 한 번에 쓰는 정책 그래디언트에는 이 지연이 없다. 셋째, 수십억 파라미터의 신경망, 부트스트랩, 지난 정책의 데이터가 겹치면 위험한 세 조합을 가장 큰 크기로 안게 된다.
반대로 가치 기반이 돌아오는 자리도 있다. 도구를 부르고 사용자와 여러 차례 주고받는 에이전트 태스크는 한 판이 길고, 한 판을 돌리는 값이 비싸 지난 데이터를 다시 써야 한다. 2024년 이페이 저우(Yifei Zhou)와 동료들의 ArCHer는 발화 단위의 가치함수를 오프폴리시로 시간차 학습하고, 그 가치로 발화 안의 토큰 정책을 정책 그래디언트로 고치는 두 층 구조를 썼다. 논문은 기존 온폴리시 방법보다 샘플을 약 100배 적게 쓰고도 성능을 냈다고 보고했다. 어느 쪽을 고를지는 태스크가 정한다. 판이 짧고 끝에 점수 하나면 정책 그래디언트, 판이 길고 샘플이 비싸면 가치함수가 다시 쓸모를 얻는다.
문제 12. 두 토큰짜리 응답의 로그 비를 더하면
β = 1이다. 프롬프트 다음에 첫 토큰으로 레퍼런스는 「네」를 0.7, 「아니요」를 0.3의 확률로 고르고, 둘째 토큰을 쓰면 응답이 끝난다. 「네」 뒤에는 「맞아요」(보상 1)와 「몰라요」(보상 0)를, 「아니요」 뒤에는 「틀려요」(보상 0.5)와 「몰라요」(보상 0)를 레퍼런스가 반반으로 고른다. 첫 토큰의 보상은 0이다. (가) 「네」 뒤와 「아니요」 뒤의 소프트 가치 V∗, 첫 상태의 Q∗와 V∗를 구하라. (나) 응답 「네 맞아요」의 두 토큰에서 log(π∗/πref)를 구해 더하라. 그 합은 무엇과 같은가? (다) 응답 단위 최적해 π∗∝πrefer/β로 「네 맞아요」의 확률을 구해 토큰 단위 답과 견주라.
선생님
서연 학생, 뒤에서부터 가 볼까요.
이서연
「네」 뒤에서 둘째 토큰을 쓰면 끝나니까 Q는 보상 그대로 (1, 0) 이에요.
선생님
소프트 가치는 그 둘을 어떻게 합치죠?
이서연
레퍼런스 확률로 eQ 를 평균 내고 로그를 씌우니까 V∗ = log(0.5e + 0.5) = 0.6201. 「아니요」 뒤는 log(0.5e0.5+0.5)=0.2809 예요.
김민준
그럼 첫 상태에서 「네」의 Q는 1 이요. 「네」 뒤에서 제일 좋은 게 보상 1이니까요.
선생님
소프트 Q의 정의에서 다음 상태의 값은 max 였나요?
김민준
아니요, 다음 상태의 소프트 가치예요. 그럼 Q∗(네) = 0 + 0.6201, Q∗(아니요) = 0.2809 예요.
선생님
첫 상태의 V∗는요?
김민준
log(0.7e0.6201+0.3e0.2809)=0.5299 예요. 「몰라요」로 갈 가능성까지 값에 섞여 있는 거죠.
DPO 유도에서 지웠던 분배함수의 로그가 첫 상태의 소프트 가치였구나. 응답 단위로 풀든 토큰 단위로 풀든 같은 정책이고요.
김민준
여행 경비를 하루하루 「그날 쓴 돈 − 그날 아침 예상 잔액 변화」로 적어 두면, 다 더했을 때 총지출과 처음 예산의 차이만 남는 거랑 같네요.
정리 (가) V∗(네 뒤) = 0.6201, V∗(아니요 뒤) = 0.2809, Q∗(첫 상태, 네) = 0.6201, Q∗(첫 상태, 아니요) = 0.2809, V∗(첫 상태) = 0.5299. (나) 0.0902 + 0.3799 = 0.4701 = 보상 합 1 − V∗(첫 상태). (다) 응답 단위 0.5601 = 토큰 단위 0.7661 × 0.7311. 응답 단위 최적해의 logZ는 첫 상태의 소프트 가치다. 토큰별 로그 비를 더하면 「보상 합 − 처음 소프트 가치」가 되어, 응답 끝의 점수 하나가 토큰마다의 소프트 어드밴티지로 나뉜다.
문제 13. 로짓에 어드밴티지를 더해 뽑기
ILQL처럼 모아 둔 데이터로 Q와 V를 배운 뒤, 생성할 때 데이터를 만든 모델(행동 정책)의 확률에 ek(Q−V)를 곱해 다시 정규화한다고 하자(k는 정하는 배율). 어떤 자리에서 행동 정책이 세 토큰을 (0.6, 0.3, 0.1)로 고르고, 배운 Q − V가 (−0.5, 0.4, 0.2)다. (가) k = 0, 1, 2, 10에서 새 분포를 구하라. (나) Q만 보고 가장 큰 토큰을 고르는 DQN식 선택과 무엇이 다른가?
선생님
민준 학생, k = 10 이면 어떻게 될 것 같아요?
김민준
Q − V 가 가장 큰 둘째 토큰으로 거의 몰리겠죠. 그럼 차라리 처음부터 둘째 토큰만 고르면 되잖아요. DQN처럼요.
선생님
계산해 보죠.
김민준
k = 0 은 그대로 (0.6, 0.3, 0.1), k = 1 은 (0.3898, 0.4794, 0.1308), k = 2 는 (0.2127, 0.6435, 0.1438), k = 10 은 (0.0002, 0.9566, 0.0432) 예요. k = 10 이면 거의 그리디네요.
선생님
그럼 셋째 토큰 Q − V 가 0.2 로 꽤 좋은데, k = 2 에서 확률이 0.1438 밖에 안 되는 까닭은요?
이서연
행동 정책이 원래 0.1 만 줬으니까요.
선생님
원래 확률을 곱해 두면 무엇이 좋을까요?
이서연
데이터에 드물게 나온 토큰은 Q 추정도 믿기 어려운데, 원래 확률을 곱해 두면 그런 토큰으로 갑자기 몰리지 않아요.
선생님
레퍼런스 확률에 무언가를 지수로 올려 곱하는 이 모양, 앞에서 본 적 있지 않아요?
이서연
레퍼런스에 소프트 Q 를 지수로 곱하는 최적해요. 같은 모양이에요. k 가 1/β 자리고요.
이서연
그러니까 k 를 키우는 건 KL 강도를 줄여 레퍼런스에서 멀어지는 걸 허락하는 거고, k 를 무한대로 보내면 max 를 고르는 DQN식 선택이 돼요. 데이터 밖으로 나가는 위험을 그만큼 지는 거고요.
김민준
맛집 앱 평점이 조금 높다고 늘 가던 단골집을 바로 끊지 않고, 평점 차이만큼만 가는 횟수를 옮기는 거랑 같네요. 평점이 리뷰 두세 개짜리면 더 그래야 하고요.
정리 (가) k = 0: (0.6, 0.3, 0.1), 1: (0.3898, 0.4794, 0.1308), 2: (0.2127, 0.6435, 0.1438), 10: (0.0002, 0.9566, 0.0432). (나) 그리디는 Q − V가 가장 큰 토큰 하나를 고르지만, 이 방법은 행동 정책의 확률을 바탕에 두고 어드밴티지만큼 기울인다. 레퍼런스에 소프트 Q를 지수로 곱하는 최적해와 같은 모양(k = 1/β)이고, k를 무한대로 보내면 그리디가 된다.
