인과적 요인 분해를 통한 보상 모델 강화

초록

본 논문은 인간 피드백 기반 강화학습(RLHF)에서 보상 모델이 학습하는 스퓨리어스(허위) 특성을 인과관계 관점에서 분석하고, 컨텍스트 임베딩을 인과적 요인과 비인과적 요인으로 분리하는 Factored Representation Learning 프레임워크인 CausalRM을 제안한다. 인과적 요인만으로 보상을 예측하고, 비인과적 요인에는 보상 예측을 방해하는 적대적 헤드를 적용해 보상 해킹을 완화한다. 수학 및 대화 데이터셋 실험을 통해 기존 방법 대비 더 견고한 보상 모델과 향상된 RLHF 성능을 입증한다.

상세 분석

CausalRM은 기존 RLHF 보상 모델이 프리트레인된 언어 모델 백본의 임베딩을 그대로 사용하면서 선형 보상 헤드만 부착하는 구조적 한계에 주목한다. 이러한 구조는 프롬프트‑응답 쌍이 내포하는 모든 정보—정답의 논리적 근거뿐 아니라 길이, 어투, 사교적 표현 등—가 보상 예측에 직접 영향을 미치게 하여 스퓨리어스 경로(z_nc → r)를 형성한다. 논문은 이를 인과 그래프(Figure 1)로 명시하고, “보상은 비인과적 요인에 대해 조건부 독립이어야 한다(r ⟂⟂ z_nc)”는 인과적 불변성 원칙을 제시한다.

이를 구현하기 위해 CausalRM은 두 단계의 변분 인코더를 도입한다. 첫 번째 인코더 qα(z_c|h)는 입력 임베딩 h에서 보상 예측에 충분히 필요한 정보를 압축하고, 두 번째 인코더 qα(z_nc|h)는 나머지 정보를 비인과적 요인으로 전이한다. 두 잠재 변수는 정규 사전(N(0,I))을 갖는 대각 공분산 가우시안으로 모델링되며, KL 정규화 항을 통해 정보량을 제한한다.

보상 헤드 gψ는 오직 z_c에만 연결되어 선형 스칼라 r̂ = gψ(z_c)를 출력한다. 이는 구조적 편향을 통해 비인과적 경로를 차단한다. 그러나 단순히 구조를 제한하면 posterior collapse와 같은 퇴화 현상이 발생할 수 있다. 이를 방지하기 위해 CausalRM은