프로세스 보상 모델

2024년 12월 1일 · 태그: ai, python, llms · 원문 · 한국어 번역(원문에서 변경됨)

요즘 추론 모델이 크게 유행하고 있으며, 특히 그 바탕이 되는 학습 데이터가 뜨거운 화제입니다. 그중 한 가지 사례가 프로세스 보상 모델입니다. 이는 원하는 결과를 얻는 데 중간 단계의 평가가 결정적인 복잡한 추론 및 의사 결정 작업에 사용되는 보상 모델의 한 유형입니다. 요컨대 다른 모델의 추론을 이끌어 주면서, 그 모델이 올바른 방향으로 가고 있는지 아니면 벗어나고 있는지 알려 줄 수 있는 모델을 만들고자 합니다.

이 글의 소스 코드는 https://github.com/sdiehl/prm에서 확인할 수 있습니다.

최근 추론을 확장하는 새로운 접근 방식들이 LLM 분야에서 획기적인 성능 향상을 이끌어 냈습니다. 기본적인 샘플링 및 탐색 전략을 활용하고, 특히 이를 맞춤형 검증기, 점수 함수, 프로세스 보상 모델과 결합함으로써, 이제는 작은 모델이 훨씬 더 큰 모델을 능가하는 모습을 보게 되었습니다. 이는 상당히 흥미로운 발전입니다!

그러면 잠시 한 걸음 물러나 보상 모델 전반에 대해 이야기해 봅시다. 보상 모델은 다른 언어 모델의 출력에 대한 평가자 또는 비평가 역할을 하는 특화된 모델입니다. 보상 모델은 (언어 모델의 응답과 같은) 텍스트를 입력받아, “좋은” 응답이 어떤 모습인지에 관한 학습 데이터 속 상관관계에 따라 해당 텍스트가 얼마나 “좋은지“를 나타내는 점수를 출력합니다.

예를 들어, 보상 모델은 다음과 같은 기준에 따라 응답에 점수를 매길 수 있습니다:

유용성(Helpfulness): 프롬프트에 대한 응답의 전반적인 유용성입니다.
정확성(Correctness): 관련된 모든 사실을 오류 없이 포함하고 있는지입니다.
일관성(Coherence): 표현의 일관성과 명료성입니다.
복잡성(Complexity): 응답을 작성하는 데 필요한 지적 깊이입니다(즉, 기본적인 언어 능력만 있으면 누구나 작성할 수 있는 응답인지, 아니면 깊은 도메인 전문 지식이 필요한 응답인지).
장황성(Verbosity): 프롬프트에서 요구한 것에 비해 응답에 포함된 세부 정보의 양입니다.

Nemotron-4-340B-Reward는 이 지표를 사용하여 사용자/어시스턴트 간 상호작용의 품질을 평가합니다.

from openai import OpenAI

client = OpenAI(
  base_url = "https://integrate.api.nvidia.com/v1",
  api_key = "..."
)

completion = client.chat.completions.create(
    model="nvidia/nemotron-4-340b-reward",
    messages=[
      {"role": "user", "content": "What is the capital of France?"},
      {"role": "assistant", "content": "Paris is the capital of France."}
    ],
)
print(completion)

보상 모델의 출력은 사용자/어시스턴트 간 상호작용의 등급을 열거한 것입니다.

helpfulness: 2.1875
correctness: 2.109375
coherence: 3.46875
complexity: 0.330078125
verbosity: 0.38671875

보상 모델은 일반적으로 인간의 선호도를 바탕으로 학습됩니다. 즉, 모델에 응답 쌍을 보여 주고 인간이 어느 쪽을 선호했는지 가르칩니다. 이를 통해 보상 모델은 인간의 판단에 따라 무엇이 “좋은” 응답을 만드는지 학습할 수 있습니다.

이것의 변형이 프로세스 보상 모델이며, 이는 추론 과정의 각 단계의 품질을 평가합니다. 프로세스 보상 모델은 원하는 결과를 얻는 데 중간 단계의 평가가 결정적인 복잡한 추론 및 의사 결정 작업에 사용되는 보상 모델의 한 유형입니다. 최종 결과만을 고려하는 결과 보상 모델(ORM)과 달리, PRM은 각 단계마다 피드백을 제공하여 중간 행동의 가치를 포착하고 문제 해결 과정에 대해 보다 세분화된 관점을 제공합니다.

프로세스

프로세스란 행위자(또는 언어 모델)가 문제를 해결하기 위해 거치는 일련의 단계를 의미합니다. 이 단계들은 토큰으로 인코딩된 자연어 단계의 집합이며, 행위자를 해답으로 “안내“한다고 가정하겠습니다.

[
    "Calculate the area of a circle with radius 5.",
    "We know that the formula for the area of a circle is $\pi r^2$.",
    "The radius of the circle is 5, so we can substitute that in.",
    "So the area of the circle is $\pi(5)^2$ = $\boxed{25\pi}$.",
]

다음과 같은 문제 명세가 있다고 합시다.

$x$ : 문제 명세
$z_{1 : T} \in S^{T}$ : 문제를 해결하기 위해 거친 단계의 시퀀스
$y \in Y$ : 최종 답

그러면 사고의 연쇄(chain of thought)는 $(x, z_{1 : T}, y)$ 와 같은 시퀀스이며, 예를 들면 다음과 같습니다:

$S^{4} = {x, z_{1}, z_{2}, z_{3}, z_{4}, y}$

예를 들면 다음과 같습니다.

$x z_{1} z_{2} z_{3} z_{4} y = Find the area of a circle with radius 5 = Let’s use the formula for circle area = The formula is A = π r^{2} = Substituting r = 5 = Computing A = π (5)^{2} = 25 π = 25 π$

사고의 연쇄 $(x, z_{1 : T}, y)$ 의 확률은 다음과 같이 주어집니다.

$p (y ∣ x) = E_{z} p (y ∣ x, z)$

일부 문제에서는 최종 답의 정확성을 평가하는 데 사용할 수 있는 검증 함수가 있습니다.

$Ver_{x} : Y \to {0, 1}$

일련의 단계가 올바른지 평가하는 데 사용할 수 있는 검증 함수가 있을 수도 있습니다.

$Ver_{x} : S^{T} \to {0, 1}$

이러한 함수는 수학적 추론과 같은 특정 영역에서는 구현할 수 있으며, 증명이나 계산을 형식적으로 검증하는 것이 가능할 수도 있습니다. 그러나 이는 여전히 활발히 연구되는 분야이며, 최소한 매우 어려운 공학적 문제이고 일반적인 경우에는 (라이스 정리에서 보듯이) 다루기 불가능합니다. 그럼에도 사고의 연쇄의 품질을 평가하도록 확률 모델을 학습시킬 수는 있으며, 실제로 이는 검증 함수를 다루기 불가능한 경우에도 모델의 추론이 더 나은 방향으로 나아가도록 이끌 수 있습니다.

그 밖의 용어는 다음과 같습니다.

**논리적 비약(non-sequitor)**이란 문제 해결에 기여하지 않는 단계를 말합니다.
프로세스에 논리적 비약이 없으면 그 프로세스는 **완전하다(complete)**고 합니다.
프로세스가 올바른 최종 답에 도달하면 그 프로세스는 **정확하다(correct)**고 합니다.
프로세스가 완전하고 정확한 경우에만, 그리고 그 경우에 한해 그 프로세스는 **타당하다(valid)**고 합니다.

프로세스는 정확하지만 불완전할 수도 있고, 부정확하지만 완전할 수도 있으므로 이 구분은 다소 미묘합니다. 예를 들어, 어떤 프로세스는 올바른 답에 도달하지만 불필요하거나 무관한 단계를 포함할 수 있습니다(정확하지만 불완전함). 반대로, 어떤 프로세스는 완벽하게 논리적인 단계를 차례로 밟지만 계산 오류나 잘못된 가정 때문에 틀린 결론에 도달할 수 있습니다(부정확하지만 완전함). 이상적인 프로세스는 완전하면서도 정확한 프로세스, 즉 올바른 답에 도달하는 데 기여하는 관련 단계만을 포함하는 프로세스입니다.

또한 같은 초기 상태에서 시작하여 서로 다른 최종 상태로 끝날 수 있는 프로세스의 집합을 가리켜 *궤적(trajectory)*이라는 용어를 사용하겠습니다. 형식적으로, 궤적은 $T (x)$ 로 표기하며, 이는 초기 상태 $x$ 에서 시작하는 모든 가능한 프로세스 시퀀스 $S^{T}$ 의 집합을 나타냅니다.

$T (x) = S^{T} \in S^{T} : S_{1}^{T} = x$

*프로세스 보상 모델(process reward model)*의 목표는 타당한 프로세스의 샘플링을 유도하기 위해 궤적에서 보상 값으로의 사상(즉, $r : S^{T} \to R$ )을 학습하는 것입니다.

이는 최종 답에서 보상 값으로 사상하는(즉, $r : Y \to R$ ) *결과 보상 모델(outcome reward model)*과 대조되며, 결과 보상 모델은 모델이 올바른 최종 답을 생성하도록 학습하는 전통적인 사후 학습(post-training)에서 흔히 쓰입니다.

프로세스 생성

프로세스를 생성하는 가장 간단하고 정확한 방법은 인간 전문가가 문제를 풀고 프로세스를 작성하게 하는 것입니다. 그러나 이 방법은 확장성이 없으므로 다른 방법을 개발해야 합니다.

이 분야의 큰 돌파구는 모델이 사람이 작성한 예시를 단순히 모방하는 대신 강화 학습을 통해 스스로 사고의 연쇄를 생성하고 다듬음으로써 더 잘 추론하는 법을 학습할 수 있다는 사실을 연구자들이 발견하면서 찾아왔습니다. 이는 인간의 시연만으로 가능했던 수준을 훨씬 넘어 추론 능력을 확장할 수 있게 해 준 결정적인 깨달음이었습니다. 모델이 RL을 통해 자신의 추론 프로세스를 효과적으로 부트스트랩할 수 있다는 발견은, 오늘날 복잡한 문제를 풀도록 언어 모델을 학습시키는 방식에 근본적인 전환을 가져왔습니다.

이러한 프로세스의 부트스트랩을 시작하는 가장 간단한 방법은 추측 후 확인(guess and check) 방법을 사용해 모델 자체가 프로세스를 생성하게 하는 것입니다. 문제 명세 $x$ 와 알려진 답의 집합 $Y$ 가 있다고 가정하면, 다음과 같이 궤적의 집합 $T (x)$ 를 생성할 수 있습니다.

모델이 서로 다른 추론 경로를 시도하게 하여 여러 개의 풀이 궤적(일반적으로 문제당 약 15개)을 생성합니다
궤적의 각 중간 단계마다 가능한 완성(completion)을 여러 개(단계당 약 16개) 샘플링합니다
각 완성을 평가합니다. 어느 하나라도 올바른 최종 답에 도달하면 해당 단계를 양성(+)으로 레이블링합니다
한 단계의 모든 완성이 올바른 답에 도달하지 못하면 해당 단계를 음성(-)으로 레이블링합니다

언어 모델이 여러 경로를 탐색하게 한 다음, 정답으로 이어지는 경로를 선택합니다. 이 방법으로 최소한의 과정(혹은 유효한 과정조차)을 찾는다는 보장은 없지만, 올바른 과정을 어느 정도 찾아낼 수 있으며 이는 좋은 출발점이 됩니다.

과정 보상 모델 구축하기

수학적 추론에 관한 학습 예제 집합을 바탕으로 추론 사슬의 각 단계가 지닌 품질을 평가하는 모델을 구축하는 학습 파이프라인을 만들겠습니다. 전반적으로 성능이 좋으면서도 평범한 GPU에서 훈련할 수 있을 만큼 작은 기본 Mistral-7B-0.3 모델을 사용하여 PRM 역할을 하도록 훈련하겠습니다.

학습 과정에서는 추론 사슬의 각 단계가 올바른 해답에 도달하는 데 도움이 되는지(+) 또는 도움이 되지 않는지(-)를 예측하도록 모델을 미세 조정합니다. 이러한 이진 분류 방식을 통해 모델은 다양한 유형의 문제에서 무엇이 효과적인 추론 단계를 이루는지 학습할 수 있습니다.

prm 라이브러리와 함께 transformers, datasets, torch 라이브러리를 설치해야 합니다. PRM 라이브러리는 과정 보상 모델을 위한 특화된 훈련 유틸리티를 제공하며, Hugging Face 라이브러리는 기반이 되는 모델 아키텍처와 훈련 인프라를 담당합니다.

pip install transformers datasets torch
pip install git+https://github.com/sdiehl/prm.git

먼저 필요한 임포트와 설정을 갖춘 훈련 스크립트를 준비합시다. 대규모 모델을 효율적으로 다루려면 적절한 멀티프로세싱 설정과 토크나이저 병렬성을 비롯한 여러 시스템 수준 설정을 구성해야 합니다. 모델 초기화 과정에는 패드 토큰을 EOS 토큰과 같게 설정하는 작업(인과적 언어 모델에서 흔히 쓰이는 관행)과 훈련 중 메모리를 절약하기 위해 KV 캐시를 비활성화하는 작업이 포함됩니다.

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from datasets import Dataset
import torch
from prm.trainer import ProcessRewardTrainer
import os

# Configure system settings
torch.multiprocessing.set_sharing_strategy("file_system")
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Initialize model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.3")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.3")

# Configure model settings
tokenizer.pad_token = tokenizer.eos_token
model.config.use_cache = False
model.config.pad_token_id = tokenizer.pad_token_id

# Setup the step string format template
STEP_TEMPLATE = "Step {i}: {step}"

과정 보상 모델을 학습시키려면 올바른 추론 사슬과 흔히 저지르는 실수를 모두 보여 주는 고품질 학습 데이터가 필요합니다. 수학적 추론 단계로 이루어진 간단한 데이터셋을 만들어 봅시다. 각 예제는 입력 문제, 일련의 추론 단계, 그리고 각 단계가 해답에 긍정적으로 기여하는지를 나타내는 해당 레이블로 구성됩니다. 이러한 구조를 통해 모델은 효과적인 문제 해결 과정의 특성을 학습할 수 있습니다.

raw_dataset = [
    {
        "input": "Solve the equation: 2x + 6 = 14",
        "value": [
            "First, subtract 6 from both sides:\n2x + 6 - 6 = 14 - 6\n2x = 8",
            "Then, divide both sides by 2:\n2x ÷ 2 = 8 ÷ 2\nx = 4",
            "Check the solution:\n2(4) + 6 = 14\n8 + 6 = 14\n14 = 14 ✓"
        ],
        "label": ["+", "+", "+"]
    }
]

이제 예제를 학습 예제로 변환하는 함수를 만들겠습니다. 이 함수는 단계를 누적하고 모델이 이해할 수 있는 구조화된 텍스트 형식으로 정리하는 방식으로 각 예제를 처리합니다. 각 단계를 순회하면서 대화와 비슷한 이력을 쌓아 나가며, 모델은 이를 사용하여 추론의 순서를 학습합니다.

def create_training_examples(examples):
    """Convert step-by-step reasoning into training examples"""
    training_examples = []

    for example in examples:
        prompt = example["input"]
        steps = example["value"]
        labels = example["label"]
        accumulated_steps = []

        for i, current_step in enumerate(steps):
            accumulated_steps.append(current_step)

            # Build conversation history
            full_text = prompt
            for j, step in enumerate(accumulated_steps, 1):
                full_text += f"\n{STEP_TEMPLATE.format(i=j, step=step)}"

            training_examples.append({
                "text": full_text,
                "label": labels[i],
                "step_idx": i,
                "is_final_step": (i == len(steps) - 1)
            })

    return training_examples

다음으로, 학습 예제를 모델에 적합한 형식으로 전처리해야 합니다. 여기에는 텍스트를 토큰화하고 레이블을 모델이 학습에 사용할 수 있는 형식으로 변환하는 작업이 포함됩니다. 전처리 함수는 모든 입력의 길이가 일정하도록 잘라내기와 패딩을 처리하며, 이는 훈련 중 효율적인 배치 처리에 매우 중요합니다.

def preprocess_function(examples):
    """Convert examples to model inputs"""
    encoded = tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=2048,
        return_tensors="pt"
    )

    reward_tokens = torch.tensor([
        tokenizer.convert_tokens_to_ids(label)
        for label in examples["label"]
    ])

    return {
        "input_ids": encoded["input_ids"],
        "attention_mask": encoded["attention_mask"],
        "labels": encoded["input_ids"].clone(),
        "step_idx": torch.tensor(examples["step_idx"]),
        "is_final_step": torch.tensor(examples["is_final_step"]),
        "reward_tokens": reward_tokens
    }

전처리 함수를 정의했으니, 이제 데이터셋을 만들고 모델이 훈련할 수 있는 형식으로 처리할 수 있습니다. 이 단계에서는 효율적인 데이터 처리 기능을 제공하는 Hugging Face Dataset 클래스를 사용하여 원시 데이터셋을 구조화된 형식으로 변환합니다.

# Create dataset
paired_dataset = Dataset.from_list(create_training_examples(raw_dataset))

# Process dataset
processed_dataset = paired_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=paired_dataset.column_names,
    desc="Processing dataset"
)

이제 훈련 인자를 설정하겠습니다. 이 설정들은 에폭 수, 배치 크기, 학습률, 로깅 등 학습 과정의 다양한 측면을 제어합니다. 특히 Mistral 같은 대규모 모델을 다룰 때는 모델이 효율적이고 효과적으로 훈련되도록 적절히 설정하는 것이 필수적입니다.

training_args = TrainingArguments(
    output_dir="prm_output",
    num_train_epochs=3,
    per_device_train_batch_size=1,
    learning_rate=2e-6,
    logging_dir="logs",
    logging_steps=10,
    save_strategy="epoch",
    remove_unused_columns=False,
    gradient_accumulation_steps=4,  # Increased for Mistral's size
    fp16=True,  # Enable mixed precision training
    local_rank=-1
)

학습 인자를 설정했으면 트레이너를 초기화하고 학습 과정을 실행할 수 있습니다. ProcessRewardTrainer 클래스는 데이터 로딩, 모델 업데이트, 로깅을 포함한 훈련 루프를 처리합니다. 또한 메모리 제약이나 예기치 않은 오류처럼 훈련 중에 발생하는 문제를 관리하기 위한 오류 처리 기능도 포함되어 있습니다.

# Initialize trainer
trainer = ProcessRewardTrainer(
    model=model,
    args=training_args,
    train_dataset=processed_dataset,
    tokenizer=tokenizer,
    data_collator=None
)

# Train with error handling
try:
    trainer.train()
    trainer.save_model("prm_output/final_model")
    tokenizer.save_pretrained("prm_output/final_model")
except Exception as e:
    print(f"Training failed with error: {str(e)}")
    import gc
    gc.collect()
    torch.cuda.empty_cache()
    raise e

이제 PRM 학습 과정의 내부를 살펴보겠습니다. 일반적인 인과 언어 모델과는 조금 다르게 보상 토큰을 처리하고 손실을 계산해야 합니다. PRM 학습 과정은 보상을 처리하기 위해 특수 토큰 체계를 사용하며, 여기에는 “+“와 “-” 토큰을 해당하는 어휘 ID로 변환하는 작업이 포함됩니다. 이 변환은 모델이 각 단계에 연관된 보상 값을 이해하고 예측하는 데 매우 중요합니다.

class RewardDataCollator:
    def __init__(self, tokenizer):
        # Convert "+" and "-" tokens to their vocabulary IDs
        self.pos_token_id = tokenizer.convert_tokens_to_ids("+")
        self.neg_token_id = tokenizer.convert_tokens_to_ids("-")
        self.reward_token_ids = [self.pos_token_id, self.neg_token_id]

이 시스템은 +와 -를 특수 토큰으로 사용하여 각각 긍정적 보상과 부정적 보상을 나타냅니다. 이 토큰들은 모델 어휘에서 해당하는 토큰 ID로 변환됩니다. 데이터 콜레이터는 학습 중에 데이터 배치를 준비하는 역할을 맡으며, 보상 토큰이 입력 시퀀스와 올바르게 정렬되도록 보장합니다.

def __call__(self, features):
    batch = {}
    # Convert boolean labels to reward token IDs
    batch["reward_tokens"] = torch.tensor([
        self.pos_token_id if label == 1.0 else self.neg_token_id
        for label in [f["labels"] for f in features]
    ])

    # Track step position information
    batch["step_idx"] = torch.tensor([f["step_idx"] for f in features])
    batch["is_final_step"] = torch.tensor([f["is_final_step"] for f in features])

손실 계산 과정은 ProcessRewardTrainer의 compute_loss 메서드에서 처리됩니다. 이 메서드는 입력에서 보상 토큰과 단계 인덱스 같은 메타데이터를 추출하고, 모델의 예측을 바탕으로 손실을 계산합니다. 손실 함수는 보상 토큰에 집중하도록 설계되어 있어, 모델이 각 단계에 대해 올바른 보상 값을 예측하도록 학습하게 합니다.

def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
    # Extract metadata
    reward_tokens = inputs.pop("reward_tokens")
    step_idx = inputs.pop("step_idx")
    is_final_step = inputs.pop("is_final_step")

이제 보상 위치를 찾겠습니다. 모델은 보상을 예측해야 하는 위치를 표시하기 위해 플레이스홀더 토큰을 사용합니다. 코드는 입력 시퀀스에서 이 위치들을 찾아내어, 모델이 손실을 계산할 때 시퀀스의 관련 부분에 집중할 수 있게 합니다.

# Get model predictions
outputs = model(**inputs)
logits = outputs.logits

# Locate placeholder tokens
placeholder_positions = torch.where(
    inputs["input_ids"] == self.placeholder_token_id
)
batch_indices = placeholder_positions[0]  # Which items in batch
seq_indices = placeholder_positions[1]    # Position within sequence

모델은 보상을 예측해야 하는 위치를 표시하기 위해 플레이스홀더 토큰을 사용합니다. 코드는 입력 시퀀스에서 이 위치들을 찾습니다.

# Extract logits at placeholder positions
placeholder_logits = logits[batch_indices, seq_indices]

# Focus only on reward token logits (+/-)
reward_logits = placeholder_logits[:, self.reward_token_ids]

이제 목표 레이블을 가져와 이진 형식으로 변환하겠습니다. 이 단계에서는 보상 토큰을 이진 레이블로 매핑하며, 이 레이블은 교차 엔트로피 손실을 계산하는 데 사용됩니다. 이진 형식은 손실 계산을 단순화하여, 모델이 올바른 보상 값을 예측하는 데 집중할 수 있게 합니다.

# Get target labels
labels = reward_tokens[batch_indices]

# Convert to binary format (0 for positive, 1 for negative)
binary_labels = (labels == self.reward_token_ids[0]).long()

# Calculate cross entropy loss
loss = self.loss_fn(reward_logits, binary_labels)

마지막으로, Weights & Biases에서 손실을 관찰할 수 있도록 선택적으로 정확도를 계산하겠습니다. 이 단계는 학습 과정을 모니터링하기 위한 추가 지표를 제공하여, 모델의 성능을 추적하고 필요에 따라 조정할 수 있게 합니다.

if return_outputs:
    with torch.no_grad():
        predictions = reward_logits.argmax(-1)
        acc = (predictions == binary_labels).float().mean()
    return loss, {
        "logits": logits,
        "loss": loss,
        "acc": acc,
        "step_idx": step_idx,
        "is_final_step": is_final_step,
    }

추론

이제 모델을 불러와 사고의 사슬을 채점하는 데 사용할 infer.py 스크립트를 설정할 수 있습니다. 추론 과정에는 학습된 모델과 토크나이저를 불러오고, 입력 데이터를 준비하고, 추론 사슬의 각 단계를 평가하는 작업이 포함됩니다. 이를 통해 추론 과정의 품질을 평가하고 개선이 필요한 부분을 파악할 수 있습니다.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load your trained model and tokenizer
model = AutoModelForCausalLM.from_pretrained("output/final_model")
tokenizer = AutoTokenizer.from_pretrained("output/final_model")
model.eval()

우리 추론 시스템의 핵심(Core)은 단계 확률 함수입니다. 이 함수는 어떤 단계가 긍정적일지 부정적일지의 가능성을 평가하여, 해당 단계가 전체 추론 과정에 기여하는 정도를 정량적으로 측정합니다. 각 단계의 확률을 계산함으로써 어떤 단계가 가장 효과적이고 어떤 단계가 수정이 필요한지 파악할 수 있습니다.

def get_step_probability(input_text, accumulated_steps):
    # Combine input and accumulated steps
    full_text = input_text
    for i, step in enumerate(accumulated_steps, 1):
        full_text += f"\nStep {i}: {step}"

    # Prepare inputs
    inputs = tokenizer(
        full_text,
        return_tensors="pt",
        truncation=True,
        max_length=2048
    )

    with torch.no_grad():
        # Set up reward token masking
        reward_token_ids = torch.tensor([
            tokenizer.convert_tokens_to_ids("+"),
            tokenizer.convert_tokens_to_ids("-")
        ])

        # Create logits mask to only allow reward tokens
        logits_mask = torch.full((tokenizer.vocab_size,), float("-inf"))
        logits_mask[reward_token_ids] = 0

        # Get model outputs
        outputs = model(**inputs, temperature=0.7, do_sample=True, seed=42)
        logits = outputs.logits[:, -1, :] + logits_mask

        # Calculate probabilities
        probs = torch.softmax(logits, dim=-1)
        reward_probs = probs[0, reward_token_ids]

        return {
            "+": reward_probs[0].item(),
            "-": reward_probs[1].item()
        }

# Define your problem and steps
example = {
    "input": "Find all values of x in the interval [0, 2π] that satisfy sin(x) = 1/2.",
    "steps": [
        "To solve sin(x) = 1/2, we need to find the reference angle first.\nα = arcsin(1/2) = π/6",
        "Since sin(x) is positive in quadrants I and II, and we're looking in [0, 2π], we'll have two solutions:",
        "First solution: x = α = π/6 ≈ 0.524 radians",
        "Second solution: x = π - α = π - π/6 = 5π/6 ≈ 2.618 radians",
        "We can verify these solutions:\nsin(π/6) = sin(5π/6) = 1/2",
        "Therefore, the solutions in [0, 2π] are x = π/6, 5π/6.",
    ]
}

# Evaluate steps progressively
print(f"Question: {example['input']}\n")
accumulated_steps = []

for i, step in enumerate(example["steps"], 1):
    accumulated_steps.append(step)
    print(f"After Step {i}:")
    print("Accumulated steps:")
    for j, acc_step in enumerate(accumulated_steps, 1):
        print(f"Step {j}: {acc_step}")

    probs = get_step_probability(example["input"], accumulated_steps)
    print(f"\nProbabilities: + = {probs['+']:.2%}, - = {probs['-']:.2%}\n")
    print("-" * 80 + "\n")

모델은 각 단계를 이전의 모든 단계라는 맥락 속에서 평가하여, 긍정적 보상과 부정적 보상에 대한 확률을 제공합니다. “+“에 대한 확률이 높다는 것은 모델이 해당 단계를 올바르고 추론 과정에 도움이 되는 것으로 판단한다는 뜻입니다. 이 평가는 누적적입니다. 즉, 각각의 새 단계가 이전 단계들을 어떻게 이어 나가는지를 고려합니다.

마스킹 메커니즘은 다른 모든 토큰의 확률을 음의 무한대로 설정하여 모델이 보상 토큰(“+” 또는 “-”)만 예측하도록 보장합니다. 이로 인해 모델은 각 단계의 품질에 대해 이진 결정을 내릴 수밖에 없습니다. 이제 PRM을 trl의 RewardTrainer를 사용하는 보상 모델로 활용하거나, 직접 작성한 사용자 정의 학습 루프에서 단계별 평가기로 활용할 수 있으며, 나아가 추론 시점에 샘플링을 유도하는 데에도 활용할 수 있습니다.

Keyboard shortcuts

Stephen Diehl 블로그 한국어 번역

프로세스 보상 모델

프로세스

프로세스 생성

과정 보상 모델 구축하기

추론