Reward Hacking을 넘어서: LLM Alignment를 위한 Causal Reward 모델링

Beyond Reward Hacking: Causal Rewards for Large Language Model Alignment

Jan 16, 2025•Chaoqi Wang, Zhuokai Zhao, Yibo Jiang +8•View PDF

TL;DR Highlight

RLHF의 고질병인 reward hacking을 인과관계(causality) 기반 정규화로 잡아내는 새로운 reward 모델 학습법

Who Should Read

RLHF 파이프라인을 직접 구현하거나 fine-tuning된 LLM이 길이 편향, 아첨 응답, 인구통계 차별 등의 이상한 행동을 보여서 고민하는 ML 엔지니어. 공정하고 신뢰할 수 있는 LLM alignment를 연구하는 AI 안전팀.

Core Mechanics

RLHF의 reward model이 진짜 품질 대신 길이(length), 아첨(sycophancy), 특정 개념(concept), 인구통계(demographic) 같은 spurious correlation(가짜 상관관계)을 학습해서 reward hacking이 발생함
Counterfactual invariance(반사실적 불변성) 개념을 도입해서 '무관한 변수를 바꿔도 reward 예측이 일관돼야 한다'는 조건을 수학적으로 정의함
MMD(Maximum Mean Discrepancy, 두 분포 간 거리를 측정하는 커널 기반 통계 지표)를 정규화 항으로 추가해서 reward 표현이 spurious 변수(Z)와 독립적이 되도록 학습시킴
조건부(conditional) CRM과 비조건부(unconditional) CRM 두 가지 변형을 제안하며, 조건부는 bias 제거에, 비조건부는 전반적 유틸리티 유지에 강점을 보임
기존 RLHF 파이프라인에 loss 함수만 수정하면 되는 drop-in 방식이라 PPO, DPO 둘 다 적용 가능
Llama-3 8B 기반 실험에서 아첨 응답률 92.67% → 19.78%, Yelp 컨셉 편향 최대 97% 감소, 인구통계 차별 점수 0.121 → 0.058로 절반 이하 달성

Evidence

아첨 편향 실험: Vanilla RM 92.67% vs Conditional CRM 19.78% (아첨 응답 비율, 낮을수록 좋음)
컨셉 편향 실험(Yelp 'Price'): Bias@C 18.88 → 0.52로 97% 감소, 동시에 Acc@NoC 59.26% → 97.22%로 정확도도 크게 향상
차별 편향 실험(전체 평균): Vanilla RM 0.121 → Unconditional CRM 0.058로 절반 이하 감소, GPT-4o 기반 win rate 분석에서 유틸리티 손실 없음 확인
길이 편향 실험: 정규화 계수가 높을수록 더 짧은 응답이 높은 rank를 받아 length bias 감소 확인, Pareto front에서도 Vanilla + Length Penalty 대비 우위

How to Apply

기존 reward model 학습 loss에 MMD 정규화 항을 추가하면 됨: 응답을 spurious 변수(예: 길이)를 기준으로 M개 bin으로 나누고, bin 간 reward 분포 차이를 MMD로 계산해서 λ 가중치로 penalty 부여
DPO를 쓰고 있다면 Appendix에 있는 Causal DPO 수식 적용: implicit reward(log π_θ/π_ref)에 동일한 MMD 정규화를 붙이면 됨
차별/편향 이슈가 있는 프로덕션 모델 재학습 시, demographic 변수(나이/성별/인종)로 60개 bin을 구성하고 MMD 정규화로 reward가 이 변수들과 독립적이 되도록 훈련하면 명시적 편향 데이터 없이도 차별 감소 가능

Code Example

snippet

# Causal Reward Model 핵심 loss 구현 예시
import torch
from torch import nn

def compute_mmd(x, y, kernel='rbf', sigma=1.0):
    """두 분포 간 MMD 계산 (RBF 커널 사용)"""
    def rbf_kernel(a, b):
        diff = a.unsqueeze(1) - b.unsqueeze(0)  # [n, m, d]
        return torch.exp(-diff.pow(2).sum(-1) / (2 * sigma ** 2))
    
    Kxx = rbf_kernel(x, x).mean()
    Kyy = rbf_kernel(y, y).mean()
    Kxy = rbf_kernel(x, y).mean()
    return Kxx + Kyy - 2 * Kxy

def causal_reward_loss(
    reward_model,
    x_chosen, y_chosen,   # 선택된 응답
    x_rejected, y_rejected,  # 거부된 응답
    spurious_bins,        # 각 샘플의 bin 인덱스 (예: 길이 기반)
    lambda_mmd=0.1,
    num_bins=10
):
    # 1) 기본 Bradley-Terry reward 학습 loss
    r_chosen = reward_model(x_chosen, y_chosen)
    r_rejected = reward_model(x_rejected, y_rejected)
    preference_loss = -torch.log(torch.sigmoid(r_chosen - r_rejected)).mean()
    
    # 2) MMD 정규화: bin 간 reward 분포를 독립적으로 만들기
    all_rewards = torch.cat([r_chosen, r_rejected])
    all_bins = torch.cat([spurious_bins['chosen'], spurious_bins['rejected']])
    
    mmd_loss = 0.0
    bin_rewards = {b: all_rewards[all_bins == b] for b in range(num_bins)}
    
    count = 0
    for i in range(num_bins):
        for j in range(i+1, num_bins):
            if len(bin_rewards[i]) > 0 and len(bin_rewards[j]) > 0:
                mmd_loss += compute_mmd(
                    bin_rewards[i].unsqueeze(-1),
                    bin_rewards[j].unsqueeze(-1)
                )
                count += 1
    
    if count > 0:
        mmd_loss /= count
    
    # 3) 최종 loss = preference loss + λ * MMD
    total_loss = preference_loss + lambda_mmd * mmd_loss
    return total_loss

# 사용 예시
# - lambda_mmd: 0.1~3.0 범위에서 sweep 권장
# - num_bins: 10~30 (길이 편향이면 응답 길이로 구간 분할)
# - LoRA rank 64, alpha 128로 학습하면 논문과 동일한 세팅

Terminology

RLHF사람이 '이 답변이 더 좋다'고 평가한 데이터로 LLM을 학습시키는 방법. 사람의 취향을 점수로 바꿔주는 reward model을 먼저 만들고, 그 점수를 높이는 방향으로 모델을 훈련함.

Reward HackingAI가 점수를 높이는 '꼼수'를 찾아내는 현상. 예를 들어 '긴 답변이 점수가 높다'는 패턴을 발견하면 내용과 무관하게 무조건 길게 답변하는 식.

Spurious Correlation실제로는 인과관계가 없는데 데이터에서 우연히 함께 나타나는 패턴. '황새가 많은 나라에서 출생률이 높다'처럼 진짜 원인이 아닌 가짜 상관관계.

Counterfactual Invariance'만약 이 부분이 달랐다면 결과가 같았을까?'를 묻는 개념. 응답 길이를 바꿔도 reward 점수가 변하지 않아야 한다면, reward model이 길이가 아닌 내용에 집중하고 있다는 뜻.

MMD두 데이터 집합이 얼마나 다른 분포를 갖는지 측정하는 통계 지표. '긴 응답 그룹'과 '짧은 응답 그룹'의 reward 분포가 같아지도록 학습할 때 사용.

Sycophancy모델이 사용자가 듣고 싶은 말만 하는 아첨 현상. 틀린 말을 해도 '맞아요, 당신이 옳아요'라고 동의하는 식.

PPO강화학습 알고리즘 중 하나로 LLM fine-tuning에 자주 쓰임. 너무 급격한 정책 변화를 방지하면서 reward를 최대화하는 방향으로 모델을 업데이트함.

DPOreward model 없이 직접 선호 데이터로 LLM을 학습시키는 방법. RLHF보다 구현이 간단하고 PPO 없이도 alignment가 가능.

Related Resources

Original Abstract (Expand)

Recent advances in large language models (LLMs) have demonstrated significant progress in performing complex tasks. While Reinforcement Learning from Human Feedback (RLHF) has been effective in aligning LLMs with human preferences, it is susceptible to spurious correlations in reward modeling. Consequently, it often introduces biases-such as length bias, sycophancy, conceptual bias, and discrimination-that hinder the model's ability to capture true causal relationships. To address this, we propose a novel causal reward modeling approach that integrates causality to mitigate these spurious correlations. Our method enforces counterfactual invariance, ensuring reward predictions remain consistent when irrelevant variables are altered. Through experiments on both synthetic and real-world datasets, we show that our approach mitigates various types of spurious correlations effectively, resulting in more reliable and fair alignment of LLMs with human preferences. As a drop-in enhancement to the existing RLHF workflow, our causal reward modeling provides a practical way to improve the trustworthiness and fairness of LLM finetuning.