확산 모델로 보상을 생성하는 DRM, 인간 선호를 단일 점수에서 분포로

배경

RLHF에서 보상 모델은 사람의 평가를 언어 모델을 최적화하기 위한 신호로 바꾼다. 기존 방식은 대체로 하나의 프롬프트와 응답에 단일 스칼라 보상을 부여하거나, 보상이 특정한 모수적 분포를 따른다고 가정한다. 이 방식은 효율적이지만 인간의 선호가 가진 불확실성을 충분히 표현하지 못할 수 있다. 같은 응답도 정확성, 안전성, 유용성 등 서로 다른 기준에 따라 여러 방식으로 평가될 수 있기 때문이다.

칭화대 NLP 그룹이 제안한 Diffusion Reward Model, 즉 DRM은 이 문제를 분포 모델링의 관점에서 다룬다. 연구진은 보상 모델링을 입력과 응답에 조건화된 p(r|x,y)의 조건부 밀도 추정 문제로 다시 정의했다. 목표는 하나의 점수를 출력하는 것이 아니라 가능한 보상들의 구조를 학습하는 것이다.

DRM의 작동 방식

• 확산 과정으로 보상 생성: 동결된 대규모 언어 모델 인코더가 프롬프트와 응답을 표현하고, 경량 Diffusion Transformer가 가우시안 노이즈에서 보상 벡터를 단계적으로 복원한다.
• 고정된 분포 가정 회피: 가우시안과 같은 특정 분포를 미리 지정하지 않기 때문에 여러 개의 모드를 가진 선호 구조를 표현하기에 적합하다.
• 다양한 데이터 형식 지원: 하나의 아키텍처로 다중 속성 회귀와 쌍대 선호 데이터를 모두 처리한다.
• 분포 기반 추론: 같은 입력에서 여러 보상 샘플을 생성하면 경험적 보상 분포를 얻을 수 있다. 이를 평균에 가까운 대표값, 분산, 분위수 등으로 집계해 목적에 맞게 사용할 수 있다.

실험 결과와 의미

논문은 5개 벤치마크에서 데이터와 백본을 맞춘 기준 모델들과 DRM을 비교했다. 보고된 결과에 따르면 DRM은 일부 기준 모델과 같거나 더 나은 성능을 보였고, 비교적 작은 학습 규모로도 훨씬 큰 판별형·분포형·생성형 보상 모델과 경쟁력을 유지했다. 또한 일반적인 보상 헤드는 다중 모드 선호를 하나의 점으로 축소하는 반면, DRM은 이러한 구조를 회복할 수 있는 것으로 나타났다.

보상 분포는 단순한 분석 도구에 그치지 않는다. 연구진은 불확실성을 고려한 거부 전략과 하한 신뢰 경계, 즉 LCB 집계를 실험했다. 이 방식은 예상 보상이 높은지만 보는 대신 추정의 불확실성이나 위험까지 판단에 반영할 수 있게 한다. 다운스트림 RLHF 실험에서는 DRM을 학습 중 보상으로 사용했을 때 정책 성능이 향상됐다고 보고돼, 분포 정보가 실제 최적화에도 기여할 가능성을 보여준다.

영향과 한계

DRM은 보상 모델을 결정론적인 채점기가 아니라 인간 판단을 조건부로 생성하는 모델로 보는 관점을 제시한다. 평가자 간 의견 차이가 크거나 여러 품질 기준이 동시에 작동하는 정렬 작업에서는 단일 점수보다 더 풍부한 신호를 제공할 수 있다. 특히 불확실성을 이용하면 지나치게 자신 있는 보상 판단을 완화하는 데 도움이 될 수 있다.

다만 제공된 자료에는 각 벤치마크의 구체적인 점수, 반복 샘플링 비용, 집계 방식별 세부 비교가 포함돼 있지 않다. 따라서 DRM이 전통적인 보상 헤드를 항상 대체한다고 해석해서는 안 된다. 실제 적용에서는 추가 추론 비용과 분포 정보의 효용을 비교하고, 학습된 분포가 목표 사용자 집단의 선호를 제대로 반영하는지 검증해야 한다. 그럼에도 DRM은 RLHF를 스칼라 보상에서 구조화되고 불확실성을 인식하는 보상으로 확장하는 구체적인 경로를 제시한다.

출처: Hugging Face Daily Papers