기술 요약: 보상 불확실성을 이용한 강화 학습에서의 다양한 행동 유도

문제 정의

전형적인 강화 학습(RL)은 일반적으로 스칼라 보상의 기대 합계를 최대화하는 결정론적 정책을 추구한다. 이러한 가정은 최적의 결정론적 정책이 존재함을 보장하지만, 대규모 언어 모델(LLM) 미세 조정이나 과학적 발견과 같은 현대적 응용 분야에서는 상당한 한계를 나타낸다. 이러한 도메인에서는 사전 학습된 정책의 다양성을 보존하는 것이 매우 중요하다. 예를 들어, 창의적인 작업이나 사고 사슬(chain-of-thought) 추론을 위해서는 단일 행동으로의 수렴이 바람직하지 않을 수 있다. 또한, 보상 모델이 실제 기저의 선호도를 완벽하게 반영하지 못하는 불완전한 대리 지표일 때, 단일 결정론적 정책에 대한 과최적화는 생성 품질의 저하(보상 해킹, reward hacking)를 초래할 수 있다.

엔트로피 정규화(entropy regularization)나 다양성 보너스(diversity bonuses)와 같은 기존의 해결책들은 확률성을 강제하려고 시도하지만 다음과 같은 결정적인 단점을 갖는다:

1. 취약한 트레이드오프: 엔트로피 정규화는 증가된 확률성이 기대 보상의 감소를 대가로 이루어지는 트레이드오프를 강요한다.
2. 부적절한 정책 순위 결정: 다양성 보너스는 종종 높은 평균 보상을 가진 정책보다 다양성이 높은 하위 최적 정책을 선호하게 만드는 바람직하지 않은 순위 결정을 초래한다.
3. 표현력 부족: 바닐라 정책 경사(vanilla policy gradient) 및 기대 스칼라 반환(expected scalarized returns, ESR)을 포함한 많은 목적 함수들은 본질적으로 결정론적 최적 정책을 가지며, 필요한 다양성을 유도하지 못한다.
4. 최적화 문제: 스칼라화된 기대 반환(scalarized expected return, SER) 정식화는 비선형 목적 함수에 대해 편향되지 않은 샘플 기반 경사 추정량을 결-결여하는 경우가 많다.

저자들은 다양성을 인위적인 제약이 아니라, 보상 함수에 대한 불확실성에 대한 합리적인 대응으로 이해해야 한다고 주장한다. 보상 함수가 완벽하게 알려져 있지 않을 때(모호한 선호도나 불완전한 보상 모델로 인해), 단일 행동에 전념하는 것은 최적이 아닐 수 있다.

방법론: ROSA

본 논문은 보상을 보상 함수에 대한 분포로 대체하고 행동 집합에 대한 비선형 목적 함수를 적용함으로써 RL 목적 함수를 근본적으로 재구성할 것을 제안한다. 이 프레임워크를 **Randomized Objectives, Set Actions (ROSA)**라고 부른다.

ROSA+Max 기준

컨텍스추얼 밴딧(contextual bandit) 설정에 집중하여, 저자들은 주요 목적 함수인 ROSA+Max를 정의한다. $\rho$를 보상 함수 $R$에 대한 분포라고 하고, $\mathbf{Y} = (Y_i)_{i=1}^n$을 정책 $\pi_\theta$로부터 i.i.d.로 샘플링된 $n$개의 행동 멀티셋이라고 하자. 목적 함수는 다음과 같다:

$$J_{\text{ROSA+Max}}(\pi_\theta) = \mathbb{E}_{X \sim \mu, \mathbf{Y} \sim \pi_\theta} \left[ \mathbb{E}_{R \sim \rho} \left[ \max_{1 \le i \le n} R(X, Y_i) \right] \right]$$

직관적으로, 이 기준은 각 샘플링된 보상 함수에 대해 샘플링된 행동 집합 내에서 달성된 최대 성능을 집계한다. 이는 정책이 평균 보상을 최적화하기보다는, 모든 가능한 보상 함수 실현(realization)에 대해 집합 내의 적어도 하나의 행동이 잘 수행되도록 장려한다.

경사 추정(Gradient Estimation)

저자들은 이 목적 함수를 위한 원칙적이고 편향되지 않으며 분산을 줄인 정책 경사 추정량을 도출한다. $m$개의 샘플링된 보상 함수와 $n$개의 샘플링된 행동이 주어졌을 때, 경사 추정량은 다음과 같다:

$$\hat{g}(\theta) = \frac{1}{m} \sum_{k=1}^m \left[ \sum_{i=1}^n \left( \max_{1 \le j \le n} R_k(X, Y_j) - \max_{1 \le j \le n, j \ne i} R_k(X, Y_j) \right) \nabla_\theta \log \pi_\theta(Y_i|X) \right]$$

이 추정량은 특정 보상 함수 샘플에 대해 최대 보상을 달라는 행동에만 크레딧을 할당하며, 두 번째 항은 분산 감소를 위한 제어 변수(control variate) 역할을 한다.

일반화

이 프레임워크는 집합 함수(set function) $f$의 가족으로 일반화된다. 저자들은 이진 보상 함수에 대해, 보상의 합에 대해 엄격히 증가하고 엄격히 오목한(예: Softmax) 모든 집합 함수 $f$가 최대의 다양성을 가진 최적 정책을 생성함을 증명한다.

주요 기여

1. ROSA 프레임워크: 보상 함수에 대한 분포를 통해 보정된 다양성 제어를 제공하는 새로운 RL 목적 함수 계열의 도입.
2. 이론적 분석:
   • 최적 정책: 균등 분포를 가진 이진 보상 함수에 대해, ROSA+Max의 유일한 최적 정책은 모든 서로 다른 최적 행동에 동일한 확률 질량을 할당하는 정책임을 증명했다.
   • 제어 가능한 다양성: 비균등 보상 함수 분포가 어떻게 특정한 제어 가능한 확률적 정책을 유도하는지에 대한 공식을 도출하여, 보상 분포 가중치를 통해 특정 정답을 샘플링할 확률을 조정할 수 있음을 보여주었다.
   • 지형 분석(Landscape Analysis): ROSA+Max가 엔트로피 정규화의 "평탄한" 지형과 다양성 보너스의 "재순위 결정(re-ordering)" 문제를 피하며, 보상과 다양성을 동시에 높이는 경사 방향을 제공함을 입증했다.
3. 실험적 검증: 토이 도메인 및 대규모 설정(LLM 포스트 트레이닝 포함)에서의 광범위한 실험을 통해, ROSA가 상충하는 선호도와 알려지지 않은 실제 보상이 있는 시나리오에서 표준 정책 경사 및 다양성 보너스 베이스라인보다 우수한 성능을 보이며 다양한 행동을 유도함을 입 dimostrated 했다.

실험 결과

논문은 세 가지 주요 시나리오에서 ROSA를 평가한다:

1. 상충하는 선호도: 두 가지 모순되는 보상 함수(예: 패리티 vs 안티-패리티)가 있는 작업에서, 바닐라 정책 경사는 보상 상쇄로 인해 실패한다. ROSA는 두 선호도를 모두 균형 있게 맞추는 정책을 성공적으로 학습하여 다양한 유효 출력을 생성한다.
2. 다수의 정답: 여러 정답 길이가 모두 유효한 수학적 추론 작업에서, ROSA로 학습된 정책은 지정된 모든 길이 선호도로부터 샘플링하는 반면, 표준 정책 경사는 단일 길이 범위로 붕괴된다.
3. 알려지지 않은 실제 보상: 노이즈가 있는 보상 모델(다양한 정확도를 가진 앙상블 판사로 시뮬레이션됨) 환경에서, ROSA는 탁월한 강건성을 보여준다. ROSA는 보상 함수의 불확실성을 포착하여 보상 모드에 비례하여 확률 질량을 분배하는 반면, 엔트로피 정규화된 PG는 가장 가능성 높은 보상에 과적합되거나 하위 최적 행동으로 너무 넓게 탐색한다.

분포 외(out-of-distribution) 평가(AIME 2025)에서, 노이즈가 있는 MATH 데이터로 ROSA를 통해 학습된 모델은 표준 PG 및 비관적 앙상블 베이스라인에 비해 더 높은 pass@k 비율을 달립했다.

의의 및 주장

본 논문은 전통적인 정식화가 원하는 행동의 폭을 유도하는 데 실패하는 복잡한 RL 작업에 대해 ROSA가 강건하고 이론적으로 근거 있는 대안을 제공한다고 주장한다. 다양성을 보상 불확실성에 대한 합리적인 대응으로 재정의함으로써, 이 프레임워크는 보상 극대화와 정책 엔트로피 사이의 긴장을 해결한다.

저자들은 ROSA가 다음을 수행함을 강조한다:

• 바닐라 정책 경사와 최근의 액션 셋 접근 방식(best-of-$n$ 등)을 자연스럽게 일반화한다.
• 엔트로피 계수의 취약한 하이퍼파라미터 튜닝 대신 보상 함수 분포를 통해 다양성을 제어하는 메커니즘을 제공한다.
• 보상 모델이 불완전하더라도 단일 결정론적 행동으로의 붕괴를 방지하면서, 가능한 보상 함수 분포를 반영하는 확률적 정책을 보장한다.

궁극적으로, 이 연구는 여러 유효한 솔루션이 존재하는 도메인(특히 LLM 미세 조정 및 과학적 발견)에서, 에이전트가 처한 환경만큼 다양하고 적응력 있는 확률적 정책을 훈련하기 위한 원칙적인 토대를 제공한다.