RLHF에서 DPO까지: LLM 정렬(Alignment) 기술 논문 심층 분석

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer

class RewardModel(nn.Module):
    def __init__(self, model_name="gpt2"):
        super().__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=1
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.logits.squeeze(-1)

def reward_model_loss(reward_model, chosen_ids, chosen_mask, rejected_ids, rejected_mask):
    """Bradley-Terry 선호 모델 기반 보상 모델 손실 함수"""
    r_chosen = reward_model(chosen_ids, chosen_mask)
    r_rejected = reward_model(rejected_ids, rejected_mask)

    # 선호 응답의 보상이 비선호 응답보다 높도록 학습
    loss = -torch.log(torch.sigmoid(r_chosen - r_rejected)).mean()
    accuracy = (r_chosen > r_rejected).float().mean()

    return loss, accuracy

from datasets import Dataset

def build_preference_dataset(prompts, model, tokenizer, num_samples=4):
    """프롬프트별로 여러 응답을 생성하고 선호 쌍을 구성하는 파이프라인"""
    preference_data = []

    for prompt in prompts:
        inputs = tokenizer(prompt, return_tensors="pt")
        responses = []

        # 다양한 응답 생성 (temperature sampling)
        for _ in range(num_samples):
            output = model.generate(
                **inputs,
                max_new_tokens=256,
                temperature=0.8,
                do_sample=True,
            )
            response = tokenizer.decode(output[0], skip_special_tokens=True)
            responses.append(response)

        # 실제 파이프라인에서는 인간 레이블러가 순위를 매김
        # 여기서는 보상 모델 점수로 대체
        scored = [(r, reward_model_score(prompt, r)) for r in responses]
        scored.sort(key=lambda x: x[1], reverse=True)

        # 최상위-최하위 쌍을 선호 데이터로 구성
        preference_data.append({
            "prompt": prompt,
            "chosen": scored[0][0],
            "rejected": scored[-1][0],
        })

    return Dataset.from_list(preference_data)

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

def compute_dpo_loss(
    policy_model,
    reference_model,
    chosen_ids,
    chosen_mask,
    rejected_ids,
    rejected_mask,
    beta=0.1,
):
    """DPO 손실 함수: 보상 모델 없이 직접 정책 최적화"""
    # 정책 모델의 로그 확률 계산
    policy_chosen_logps = get_log_probs(policy_model, chosen_ids, chosen_mask)
    policy_rejected_logps = get_log_probs(policy_model, rejected_ids, rejected_mask)

    # 참조 모델의 로그 확률 계산 (gradient 불필요)
    with torch.no_grad():
        ref_chosen_logps = get_log_probs(reference_model, chosen_ids, chosen_mask)
        ref_rejected_logps = get_log_probs(reference_model, rejected_ids, rejected_mask)

    # DPO 로그 비율 계산
    chosen_logratios = policy_chosen_logps - ref_chosen_logps
    rejected_logratios = policy_rejected_logps - ref_rejected_logps

    # DPO 손실: -log(sigmoid(beta * (chosen_logratio - rejected_logratio)))
    logits = beta * (chosen_logratios - rejected_logratios)
    loss = -F.logsigmoid(logits).mean()

    # 메트릭: 선호 응답의 암묵적 보상이 더 높은 비율
    reward_accuracy = (logits > 0).float().mean()

    return loss, reward_accuracy

def get_log_probs(model, input_ids, attention_mask):
    """시퀀스의 토큰별 로그 확률 합산"""
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits[:, :-1, :]
    labels = input_ids[:, 1:]
    log_probs = F.log_softmax(logits, dim=-1)
    selected_log_probs = torch.gather(
        log_probs, dim=-1, index=labels.unsqueeze(-1)
    ).squeeze(-1)
    mask = attention_mask[:, 1:]
    return (selected_log_probs * mask).sum(dim=-1)

import torch
import torch.nn.functional as F

def ppo_step(
    policy_model,
    value_model,
    old_log_probs,
    states,
    actions,
    rewards,
    advantages,
    clip_epsilon=0.2,
    value_clip=0.2,
    kl_coeff=0.1,
    ref_log_probs=None,
):
    """PPO 업데이트 스텝 (LLM RLHF용)"""
    # 현재 정책의 로그 확률
    new_log_probs = policy_model.get_log_probs(states, actions)

    # 중요도 비율 계산
    ratio = torch.exp(new_log_probs - old_log_probs)

    # 클리핑된 대리 목적함수
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1.0 - clip_epsilon, 1.0 + clip_epsilon) * advantages
    policy_loss = -torch.min(surr1, surr2).mean()

    # 가치 함수 손실 (클리핑 적용)
    new_values = value_model(states)
    value_loss = F.mse_loss(new_values, rewards)

    # KL 발산 페널티 (참조 모델 대비)
    kl_penalty = 0.0
    if ref_log_probs is not None:
        kl_div = (old_log_probs - ref_log_probs).mean()
        kl_penalty = kl_coeff * kl_div

    total_loss = policy_loss + 0.5 * value_loss + kl_penalty

    return total_loss, {
        "policy_loss": policy_loss.item(),
        "value_loss": value_loss.item(),
        "kl_penalty": kl_penalty if isinstance(kl_penalty, float) else kl_penalty.item(),
        "mean_ratio": ratio.mean().item(),
    }

import torch
import torch.nn.functional as F

def compute_orpo_loss(model, chosen_ids, chosen_mask, rejected_ids, rejected_mask, alpha=1.0):
    """ORPO 손실: SFT + Odds Ratio 기반 선호 최적화 통합"""
    # SFT 손실 (선호 응답에 대한 NLL)
    chosen_logits = model(input_ids=chosen_ids, attention_mask=chosen_mask).logits
    sft_loss = F.cross_entropy(
        chosen_logits[:, :-1, :].reshape(-1, chosen_logits.size(-1)),
        chosen_ids[:, 1:].reshape(-1),
        reduction="mean",
    )

    # 로그 확률 계산
    chosen_logps = get_sequence_log_probs(model, chosen_ids, chosen_mask)
    rejected_logps = get_sequence_log_probs(model, rejected_ids, rejected_mask)

    # 승산비(Odds Ratio) 계산
    chosen_odds = chosen_logps - torch.log1p(-torch.exp(chosen_logps))
    rejected_odds = rejected_logps - torch.log1p(-torch.exp(rejected_logps))
    log_odds_ratio = chosen_odds - rejected_odds

    # ORPO 최종 손실
    orpo_loss = sft_loss - alpha * F.logsigmoid(log_odds_ratio).mean()

    return orpo_loss

import numpy as np
from typing import List, Dict

def compute_alignment_metrics(
    model_responses: List[str],
    reference_responses: List[str],
    reward_model,
    tokenizer,
) -> Dict[str, float]:
    """정렬 품질 평가를 위한 종합 메트릭 계산"""
    metrics = {}

    # 1. 보상 모델 점수
    rewards = []
    for response in model_responses:
        tokens = tokenizer(response, return_tensors="pt")
        score = reward_model(**tokens).item()
        rewards.append(score)
    metrics["mean_reward"] = np.mean(rewards)
    metrics["reward_std"] = np.std(rewards)

    # 2. 승률 (win rate) - 참조 모델 대비
    wins = 0
    for model_resp, ref_resp in zip(model_responses, reference_responses):
        model_tokens = tokenizer(model_resp, return_tensors="pt")
        ref_tokens = tokenizer(ref_resp, return_tensors="pt")
        model_score = reward_model(**model_tokens).item()
        ref_score = reward_model(**ref_tokens).item()
        if model_score > ref_score:
            wins += 1
    metrics["win_rate"] = wins / len(model_responses)

    # 3. 응답 길이 분포 (길이 해킹 탐지)
    lengths = [len(r.split()) for r in model_responses]
    metrics["mean_length"] = np.mean(lengths)
    metrics["length_std"] = np.std(lengths)

    # 4. 다양성 지표 (모드 붕괴 탐지)
    unique_bigrams = set()
    total_bigrams = 0
    for response in model_responses:
        words = response.split()
        for i in range(len(words) - 1):
            bigram = (words[i], words[i + 1])
            unique_bigrams.add(bigram)
            total_bigrams += 1
    metrics["distinct_2"] = len(unique_bigrams) / max(total_bigrams, 1)

    return metrics