How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis

"How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis" 논문에 대한 상세한 기술적 요약입니다.

문제 제기

보상 모델링 (Reward Modeling, RM) 은 대규모 언어 모델의 정렬 (예: RLHF 를 통한) 에 있어 핵심 구성 요소이지만, 표준 지도 예측과는 근본적으로 다릅니다. KL 정규화 정책 최적화에서 학습된 보상 함수 $r(x)$는 배포된 정책을 정의하기 위해 지수화되며, $\pi_\beta(dx) \propto e^{r(x)/\beta}\pi_{ref}(dx)$가 됩니다. 따라서 시스템의 하류 가치는 참조 분포 $\pi_{ref}$하의 예측 정확도뿐만 아니라, 보상이 높고 정책이 집중되는 영역에서의 오차에 의존합니다. 이로 인해 정책 유도 분포 이동 (policy-induced distribution shift) 이 발생합니다.

기존 오프라인 정렬에 대한 이론적 분석들은 종종 고정된 알려진 선형 특징을 가정하여, 신경망이 이러한 분포 이동 하에서 관련 특징을 어떻게 발견하는지 (특징 학습) 그리고 비선형 보상 지형을 어떻게 근사하는지에 대한 과제를 우회해 왔습니다. 본 논문은 훈련 목표가 지수 보상 가중치를 포함할 때 2 층 신경망이 특징을 학습하고 비선형 보상 함수를 적합시키는 방식을 분석함으로써 이 격차를 해소합니다.

방법론

저자들은 문제를 가우스 단일 인덱스 모델 (Gaussian Single-Index Model, SIM) 설정 내에서 공식화합니다:

• 데이터 생성: 입력 $x \sim \mathcal{N}(0, I_d)$, 그리고 레이블 $y = r^*(x) + \zeta$이며, 여기서 실제 보상은 $r^*(x) = \sigma^*(\langle \theta^*, x \rangle)$입니다. 여기서 $\theta^*$는 숨겨진 목표 방향, $\sigma^*$는 상한이 있는 다항식 링크 함수, $\zeta$는 유계 잡음입니다.
• 모델 아키텍처: $N$개의 뉴런을 가진 2 층 신경망: $r_a(x) = \frac{1}{N} \sum_{j=1}^N a_j \sigma(\langle w_j, x \rangle + b_j)$.
• 2 단계 훈련 프로토콜:
  1. 특징 회복 (1 단계): 첫 번째 층의 방향 $w_j$는 온도 $\beta_1$을 가진 보상 가중치 목적 함수를 사용하여 온라인 구형 SGD 를 통해 훈련됩니다. 두 번째 층의 계수 $a$는 초기화 시 고정됩니다. 목표는 $w_j$를 숨겨진 방향 $\theta^*$와 정렬시키는 것입니다.
  2. 리드아웃 적합 (2 단계): 특징이 회복되면 방향 $w_j$와 편향 $b_j$는 고정됩니다. 두 번째 층의 계수 $a$는 새로운 샘플에 대한 가중치 릿지 회귀를 사용하여 적합됩니다. 가중치 방식은 배포 온도 $\beta_2$를 사용하며, 이는 $\beta_1$과 다를 수 있습니다. 두 가지 가중치 방식이 분석됩니다:
     • 레이블 가중치: 가중치는 $e^{y/\beta_2}$입니다.
     • 대리 가중치: 가중치는 $e^{r_{proxy}(x)/\beta_2}$이며, 여기서 $r_{proxy}$는 고정된 대리 보상 (예: 특징 회복 단계에서 나온 것) 입니다.

주요 기여 및 결과

1. 보상 유도 분포 이동 하의 특징 학습
본 논문은 지수 보상 가중치 하의 첫 번째 층 특징 회복을 특성화합니다.

• 신호 변환: 지수 가중치는 첫 번째 층에 이용 가능한 헤르미트 (Hermite) 신호를 변경하는 비선형 레이블 변환으로 작용합니다.
• 지수 감소: 분석은 차원 무관한 $O(1)$ 임계값 이상의 모든 특징 학습 온도 $\beta_1$에 대해, 효과적인 학습 병목이 **정보 지수 (Information Exponent, IE)**에서 **생성 지수 (Generative Exponent, GE)**로 감소함을 보여줍니다.
• 회복 보장: 상한이 있는 다항식 가정 하에서, 뉴런의 일정한 비율이 높은 확률로 숨겨진 방향 $\theta^*$를 회복합니다. 약한 회복 (weak-recovery) 복잡도는 생성 지수에 의해 지배되며 $\tilde{O}(\beta_1^{2(I^*-1)} d^{(p_{gen}-1)\vee 1})$로 스케일링되며, 여기서 $I^*$는 GE 에 도달하는 데 필요한 지수입니다.
• $\beta_1$의 중요성: 결정적으로, $\beta_1$에 필요한 하한은 차원 $d$와 무관합니다. 이는 차원 증가 다로그 (polylogarithmic) 온도를 요구했던 이전 잘린 지수 분석을 개선한 것으로, 그렇지 않으면 지수 감소의 이점이 무효화되었을 것입니다.

2. 가중치 예측에서 기울어진 정책 후회 (Tilted-Policy Regret) 로
저자들은 표준 예측 오차 경계를 넘어 **기울어진 정책 가치 격차 (tilted-policy value gap)**를 평가합니다:
$$\mathcal{R} = \mathbb{E}_{x \sim \pi_{\beta^*}}[r^*(x)] - \mathbb{E}_{x \sim \hat{\pi}_{\beta_2}}[r^*(x)]$$
여기서 $\pi_{\beta^*}$는 목표 정책이고 $\hat{\pi}_{\beta_2}$는 학습된 보상에 의해 유도된 정책입니다.

• 분해: 후회는 세 가지 항으로 분해됩니다:
  1. 온도 불일치 ($T_{temp}$): 목표 온도 $\beta^*$와 배포 온도 $\beta_2$ 사이의 차이에 의존합니다.
  2. 투영 단절 ($T_{cut}$): 적분 가능성을 보장하기 위해 분석을 투영된 부분 공간으로 제한함으로써 발생합니다.
  3. 학습 항 ($T_{learn}$): 적합된 보상의 가중치 예측 오차에 의존합니다.
• 저온 거동: 본 논문은 기울어진 정책 하에서 보상의 분산이 $\beta_2^2 / (2p_{max})$로 스케일링됨을 보여주는 명시적 경계를 유도합니다. 여기서 $p_{max}$는 보상 최대값의 국소 차수입니다.

3. 레이블 및 대리 가중치 후회 경계
본 논문은 두 가지 적합 방식에 대한 명시적 가치 격차 경계를 제공합니다:

• 레이블 가중치 경계: 이론적 벤치마크 역할을 합니다. 이 경계는 온도 불일치, 단절, 학습 비용을 분리하여 정책 가치율을 비선형 보상 최대값의 국소 모양과 연결합니다.
• 대리 가중치 경계: 플러그인 파이프라인에 더 실용적입니다. 이 경계는 대리 유도 측도 (measure) 와 참조 분포 사이의 밀도 비율을 나타내는 추가 인자 $C_{0,R}(\beta_2)$를 도입합니다.
  • 대리 불일치 메커니즘: 대리 보상이 불완전한 경우, $C_{0,R}(\beta_2)$는 커질 수 있으며 ($1/\beta_2$에 대해 다항식으로 스케일링), 이는 근사 오차와 유한 샘플 오차 모두를 증폭시킵니다. 이는 "대리 불일치"와 보상 과최적화 현상 뒤에 있는 통계적 메커니즘을 식별합니다.
• 허용 가능한 온도: 분석은 배포 온도 $\beta_2$의 허용 가능한 집합을 산출합니다. 더 작은 $\beta_2$를 선택하면 목표 정책과의 일치를 개선하지만 지수 가중치를 통해 학습 비용을 증폭시킵니다. 최적의 선택은 이러한 트레이드오프를 균형 있게 맞춥니다.

중요성 및 주장

본 논문은 통제된 이론적 설정 내에서 보상 모델링과 정책 가치 간의 상호작용에 대한 기초적 특성을 제공한다고 주장합니다.

• 특징 학습: 온도가 $\beta_1$이 너무 작지 않다면, 보상 유도 지수 가중치가 IE 에서 GE 로 효과적인 학습 지수를 감소시킴으로써 특징 학습을 촉진할 수 있음을 보여줍니다. 이는 차원 의존 온도를 요구했던 이전 가정들을 수정하여 차원 무관한 $O(1)$ 온도가 회복에 충분함을 확립합니다.
• 평가 지표: 보상 모델링의 관련 통계적 목표는 참조 분포 예측 오차가 아니라 기울어진 정책 가치라고 주장합니다. 유도된 경계는 가치 격차를 보상 지형의 국소 기하학과 배포 온도 선택과 명시적으로 연결합니다.
• 대리 위험: 대리 보상을 가중치에 사용하는 것 (대리 가중치) 이 불완전한 경우 오차를 증폭시킬 수 있는 특정 통계적 메커니즘을 식별하며, 이는 실제에서 관찰된 보상 과최적화 현상에 대한 이론적 설명을 제공합니다.

저자들은 가우스 SIM 이 이론적 추상화이지만, 특징 학습과 정책 유도 분포 이동의 핵심 메커니즘을 분리하여 더 풍부한 입력 분포와 다중 인덱스 구조에 대한 향후 연구의 기초를 제공한다고 주장합니다.