기술 요약: 거리가 주의를 분산시킬 때: BT-Loss의 표현 거리 편향(Representation Distance Bias)에 대하여

1. 문제 정의

보상 모델(Reward Models, RM)은 인간의 선호도를 대리하여 대규모 언어 모델(LLM)의 정렬을 가이드하는 강화 학습으로부터 인간 피드로백(RLHF)의 핵심적인 역할을 합니다. 이러한 모델을 학습시키기 위한 표준 목적 함수는 "선택된(chosen)" 응답($y_w$)과 "거절된(rejected)" 응답($y_l$)으로 구성된 쌍체 데이터로부터 학습하는 브래들리-테리(Bradley-Terry, BT) 손실입니다.

저자들은 표준 BT 손실의 업데이트 역학에서 발견된 결정적인 결함을 지적합니다: 바로 **표현 거리 편향(Representation Distance Bias)**입니다. BT 그래디언트는 예측 오차(순위에 대한 실수 크기)에 의해 유도되어야 하지만, 분석 결과 그래디언트 노름(gradient norm)이 두 가지 별개의 성분에 의해 스케일링된다는 것이 밝혀졌습니다:

1. 예측 오차(Prediction Error): 선택된 응답과 거절된 응답 사이의 예측된 보상 차이.
2. 표현 거리(Representation Distance): 최종 레이어의 출력 공간에서 두 응답 사이의 거리($\|h_w - h_l\|$).

이러한 결합은 업데이트의 크기가 단순히 모델이 얼마나 틀렸느냐가 아니라, 두 응답의 표현이 얼마나 떨어져 있는지에 의해 결정되는 가짜 학습 신호를 생성합니다.

• 작은 거리의 쌍(Small-Distance Pairs): 의미적 또는 구조적으로 유사한 쌍(예: 올바른 코드 스니펫과 단 하나의 논리적 오류가 있는 코드 스니펫)은 종종 작은 표현 거리를 가집니다. 표준 BT 손실 하에서, 이러한 쌍들은 모델이 순위를 잘못 매겼음에도 불구하고 매우 약한 업데이트를 받게 됩니다. 이는 추론 및 코딩과 같은 작업에서 미세한 구분이 매우 중요하기 때문에 문제가 됩니다.
• 큰 거리의 쌍(Large-Distance Pairs): 매우 다른 쌍(예: 거절 문구와 유해한 질의에 대한 답변)은 자연스럽게 큰 표현 거리를 가집니다. 이러한 쌍들은 불균형하게 강한 업데이트를 받으며, 잠재적으로 더 미세하고 정교한 쌍들로부터 오는 학습 신호를 압도할 수 있습니다.

결과적으로, 모델은 가장 어려움을 겪는 미세하고 정교한 구분에 대해서는 학습이 부족해지는 반면, 표면적인 차이에는 과도하게 집중하게 됩니다.

2. 방법론: NormBT

이러한 편향을 해결하기 위해, 저자들은 각 선호도 쌍의 그래디언트 기여도를 재조정하는 적응형 쌍체 정규화 방식인 NormBT를 제안합니다.

이론적 기초

저자들은 BT 그래디언트 노름 $\|\nabla_\theta L_{BT}\|$을 분해하여 이것이 다음의 비례 관계에 있음을 보여줍니다:
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
여기서 $d$는 보상 차이이며, $\|h_w - h_l\|$은 표현 거리입니다. 업데이트 크기가 주로 예측 오차 항인 $|\sigma(d) - 1|$에 의존하도록 만들기 위해, 저자들은 표현 거리에 반비례하는 가중치 $w_i$를 도입합니다.

구현 세부 사항

그래디언트 노름을 정확히 계산하여 가중치를 부여하는 것은 계산 비용이 많이 듭니다. 대신, NormBT는 최종 레이어 표현 거리($\|h_w - h_l\|$)를 사용하는데, 이는 선형 스코어 헤드(linear score head)에 대한 정확한 대리물이며 전체 모델에 대한 강력한 근사치임이 입증되었습니다.

제안된 손실 함수는 다음과 같습니다:
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

여기서 가중치 $\tilde{w}$는 다음과 같이 정의됩니다:
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

주요 설계 선택 사항은 다음과 같습니다:

• Stop-Gradient: 정규화 가중치 자체가 최적화되는 것을 방지하기 위해 거리 항에 stop-gradient 연산자를 적용합니다.
• EMA 안정화: (임베딩 크기가 변화함에 따라 분포가 이동할 수 있으므로) 스케일 드리프트를 방지하기 위해, 저자들은 가중치를 정규화하기 위해 배치 평균 표현 거리($\mu_t$)의 **지수 이동 평균(EMA)**을 사용합니다. 이를 통해 가중치가 분포 변화에 적응하면서도 안정적으로 유지되도록 합니다.
• 경량 통합: NormBT는 아키텍처 변경을 요구하지 않으며, BT의 확률적 기초를 보존하고, 무시할 만한 수준의 계산 오버헤드만을 발생시킵니다.

3. 주요 기여

1. 그래디언트 분석: 본 논문은 BT 그래디언트 노름이 예측 오차와 표현 거리에 의해 공동으로 결정되며, 이로 인해 편향된 학습 신호가 발생한다는 상세한 이론적 및 경험적 분석을 제공합니다.
2. 편향 식별: 저자들은 이러한 결합이 미세한 구분이 필요한 중요한 작업(작은 거리의 쌍)에는 약한 업데이트를 주고, 큰 거리의 쌍이 학습 신호를 지배하게 만든다는 점을 입증합니다.
3. NormBT 제안: 표현 거리에 따라 업데이트를 정규화하여 학습 신호가 예측 오차에 집중되도록 하는 드롭인(drop-in) 수정 방식인 NormBT를 도입합니다.
4. 경험적 검증: 다양한 LLM 백본(Gemma-2b-it, Llama-3.2-3B)과 데이터셋(Unified-Feedback, Skywork-Reward)에 걸친 광범로한 실험을 통해 일관된 성능 향상을 보여줍니다.

4. 실험 결과

저자들은 RewardBench 벤치마크와 다운스트림 Best-of-N (BoN) 선택 작업에서 NormBT를 표준 BT 베이스라인 및 여러 변형 모델(BT + Margin, BT + Label Smoothing)과 비교 평가했습니다.

• RewardBench 성능: NormBT는 모든 설정에서 BT 베이스라인보다 일관되게 우수한 성능을 보였습니다. 가장 눈에 띄는 이득은 추론(Reasoning) 카테고리에서 관찰되었으며, 여기서 정확도가 평균 5% 이상 향상되었습니다. 이는 표준 BT에 의해 억제되었던 미세한 구분(작은 표현 거리)이 추론 작업에 중요하다는 가설과 일치합니다.
• 다운스트림 RLHF (BoN): Best-of-N 선택 실험에서, NormBT로 학습된 모델은 BT 베이스라인에 비해 더 높은 "골드 스코어(gold scores, 선택된 응답의 품질)"를 달s성했습니다.
• 절제 연구(Ablation Studies):
  • NormBT는 상대적인 가중치 문제를 해결하지 못한 LayerNorm이나 Gradient Clipping과 같은 다른 학습 기법보다 우수한 성능을 보였습니다.
  • 이 방법은 보완적인 성격을 띠어, 다른 BT 변형(예: NormBT + Margin) 위에 적용했을 때도 성능을 향상시켰습니다.
  • EMA 안정화를 제거하거나 대체 유사도 메트릭(코사인 유사도, 평균 풀링 임베딩)을 사용하면 성능 저하가 발생하여, 특정 설계 선택의 타당성을 입증했습니다.

5. 의의 및 주장

본 논문은 NormBT가 표준 BT 손실의 구조적 한계에 대한 원칙적이고 경량화된 솔루션을 제공한다고 주장합니다. 업데이트 크기를 표현 거리로부터 분리함으로써, NormBT는 보상 모델이 미세한 구분에 대해 더 충실하게 학습할 수 있도록 보장하며, 이러한 미세한 구분은 고품질 정렬을 위해 가장 어렵지만 중요한 부분입니다.

저자들은 자신들의 연구가 데이터 품질이나 어노테이터의 편향을 직접 다루는 것이 아니라, 기존의 쌍체 데이터로부터 추출되는 학습 신호의 효율성과 충실도를 개선하는 것임을 강조합니다. 그들은 NormBT를 판별적 보상 모델링 패러다임에 대한 기초적인 개선책으로 위치시키며, 특히 정밀한 추론과 미묘한 오류 탐지가 필요한 영역에서 다운스트림 RLHF 애플리케이션에 더 신뢰할 수 있는 보상 신호를 제공할 수 있다고 봅니다.