보상 모델을 조건부 확률분포로 다시 정의한 확산 보상 모델(DRM)

무엇인가

이 논문이 겨냥하는 문제는 보상 모델(RM)의 출력 형태 자체다. RLHF에서 RM은 최적화 신호를 정의하므로 그 품질이 정렬의 성패를 좌우하는데, 지배적인 설계인 판별형 RM(Bradley–Terry 손실)과 생성형 RM(다음 토큰 예측)은 결국 출력을 결정론적 스칼라 r(x,y)로 붕괴시킨다. 저자들은 이 가정이 인간 선호와 어긋난다고 본다. Anthropic-HH의 주석자 간 일치도가 약 63%에 불과하고, HelpSteer3-Preference에서도 세심한 필터링 후에도 불일치가 남는다. 즉 같은 응답이 여러 방식으로 합리적으로 평가될 수 있고, 이는 통계적 의미의 다봉(multimodal) 분포다. 기존의 탈출 시도들도 부분적이다. 다목적 RM은 고정된 속성 스키마를 요구하고, 생성형·루브릭 기반 판사는 긴 추론으로 비용을 치르면서도 단일 판정을 낸다. 파라미터 분포 헤드들은 분포족을 미리 확정한다. DPL과 URM은 가우시안 평균·분산을 예측해 구조적으로 단봉이고, DPRM은 범주형 분포라 빈(bin) 입자성에 묶이며, QRM은 고정 분위 격자를 예측해 분위 교차와 불안정한 꼬리 문제를 안는다.

어떻게 동작하나

제안 방법 DRM은 보상 모델링을 조건부 밀도 추정 p(r|x,y)으로 재정식화한다. 구조는 두 부분이다. 먼저 동결된 LLM 인코더가 프롬프트와 응답을 하나의 시퀀스로 포맷해 마지막 토큰의 은닉 상태 h를 만든다. 이 인코딩은 오프라인으로 미리 계산해 두고, 그 위에 경량 Diffusion Transformer(DiT) 보상 헤드를 학습시킨다. 이 헤드는 노이즈가 섞인 보상 r_t와 타임스텝 t, 조건 h를 받아 순방향 과정의 노이즈를 예측하며, 텍스트 조건과 타임스텝은 각 DiT 블록에 adaptive layer normalization으로 주입된다. 출력 분포에 어떤 파라미터 가정도 두지 않는다는 점이 핵심이다.

무엇과 다른가

학습은 하나의 확산 프레임워크 안에서 두 감독 방식으로 나뉜다. 다속성 회귀에서는 주석이 달린 차원만 마스크 m으로 선택해 노이즈 예측의 평균제곱오차를 최소화한다. 가우시안 확산에서 이 목적함수는 조건부 보상 분포의 스코어 함수를 학습하는 것과 동치다. 페어 선호 데이터에서는 점별 회귀만으로는 어느 응답이 선호되는지 담기 어려우므로, 복원식으로 얻은 디노이즈 추정치에 Bradley–Terry 스타일 랭킹 손실을 걸고, 최종 손실을 L_pair = L_denoise + λ_BT·L_BT로 결합한다. 절대 보상 라벨이 없는 선호 전용 데이터에는 0을 중심으로 선호·비선호 응답 사이에 고정 마진 Δ를 둔 대칭 의사 보상 타깃을 만들어 디노이징 감독으로 쓴다. 학습 중에는 확률적으로 텍스트 조건을 학습 가능한 무조건 벡터로 치환하는 classifier-free guidance도 적용한다.

어떻게 쓰나

추론에서는 DDIM과 CFG로 표준 가우시안 노이즈에서 N개의 보상을 독립 샘플링해 경험적 보상 분포를 만든다. 이 분포는 표준 RLHF·Best-of-N·보상 벤치마크와 호환되는 스칼라(평균)로 집계할 수도 있고, 분산이나 분위수로 요약할 수도 있다. 여기서 DRM만의 테스트타임 스케일링 축이 하나 더 생긴다. 응답 축은 후보 응답 N개를 생성해 최고 점수를 고르는 기존 Best-of-N과 같지만, 보상 축은 고정된 (x,y)에 대해 확산 샘플 수 N을 늘려 평균 추정을 조여 점수 정밀도를 올리는 것으로, 결정론적 스칼라 RM에는 존재하지 않는 레버다.

전제와 한계

실험은 RewardBench v2, PPE, RMB, RM-Bench, JudgeBench 다섯 벤치마크에서 이뤄졌다. DRM-Multi-8B는 ArmoRM의 다속성 코퍼스(569K 샘플, 19개 속성)로, DRM-Pref-8B는 Tulu3 선호 혼합(273K chosen/rejected 쌍, 마진 Δ=1)으로 학습했고, 둘 다 FsfairX-LLaMA3-RM-v0.1의 인코더를 동결 백본으로 쓰고 스칼라 값 헤드만 확산 헤드로 교체했다. 데이터와 백본을 고정한 통제 비교에서 DRM-Multi-8B는 평균 66.2로 스칼라 다속성 헤드 ArmoRM(62.3)과 파라미터 분위 헤드 QRM(64.1)을 모두 앞섰고, RMB에서 78.0으로 격차가 가장 컸다. 더 큰 모델들과도 경쟁한다. Skywork-Reward-Llama-3.1-8B-v0.2(64.8), Llama-3.1-Nemotron-70B(67.8), DeepSeek-GRM-27B(65.6), GPT-4o(67.7)와 비슷한 범위이며, URM(66.1)과는 동등하면서도 분포 가정을 두지 않는다. 두 감독 방식의 성능도 66.2 대 65.8로 유사해 프레임워크가 전이됨을 보였다. JudgeBench에서는 58.6으로 GPT-4o와 Claude-3.5-Sonnet을 상회했고, PPE Correctness에서도 63.8을 기록해 수학·MMLU·MBPP 하위 과제에서 안정적이었다.

분포 자체를 검증한 결과도 제시된다. 반복 주석이 있는 HelpSteer2-Disagreements에서 DRM의 샘플 분포와 인간 평점 경험분포의 Wasserstein 거리는 helpfulness에서 0.804로, 경험적 전역 사전(1.030), 전역 가우시안(1.032), 점별 베이스라인(1.032)보다 낮았고 JS 발산 0.215, L1 거리 0.907이었다. correctness에서도 Wasserstein 0.846으로 최선이었다. 인간 불일치가 커질수록 DRM 출력의 다봉 비율도 37.6% → 55.5% → 63.2%(helpfulness), 37.6% → 54.7% → 62.0%(correctness)로 증가했다. 이 분포 정보는 의사결정에도 쓰인다. 불확실성 기반 거부에서 커버리지를 100%에서 70%로 낮추면 PPE Correctness가 평균 2.81%포인트, RMB가 4.56~7.31점 올랐다. 평균과 표준편차를 함께 쓰는 LCB(λ=0.4) 집계는 전 구간 커버리지에서 평균 집계를 앞섰다. PPE BoN N=32에서 57.637→57.797, RMB K=3에서 Helpfulness 68.231→68.611, Harmlessness 62.493→62.791이다. 보상 축 스케일링도 RewardBench v2에서 단일 샘플 56.5%가 32샘플에서 65.6%로 오르는 식으로 확인됐다.

실무적으로 중요한 것은 이 분포가 오프라인 평가에만 머물지 않는다는 점이다. 저자들은 allenai/Llama-3.1-Tulu-3-8B-SFT를 공통 액터로, UltraFeedback 프롬프트로 RLHF를 돌려 FsfairX, ArmoRM, DRM-Multi를 같은 조건에서 비교했다. Arena-Hard v2에서 FsfairX-RLHF의 1.3이 DRM-Multi-RLHF에서 2.0으로, MT-Bench에서 73.6이 74.8로 올랐다. 즉 보상 모델을 학습 시점의 보상으로 쓸 때 정책 성능이 실제로 개선된다. 개발자 입장에서는 보상 헤드를 스칼라 회귀에서 확산 헤드로 바꾸는 것만으로 분산·분위·불확실성이라는 추가 신호를 얻고, 샘플 수를 늘려 점수 정밀도를 사는 추론 비용 트레이드오프를 선택할 수 있다는 뜻이다. 다만 도입 시에는 인코더를 동결하고 헤드만 학습한다는 전제, 그리고 분산 기반 거부·LCB 같은 집계 규칙의 λ 값이 결과에 영향을 준다는 점을 함께 확인해야 한다.

저자들이 밝힌 한계는 분명하다. 이 연구는 새 패러다임의 초기 탐색으로, DRM은 수십만 개 수준의 오픈소스 샘플로만 학습돼 수천만 개 큐레이션 선호로 학습되는 산업용 RM 규모에 크게 못 미치며, 절대 성능에서 최강 오픈소스 스칼라 RM을 아직 따라잡지 못한다. 또한 단일 8B 인코더와 고정된 학습 데이터 규모만 실험했으므로 백본 크기·모델 계열·더 큰 데이터 체제에서의 거동은 미탐색이며, 이 스케일링 경향을 규명하는 것이 향후 과제다. 다른 보상 모델과 마찬가지로 학습 데이터의 편향을 물려받을 수 있고, 문체적으로 설득력 있지만 사실이 틀리거나 사회적으로 유해한 출력에 높은 점수를 줄 수 있다. 인간 감독이나 안전 제약 없이 하위 언어 모델의 최적화 목적으로 쓰면 이런 편향을 증폭시키거나 리워드 해킹을 유도할 위험이 있다고 저자들은 경고한다.