reward-lens: A Mechanistic Interpretability Library for Reward Models

다음은 논문 **"reward-lens: A Mechanistic Interpretability Library for Reward Models"**에 대한 상세한 기술 요약입니다.

1. 문제 제기

인간 피드백을 통한 강화학습 (RLHF) 은 대형 언어 모델 (LLM) 의 행동을 형성하기 위해 **보상 모델 (Reward Models, RMs)**에 의존합니다. 보상 모델이 잘못된 선호도를 인코딩할 경우, 결과적으로 도출된 정책은 이러한 오류를 물려받아 증폭시킵니다. 보상 모델의 행동 평가 (예: RewardBench) 는 성숙해졌지만, 보상 모델을 위한 **기계적 해석 가능성 (mechanistic interpretability)**은 사실상 존재하지 않습니다.

기존 해석 가능성 도구들 (예: TransformerLens, Logit Lens, Direct Logit Attribution) 은 생성형 LLM 을 위해 설계되었습니다. 이들은 토큰 분포를 분석하기 위해 중간 모델 상태를 어휘 역임베딩 행렬 ($W_U$) 에 투영하는 방식에 의존합니다. 그러나 보상 모델은 $W_U$를 **스칼라 회귀 헤드 (scalar regression head, $w_r$)**로 대체합니다. 이러한 구조적 차이로 인해 표준 원시 연산자들이 적용 불가능해지며, 안전이 중요한 보상 모델의 내부 메커니즘이 여전히 블랙박스인 "구조적 간극"이 발생합니다.

2. 방법론: $w_r$ 중심 관점

본 논문은 보상 모델을 위한 기계적 해석 가능성 원시 연산자를 적용한 오픈소스 라이브러리인 reward-lens를 소개합니다. 핵심 이론적 통찰은 보상 헤드의 가중치 벡터인 $w_r$가 모든 해석 가능성 질문을 위한 자연스러운 축이라는 점입니다.

이 라이브러리는 최종 보상 $r$이 최종 잔류 스트림 $h_{final}$의 선형 함수라는 관찰에 기반합니다:
$$r = w_r^\top h_{final} + b_r$$
잔류 스트림은 레이어 간에 가법적으로 분해되므로 ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), 보상은 다음과 같이 정확하게 분해될 수 있습니다:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

이를 바탕으로 해당 라이브러리는 $W_U$를 $w_r$로 대체하여 다음과 같은 원시 연산자들을 구현합니다:

• Reward Lens: 레이어 $\ell$에서의 중간 잔류 스트림을 $w_r$에 투영하여 레이어를 거치며 선호도 신호가 어떻게 "결정화 (crystallization)"되는지 추적합니다.
• Component Attribution: $w_r^\top \cdot \text{contribution}$을 통해 각 하위 레이어 (Attention/MLP) 가 최종 스칼라 보상에 기여하는 정도를 계산합니다.
• Contrastive Activation Patching: 선호되는 완료와 선호되지 않는 완료 간의 활성화 값을 교환하고 $w_r^\top h_{final}$의 변화를 측정합니다. 이는 선호 데이터의 고유한 대조적 구조로 인해 보상 모델에 본질적으로 적합합니다.
• SAE Feature Attribution: 희소 오토인코더 (SAE) 를 사용하여 잔류 스트림을 특징으로 분해하고, 특징 정렬도를 $w_r^\top d_i$ (여기서 $d_i$는 사전 벡터) 로 계산합니다.

아키텍처 무관성: 이 라이브러리는 10 가지 방법 어댑터 프로토콜을 갖춘 HuggingFace 네이티브 방식을 사용합니다. 이를 통해 가중치 수술 없이도 표준 모델 (Llama, Mistral, Gemma-2) 과 복잡한 다목적 모델 (MoE 게이트가 있는 ArmoRM 등) 을 래핑할 수 있습니다.

이론 기반 확장: 이 라이브러리는 최신 정렬 이론 결과를 코드로 구체화합니다:

• Distortion Index: 평가 프로브로 포착되지 않는 품질 차원을 예측하여 (보상 해킹의 위험) 어떤 차원이 과소 포착되었는지 파악합니다.
• Divergence-Aware Patching: 패칭이 분포 외 (out-of-distribution) 활성화를 유발할 때 이를 감지합니다.
• Misalignment Cascade Detector: 서로 다른 정렬 불일치 차원 간의 상관된 실패 모드를 식별합니다.
• Reward-Term Conflict Analyzer: 서로 다른 보상 목표가 정렬되어 있는지, 직교하는지, 아니면 충돌하는지 확인합니다.
• Concept-Vector Analysis: 선형 개념 방향을 추출하고 보상에 대한 그들의 인과적 용량 - 반응 (dose-response) 을 측정합니다.

3. 주요 기여

1. 라이브러리 (reward-lens): 보상 모델의 기계적 해석 가능성을 위한 최초의 설치 가능하고 아키텍처 무관한 툴킷입니다. 이는 관찰적 (attribution) 인과적 (patching) 관점을 단일 대수적 프레임워크 아래 통합합니다.
2. 통합적 프레임: 모든 유용한 보상 모델용 해석 가능성 원시 연산자가 $w_r$에 대한 투영이거나 $w_r$을 따른 분해임을 확립합니다.
3. 실증적 검증: 두 가지 프로덕션 모델에 대한 체계적 적용: Skywork-Reward-Llama-3.1-8B-v0.2(단일 스칼라 헤드) 및 ArmoRM-Llama3-8B-v0.1(다목적 헤드).
4. 부정적 실증 발견: 논문은 선형 기여도 (linear attribution) 가 보상 모델에서 인과적 중요성을 예측하지 못함을 명시적으로 보고합니다.

4. 주요 결과 및 발견

저자들은 모델당 약 695 개의 RewardBench 선호도 쌍에 대해 라이브러리를 평가했습니다.

• 결정화 깊이 (Crystallization Depth):
  • Skywork: 선호도가 균일하게 늦게 결정화됩니다 (분수 레이어 깊이 $\approx 0.90$). 이는 출력 근처에서 조립된 통합된 메커니즘을 시사합니다.
  • ArmoRM: 선호도가 더 일찍 결정화되며 변동성이 더 크며, 이는 다목적, MoE 게이트 아키텍처를 반영합니다.
• Attribution vs. Causality (신실성 간극):
  • 관찰적 Attribution: 최종 MLP 레이어 (특히 mlp_L31) 가 두 모델 모두에서 보상 신호를 지배합니다.
  • 인과적 Patching: 초기 레이어 구성 요소 (예: mlp_L0, attn_L0) 가 종종 최종 보상에 가장 중요합니다.
  • 상관관계: 기여도 크기와 패칭 효과 간의 스피어만 순위 상관관계는 음수이거나 거의 0 입니다 (Skywork: $\rho = -0.256$; ArmoRM: $\rho = -0.027$).
  • 해석: 이는 버그가 아니라 중복된 잔류 스트림의 특징입니다. 후기 레이어는 병렬 경로가 보상할 수 있으므로 제거 (ablation) 될 수 있지만, 초기 레이어는 "하중을 지는 (load-bearing)" 역할을 합니다.
• 보상 해킹 프로파일:
  • Skywork는 길이, 자신감, 포맷팅 편향을 패널티 처리합니다 (일반적으로 RLHF 모델이 이를 보상하는 것과 대조적).
  • ArmoRM은 자신감과 포맷팅을 보상하지만 아첨 (sycophancy) 은 패널티 처리합니다.
• 개념 정렬:
  • Skywork 는 "동의 (agreement)"와 같은 개념과 강한 선형 정렬을 보입니다 (기울기 $\approx 1.0$).
  • ArmoRM 은 선형 개념과 약한 정렬 ($\approx 1\%$ 크기) 을 보여주며, 더 복잡하고 비선형적인 보상 기하학을 시사합니다.

5. 의의

• 패러다임 전환: 이 논문은 연구자들이 보상 모델에서 관찰적 기여도를 인과적 중요성의 대리 지표로 취급하는 것을 중단해야 한다고 주장합니다. 두 가지 간의 간극은 구조적이며 무시해서는 안 되고 측정되어야 합니다.
• 안전 도구: "왜곡 (distortion, 과소 포착된 차원)"과 "정렬 불일치 연쇄 (misalignment cascades)"를 감지할 수 있는 도구를 제공함으로써 학습 중 모니터링을 가능하게 합니다. 이는 보상 모델 감사를 사후 활동에서 RLHF 학습 루프 중 실시간 안전 장치로 전환시킵니다.
• 편집 가능성: 이 프레임워크는 보상 신호를 주도하는 특정 특징과 개념을 식별하여, 하위 정책 학습 전에 해석 가능성 기반 보상 모델 편집 (예: 보상 헤드에서 "아첨" 개념 억제) 을 위한 길을 엽니다.
• 개방 과학: reward-lens를 공개함으로써 저자들은 보상 모델의 내부 메커니즘에 대한 접근을 민주화하여 최첨단 AI 시스템의 안전을 보장하는 중요한 단계를 밟았습니다.

결론적으로, reward-lens는 보상 모델의 스칼라 출력 특성에 맞춰 기계적 해석 가능성을 적용함으로써 AI 안전 분야의 중요한 공백을 메웠습니다. 이는 보상 모델의 내부 논리가 종종 직관에 반함 (후기 기여도 $\neq$ 초기 인과성) 을 드러냈으며, 정렬 실패가 정책 모델로 전파되기 전에 이를 진단하고 수정하는 데 필요한 도구를 제공합니다.