기술 요약: 테스트 시간 조종 가능한 보상 모델을 위한 베이지안 선호 학습

문제 정의

보상 모델 (RMs) 은 인간 피드백을 통한 강화 학습 (RLHF) 을 통해 언어 모델을 인간의 선호도와 정렬시키는 데 핵심적인 역할을 합니다. 전통적으로 분류기 기반 보상 모델은 고정된 데이터셋을 사용하여 Bradley-Terry (BT) 모델 파라미터의 정적 추정기로 훈련됩니다. 일단 훈련되면 이러한 모델은 적응력이 부족하여 재훈련 없이 미지의 선호도 분포나 동적 트레이드오프 (예: 유용성 대 안전성 간의 균형) 에 적응할 수 없습니다. 또한 표준 분류기 기반 보상 모델은 선호도 확률이 1 로 수렴하여 인간 선호도의 실제 복잡성을 포착하지 못하는 국소 최적점에 빠지는 "과최적화 (over-optimization)"에 취약합니다. 반면 "LLM-as-a-Judge"접근법은 유연성을 제공하지만, 종종 높은 계산 비용이 드는 독점 모델에 의존합니다. 본 논문은 테스트 시간에 적응 가능하며 과최적화를 피하는 효율적인 분류기 기반 보상 모델의 필요성을 다룹니다.

방법론: 변분 문맥 내 보상 모델링 (ICRM)

저자들은 잠재 선호도 확률에 대한 감마 변분 추론으로 보상 모델링을 다루는 베이지안 프레임워크인 **변분 문맥 내 보상 모델링 (ICRM)**을 제안합니다.

핵심 공식화

1. 잠재 변수 정의: 모델은 프롬프트 $x$와 일련의 문맥 내 시연 (in-context demonstrations) $C$가 주어졌을 때, 선택된 응답 $y_w$가 거부된 응답 $y_l$보다 선호될 확률을 나타내는 잠재 확률 변수 $z \in [0, 1]$를 도입합니다.
2. 베타 사전 분포와 사후 분포:
   • BT 모델은 $N=2$(한 쌍) 에 특화되어 있으며, 선호도 결과의 가능도는 $z$로 매개변수화된 베르누이 분포를 따릅니다.
   • $z$에는 베타 사전 분포 $p(z) = \text{Beta}(\alpha_0, \beta_0)$가 적용됩니다. 저자들은 무정보적 시작점을 나타내기 위해 일반적으로 균일 사전 분포 $(\alpha_0, \beta_0) = (1, 1)$를 사용합니다.
   • 근사 사후 분포 $q_\theta(z | x, y_w, y_l, C)$는 2 차원 프로젝션 헤드가 있는 신경망으로 매개변수화되어 유틸리티 점수 $u_\theta$와 신뢰도 (증거) 점수 $s_\theta$를 출력합니다.
   • 사후 분포는 다음과 같이 재매개변수화됩니다: $\text{Beta}(\alpha_q, \beta_q)$
     • $\mu = \sigma(u_w - u_l)$는 사후 평균 (선호도 확률) 을 나타냅니다.
     • $\tau = \text{Softplus}(s_w) + \text{Softplus}(s_l) + 1$는 농도 (신뢰도) 를 나타냅니다.
     • $\alpha_q = \mu\tau$이고 $\beta_q = (1-\mu)\tau$입니다.

학습 목적 함수

모델은 증거 하한 (ELBO) 을 최대화하여 음수 ELBO 손실을 최소화함으로써 훈련됩니다:
$$\mathcal{L}_{ICRM} = -\mathbb{E}_{q_\theta}[\log z] + \lambda(N) D_{KL}(q_\theta(z|\cdot) \parallel p(z))$$

• 재구성 항: 첫 번째 항은 사후 분포가 관찰된 선호도를 얼마나 잘 설명하는지 측정합니다. 베타 분포의 경우, 이는 디가마 함수 $\psi$를 포함하는 폐쇄형 식을 가집니다. 이를 최소화하면 선택된 응답에 대해 $\mu$를 1 로 밀어붙입니다.
• 정규화 항: 두 번째 항은 학습된 사후 분포와 고정된 사전 분포 간의 KL 발산입니다. 하이퍼파라미터 $\lambda(N)$은 문맥 내 시연 수 $N$에 대해 단조 감소하는 함수입니다.
  • $N$이 작을 때 $\lambda(N)$은 커져 사후 분포가 사전 분포에 가깝게 유지되도록 강제합니다 (높은 불확실성).
  • $N$이 증가함에 따라 $\lambda(N)$은 감소하여 재구성 항이 지배적이 되게 하고, 모델이 데이터 기반의 확신 있는 사후 분포를 형성하도록 합니다.

이론적 보장

저자들은 KL 정규화 목적 함수가 표준 BT 모델에서 흔히 발생하는 "과최적화"실패 모드를 방지함을 증명합니다.

• 보조정리 3.1: 선호도 확률 $\mu \to 1$로 갈 때, 유틸리티 점수에 대한 손실의 기울기는 $\epsilon = 1-\mu$인 $\epsilon$에 비례하는 $1/\epsilon$ 인자로 페널티를 받아 통제되지 않는 최대화를 방지합니다.
• 정리 3.2: 손실 함수의 전역 최소해는 도메인 내부 ($0 < \mu^* < 1$ 및 $0 < \tau^* < \infty$) 에 존재함이 보장되어 모델이 경계 해로 수렴하지 않도록 합니다.

주요 기여

1. 원칙적인 변분 선호 학습: ICRM 은 RMs 가 소수 샷 문맥 내 샘플로부터 테스트 시간 선호도를 추론할 수 있게 합니다. KL 정규화 목적 함수는 전역 내부 최적해를 보장하여 선호도 평균의 과도한 최대화를 억제합니다.
2. 단일 및 다중 목적 조종 가능성: 변분 설계는 특정 선호도 방향에 적응하는 방향성 조종 가능성과 선호도 변화 정도를 인코딩하는 보정된 조종 가능성을 모두 허용합니다. 이는 문맥 내 시연을 혼합하여 상충되는 목표 간 (예: 안전성 대 유용성) 의 트레이드오프를 탐색할 수 있도록 다중 목적 설정으로 확장됩니다.
3. 검증 가능한 보상 매개변수화: 이 프레임워크는 검증 가능한 보상을 통한 강화 학습 (RLVR) 에 적용됩니다. 모델은 유틸리티 및 신뢰도 점수를 해석하여 방향성이 있고 문맥의 신뢰도에 보정된 보상 신호를 생성하며, 수학 추론 작업에서 표준 검증자 및 BT 모델보다 우수한 성능을 보입니다.

실험 결과

저자들은 Qwen3-4B 및 Llama-3.2-3B 백본을 사용하여 ICRM 을 평가했으며, 표준 분류기 기반 보상 모델 (BTRM, ArmoRM, GRM) 과 생성형 심사자 (LLM-as-a-Judge, GRAM-r2) 를 포함한 베이스라인과 비교했습니다.

• 보정된 조종 가능성: 도덕적 딜레마 선호도 쌍에서 ICRM 은 문맥 내 시연의 구성에 의해 지정된 선호도 변화 정도를 따르는 능력을 보여주었습니다. 시연 수 ($N$) 가 8 에서 64 로 증가함에 따라 기대 보정 오차 (ECE) 는 0.099 에서 0.076 으로 감소하여 LLM-as-a-Judge 베이스라인을 능가했습니다.
• 방향성 조종 가능성: SafeRLHF 및 HHH-Alignment 와 같은 벤치마크에서 ICRM 은 해로운 응답이 선택된 것으로 레이블링된 "반전"선호도 레이블에 높은 정확도 (최대 89%) 로 적응한 반면, 정적 베이스라인은 적응에 실패했습니다.
• RM-Bench 성능: 문맥 내 시연 수를 늘리면 RM-Bench 정확도가 크게 향상되었습니다. $N=64$에서 ICRM 은 평균 정확도 **70.8%**를 기록하여 $N=1$일 때의 60.5% 에서 상승했으며, LLM-as-a-Judge 베이스라인 (70.3%) 을 능가했습니다.
• 다중 목적 트레이드오프: RM-Bench 안전성 도메인에서 ICRM 은 "응답해야 함"과 "거부해야 함" 행동 사이에 매끄러운 파레토 프론티어를 생성했습니다. 이러한 프론티어의 초부피 (HV) 는 $N$이 증가함에 따라 증가하여 ICRM 이 정적 RMs 가 lacking 한 달성 가능한 트레이드오프 영역을 확장할 수 있음을 보여주었습니다.
• RLVR 적용: INTELLECT-MATH 데이터셋을 사용한 수학 추론 작업에서 ICRM 보상으로 훈련된 정책은 MATH-500 에서 **53.8%**의 정확도를 달성하여 골드 검증자 (50.4%) 와 표준 BT 보상 모델 (51.4%) 로 훈련된 정책을 능가했습니다.

중요성 및 주장

본 논문은 ICRM 이 훈련 후 미지의 선호도 표준에 적응할 수 있는 보상 모델을 향한 이론적으로 근거 있고 실용적으로 효과적인 경로를 제공한다고 주장합니다. 문맥 내 시연에 조건부인 베타 사후 분포를 가진 잠재 확률로 선호도를 취급함으로써 ICRM 은 베이지안 정규화와 문맥 내 조종 가능성을 통합합니다.

저자들은 이 접근법이 다음과 같음을 강조합니다:

• 독점 생성형 심사자의 계산 오버헤드 없이 전통적인 분류기 기반 보상 모델의 정적 성격을 해결합니다.
• KL 정규화를 통해 보상 과최적화를 완화하는 원칙적인 메커니즘을 제공합니다.
• 고정된 작동 지점이 아닌 선호도 공간의 지속적인 탐색을 가능하게 하는 동적 테스트 시간 다중 목적 정렬 제어를 가능하게 합니다.

이 연구는 미래의 보상 모델링이 고정된 데이터셋을 넘어 동적이고 시연 기반의 정렬로 이동할 수 있음을 시사하지만, 저자들은 컨텍스트 길이 제약과 해로운 행동으로 모델을 조종하는 것을 방지하기 위한 안전장치의 필요성과 관련된 한계를 지적합니다.