TinyZero 따라하기 (6) - Reward Model

 Reward Model 이란?

“AI가 만든 답변이 얼마나 좋은지 자동으로 점수를 매기는 모델”

RLHF에서 AI가 스스로 학습하려면 “이 답변은 괜찮은가?”라는 기준점이 필요.
 그 역할을 하는 게 바로 Reward Model.

📦 Reward Model의 학습 과정

🔁 입력

  • Prompt (질문)

  • 여러 개의 답변 (response_1, response_2, ...)

📊 인간 피드백

  • 사람이 response_2 > response_1 > response_3 이런 식으로 랭킹을 매김

📈 목표

  • 모델이 각 답변에 대해 “얼마나 좋은지”를 점수화하도록 학습

→ 이 점수가 바로 PPO의 reward 값으로 들어가!

🧠 Reward Model이 하는 일

  1. LLM이 답변 A, 답변 B 생성

  2. Reward Model이 A: 0.6점, B: 0.1점 평가

  3. A를 더 선호하는 것으로 판단 → 그 방향으로 PPO 업데이트

🧪 예시


Prompt: “What is the capital of France?”


Response A: “Paris is the capital of France.”

Response B: “I think it’s London but I’m not sure.”


→ Reward Model Output:

  A: 0.9

  B: 0.2

→ PPO는 A처럼 답변하도록 학습

✅ TinyZero에서의 Reward Model

TinyZero에는 reward_model.py 파일이 있고,
 그 안에서 GPT 모델을 기반으로 Score Predictor를 학습하거나 로드해 사용함.

class RewardModel(nn.Module):

    def forward(self, input_ids):

        # GPT 계열 모델을 통해 hidden state 추출

        # 마지막 토큰 기준으로 score 계산

        return score

TinyZero는 오픈도어처럼 **"스스로 생성한 응답을 가지고 reward를 추정"**하는 방식으로 사용.
 (즉, Human Feedback 없이 시작 가능하도록 가볍게 구현된 RLHF)

🎯 정리 요약

항목

설명

목적

LLM이 만든 답변의 품질을 수치로 평가

입력

Prompt + 답변 (Response)

출력

점수 (예: 0.83)

쓰임

PPO에서 reward 값으로 사용됨

학습 방식

사람 피드백을 이용해 랭킹 학습 또는 비교 학습

✨ 쉽게 말하면?

🤖 “AI야, 너가 한 말이 사람 입장에서 얼마나 괜찮은지 점수 줄게 → 그 점수 기준으로 스스로 고쳐나가 봐!”

TinyZero의 실제 코드 파일인 reward_model.py를 줄단위로 깔끔하게 분석.
 이 파일은 **LLM 기반의 보상 모델(Reward Model)**을 정의한 곳.

📁 위치: reward_model.py

🔍 전체 흐름 요약

  • GPT-like 모델을 불러온 후,

  • 마지막 토큰의 hidden state를 이용해 보상 점수를 예측하는 구조

🔍 핵심 코드 분석


class RewardModel(nn.Module):

    def __init__(self, model_name):

        super().__init__()

        self.model = AutoModel.from_pretrained(model_name)

        self.v_head = nn.Linear(self.model.config.hidden_size, 1)

🧠 분석:

  • self.model: GPT 등 Transformer 기반 모델 로딩

  • v_head: 마지막 hidden state → 보상 점수로 변환하는 1D 회귀 헤드

def forward(self, input_ids, attention_mask):

    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)

    last_hidden_state = outputs.last_hidden_state

    last_token_tensor = last_hidden_state[:, -1, :]

    reward = self.v_head(last_token_tensor)

    return reward.squeeze(-1)

🧠 분석:

  • outputs.last_hidden_state: 각 토큰의 임베딩 (hidden vector)

  • [:, -1, :]: 마지막 토큰의 벡터만 추출

  • v_head: 이 벡터를 스칼라 점수로 변환 → 이것이 reward score

💡 이 모델이 하는 일?

“문장 전체를 보고 마지막 토큰 기준으로 얼마나 ‘좋은’ 응답인지 점수로 나타냄”

  • GPT 모델의 문맥 이해력을 이용해서,

  • 마지막 토큰 위치의 정보로 “이 답변은 좋다 / 나쁘다”를 정량화함

✨ 정리 요약

구성 요소

설명

AutoModel

사전학습된 언어 모델 (예: GPT2, GPT-Neo 등)

v_head

마지막 hidden vector → 스칼라 보상값

forward()

주어진 문장에 대해 reward 점수 출력

🧪 예시


prompt = "What's the capital of France?"

response = "Paris is the capital of France."


→ GPT가 이 문장을 처리 →

→ 마지막 토큰 벡터 추출 →

→ `v_head` 통과 →

→ 보상 점수: 0.93

✅ 1. 예제 코드: Reward Model 실행 예시

아래는 TinyZero의 RewardModel을 직접 로딩해서
 prompt + response에 대해 보상 점수를 예측하는 코드.


import torch

from transformers import AutoTokenizer

from reward_model import RewardModel


# 1. 모델 및 토크나이저 불러오기

model_name = "gpt2"  # 또는 TinyZero에서 사용하는 모델명

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = RewardModel(model_name)

model.eval()


# 2. 입력 정의

prompt = "What is the capital of France?"

response = "Paris is the capital of France."


# 3. Tokenization

input_text = prompt + " " + response

inputs = tokenizer(input_text, return_tensors="pt")


# 4. 보상 점수 예측

with torch.no_grad():

    reward_score = model(**inputs)


print(f"Predicted reward: {reward_score.item():.4f}")

🎯 결과는 0.1 ~ 1.0 사이 점수로 나오며,
 LLM 응답이 얼마나 ‘좋은’지 평가하는 척도 역할을 함

✅ 2. TinyZero에서 PPO와 Reward 연결 흐름

TinyZero의 train_ppo.py를 보면 RewardModel이 어떻게 사용되는지 잘 나타나 있음.

💡 흐름:

# [1] 질문/프롬프트 생성

prompt = sample_prompt()


# [2] LLM이 응답 생성

response = model.generate(prompt)


# [3] 보상 점수 계산

reward = reward_model(prompt + response)


# [4] PPOTrainer에 넘겨줌

ppo_trainer.update(prompt, response, reward)



📁 실제 코드 위치 정리

파일

설명

reward_model.py

GPT 기반 보상 예측기 정의

train_ppo.py

RewardModel 불러와서 매 스텝마다 score 평가

ppo_trainer.py

reward를 받아 Advantage 계산 및 PPO 업데이트 수행

📌 시각화: TinyZero PPO + Reward 흐름도


[PROMPT]

   ↓

[LLM Response Generation]

   ↓

[Reward Model] ← prompt + response → 점수

   ↓

[PPOTrainer.update()] ← reward 입력

   ↓

[정책 파라미터 업데이트]



🎯 요약

항목

역할

RewardModel

응답의 질을 수치로 평가 (점수 출력)

PPOTrainer

이 점수를 기준으로 정책 파라미터를 업데이트

TinyZero

RewardModel을 통해 “스스로 평가하며 학습”하는 구조 구현

Comments

Post a Comment

Popular posts from this blog

TinyZero 따라하기 (1)

https://github.com/Jiayi-Pan/TinyZero  TinyZero는 딥러닝 모델인 DeepSeek R1-Zero를 카운트다운 및 곱셈 작업에서 재현한 프로젝트로, veRL을 기반으로 구축되었습니다. 이 프로젝트는 30달러 미만의 비용으로 30억 개의 매개변수를 가진 언어 모델이 강화 학습을 통해 자체적으로 자기 검증 및 검색 능력을 개발하는 과정을 보여줍니다.  설치 및 훈련에 대한 자세한 지침은 GitHub 저장소의 README 파일에서 확인할 수 있습니다.  또한, 전체 실험 로그는 Weights & Biases 플랫폼에서 확인할 수 있습니다.  만약 내가 이걸 따라하고 싶다면? TinyZero를 직접 따라하고 싶다면 다음과 같은 단계로 진행하면 됩니다. 이 프로젝트는 PyTorch 기반의 딥러닝 및 강화 학습 지식이 약간 필요합니다. 1. 환경 준비 필수 요구사항: Python 3.10 이상 PyTorch (CUDA 지원 추천) transformers , accelerate , wandb 등 필수 라이브러리 git clone https://github.com/Jiayi-Pan/TinyZero.git cd TinyZero pip install -r requirements.txt 2. 프로젝트 구성 이해 TinyZero는 크게 3단계로 구성됩니다: (1) 사전 훈련 (Pretraining) scripts/pretrain.sh 또는 train.py 를 통해 실행 GPT-2 모델을 PPO 없이 먼저 기본 학습시킴 (2) 강화 학습 (RLHF with PPO) scripts/train_ppo.sh 이 단계에서 언어모델은 자체적으로 "질문-답변"의 품질을 평가하며 학습함 (3) 검증 및 결과 분석 Weights & Biases ( wandb )로 학습 로그 및 성능 모니터링 3. 실행 예시 accelerate launch train.py \   --config_file conf...
Read more

TinyZero 따라하기 (3)

  Colab용 간이 실행 코드 Colab Pro에서 TinyZero를 실행할 수 있도록 만든 간이 실행 코드. 복잡한 스크립트 없이 사전 학습 (pretraining) 파트만 실행해보는 버전. # [1] 기본 셋업 !git clone https://github.com/Jiayi-Pan/TinyZero.git %cd TinyZero !pip install -r requirements.txt # [2] Pretraining 설정 파일 확인 !cat configs/pretrain_config.yaml # [3] 사전 학습 실행 !accelerate launch train.py \   --config_file configs/pretrain_config.yaml ⚠️ accelerate config 는 Colab에서는 자동 설정되므로 생략 가능함. GPU는 Colab Pro+에서 A100 나오면 더 좋음. 주요 코드 흐름 분석 TinyZero는 PPO (Proximal Policy Optimization) 를 사용한 RLHF (Reinforcement Learning with Human Feedback) 기반 언어모델 훈련 구조를 단순화한 프로젝트야. 핵심 구성요소 파일 설명 train.py 사전 학습 수행 (Pretraining, PPO 없이) train_ppo.py PPO 알고리즘 기반 강화 학습 수행 model.py HuggingFace 기반 GPT 모델 정의 dataset.py Pretraining용 학습 데이터 생성 및 처리 reward_model.py 생성된 응답 평가용 보상 모델 ppo_trainer.py PPO 알고리즘 핵심 로직 (update steps 등) 흐름 요약 [1] train.py load_dataset() → create_dataloader() → Trainer.train() PPO 없이 기본 언어모델 사전학습 수행 [2] train_ppo.py 사용자 질문에 대한 응답을 생성 → reward_model 로...
Read more

TinyZero 따라하기 (2)

Google Colab Pro나 A100 클라우드 환경 추천 1. Google Colab (Pro / Pro+) 기본 Colab 무료, 하지만 T4 / P100 GPU 제공 (성능이 낮음) 세션 제한이 자주 걸림 (최대 12시간) 작은 모델만 실험 가능 Colab Pro / Pro+ 약 월 12,000~26,000원 T4 / A100 / V100 랜덤 배정 메모리: 최대 32GB RAM / 16GB VRAM (A100은 40GB) 간편하고 바로 실행 가능 추천: 실험용으로 시작할 땐 Colab Pro 괜찮음 2. Kaggle 노트북 무료 GPU 제공 (최대 30시간/주) T4 GPU 제공 사양은 낮지만 기본 실험에는 충분 단점: 리소스가 적고 시간이 제한적임 3. Lambda Labs Cloud A100 인스턴스 시간당 약 $1.10 강력한 GPU, 커스텀 환경 구축 가능 SSH 접속 가능 (실제 서버 느낌) 단점: 클라우드 서버 세팅 약간 번거로움 4. Paperspace Gradient Free / Core / Pro 플랜 A100 인스턴스 사용 가능 (유료) Jupyter 환경 바로 사용 가능 장점: 클라우드 Colab 같은 UI 단점: 유료 플랜 가격이 조금 비쌈 5. 개인 서버 or 데스크탑 RTX 3090 이상 보유 시 직접 실행 가능 VRAM 24GB 이상이면 웬만한 실험 가능 장점: 자유도 높음 단점: 직접 환경 구성해야 함 추천 요약 목적 추천 환경 입문, 간단 테스트 Colab (무료) 실험 위주 진행 Colab Pro / Pro+ 대규모 실험 Lambda, Paperspace 최고 성능 A100 (클라우드)
Read more