이질적 인간 피드백을 위한 저차원 컨텍스트 강화학습

초록

본 논문은 사용자별 컨텍스트 정보를 활용해 인간 피드백의 이질성을 모델링하고, 파라미터 공간을 저차원으로 압축하는 Low‑rank Contextual RLHF(LoCo‑RLHF) 프레임워크를 제안한다. 저차원 서브스페이스를 추정한 뒤, 불확실성을 반영한 신뢰구간을 구성하고, 보수적(pessimistic) 정책을 도출하는 PRS 알고리즘을 설계한다. 이론적으로 서브옵티멀리티 갭을 기존 방법보다 더 타이트하게 상한을 제시하며, 합성 실험과 실제 PersonalLLM 벤치마크에서 개인화·분포 이동·고차원 상황에서 우수한 성능을 입증한다.

상세 분석

LoCo‑RLHF는 기존 RLHF가 가정하는 “모든 사용자가 동일한 보상 함수를 공유한다”는 전제를 깨고, 보상 함수를 r(x, s, a)=xᵀΘϕ(s,a) 형태의 bilinear 모델로 확장한다. 여기서 x는 사용자 컨텍스트 벡터, ϕ(s,a)는 상태‑행동 임베딩이며, Θ는 dₓ×d_ϕ 차원의 파라미터 행렬이다. 고차원 행렬 Θ를 직접 추정하면 계산량이 O(dₓ·d_ϕ)로 급증하지만, 실제 데이터에서는 Θ가 저차원 구조를 가진다는 가정을 도입한다. 즉, Θ*≈UΣVᵀ 형태로 rank r (r≪min{dₓ,d_ϕ}) 로 근사함으로써 파라미터 공간을 (dₓ+d_ϕ)·r 로 축소한다.

알고리즘은 크게 세 단계로 구성된다. 첫 번째 단계에서는 rank‑constrained MLE를 Burer‑Monteiro 형태로 풀어 저차원 서브스페이스 {U,V}를 추정한다. 비선형 최적화이지만 교대 그라디언트(Alternating Gradient Descent)를 이용해 수렴성을 확보한다. 두 번째 단계에서는 추정된 서브스페이스에 대한 불확실성을 정량화한다. 여기서는 서브스페이스 추정 오차와 MLE의 통계적 변동을 결합한 새로운 신뢰구간을 설계하여, 저차원 투영 후에도 보상 추정이 과도하게 낙관적이지 않도록 보장한다. 마지막 단계는 보수적 정책(Pessimism in Reduced Subspace, PRS)이다. 신뢰구간 하한을 보상으로 사용해 최악의 경우에 대한 정책을 최적화함으로써, 오프라인 데이터의 분포 이동(distribution shift)과 희소한 피드백에 대한 과적합을 방지한다.

이론적 결과는 두 가지 핵심을 제공한다. 첫째, 서브옵티멀리티 갭이 Õ((dₓ+d_ϕ)·r·log(1/δ)/n) 으로, 기존 Õ(dₓ·d_ϕ·log(1/δ)/n) 보다 차원 의존도가 크게 감소한다. 둘째, 그룹 구조가 사전에 정의된 특수 경우에 기존 Zhong et al. (2024)의 결과와 일치함을 보여, 제안 방법이 일반적인 저차원 가정 하에서도 최적에 가깝게 수렴함을 증명한다.

실험에서는 synthetic data에서 rank와 차원을 다양하게 변형하며 PRS 정책이 greedy 정책 및 기존 보수적 정책보다 일관되게 작은 서브옵티멀리티 갭을 기록한다. 특히, rank가 작고 차원이 클수록 성능 차이가 두드러진다. Real‑world PersonalLLM 벤치마크에서는 사용자 연령·교육 수준 등 메타데이터를 컨텍스트로 활용해 개인화된 보상 모델을 학습하고, PRS 기반 정책이 기존 단일 보상 모델 대비 응답 품질과 사용자 만족도에서 유의미하게 개선됨을 확인한다. 노이즈가 섞인 피처를 인위적으로 추가한 실험에서도 PRS는 안정적인 성능을 유지하며, 다른 방법은 급격히 성능이 저하되는 모습을 보인다.

전체적으로 이 논문은 (1) 인간 피드백의 이질성을 컨텍스트 기반 저차원 모델로 정량화, (2) 비선형 저차원 서브스페이스 추정과 불확실성 정량화를 결합한 새로운 신뢰구간 설계, (3) 보수적 정책을 통해 오프라인 RLHF의 분포 이동 문제를 해결한다는 세 가지 혁신을 제공한다. 이러한 접근은 대규모 LLM 정렬 작업에서 개인화와 안전성을 동시에 확보하는 실용적인 길을 제시한다.