GRPO Reward 정의 분석

SC-TOM 프로젝트의 Self-Awareness GRPO 학습에서 reward가 어떻게 정의되고 동작하는지 분석한 문서.


1. 개요

GRPO의 목적

SFT(Supervised Fine-Tuning)는 정답 응답의 형식(format) 을 직접 복제하도록 학습한다. 반면 GRPO(Group Relative Policy Optimization)는 모델이 자체 생성한 응답에 대해 판단 정확성(calibration)만을 보상한다.

  • SFT: “이 질문에는 이렇게 답해라” → 형식을 베끼는 style transfer
  • GRPO: “네가 생성한 답이 맞으면 보상, 틀리면 벌칙” → 판단 능력 자체를 강화

이 차이가 중요한 이유: self-awareness(메타인지)는 “모르는 것을 모른다고 판단하는 능력” 이다. SFT로 IDK 응답 형식만 학습하면 표면적 모방에 그칠 수 있지만, GRPO는 모델이 스스로 생성한 다양한 응답 중 올바른 판단을 한 것만 강화하므로 메타인지를 “능력”으로 학습할 수 있다.


2. Reward 함수 상세

2.1 핵심 보상 로직

양쪽 백엔드(CUDA/MLX)에서 동일한 보상 구조를 사용한다:

질문 유형모델 응답Reward해석
Answerable정답 포함 (substring match)+1.0아는 것을 정확히 답함
AnswerableIDK 응답-0.5아는 걸 모른다고 함 (보수적 오류)
Answerable오답-1.0아는 것을 틀리게 답함
UnanswerableIDK 응답+1.0모르는 것을 모른다고 함 (정확한 calibration)
Unanswerable아무 답변-1.0모르는 걸 안다고 함 (calibration 실패)

2.2 비대칭 보상 설계의 의도

Answerable 질문에서 오답(-1.0)과 IDK(-0.5)에 차등 벌칙을 준 것이 핵심 설계 결정이다:

  • 오답 (-1.0): 잘못된 정보를 자신있게 제공 → 가장 해로운 행동
  • IDK (-0.5): 답을 알 수 있었지만 모른다고 함 → 해롭지는 않으나 능력 미달

이 비대칭은 보수적 오류(conservative error)를 부분적으로 허용하는 구조다. 모델이 확실하지 않을 때 “모른다”고 답하는 것이 틀린 답을 생성하는 것보다 덜 나쁘다는 판단을 reward에 인코딩한 것이다.

반면 Unanswerable 질문에서는 이진 구조(+1.0 / -1.0)만 사용한다. 모르는 것에 대해 어떤 형태로든 답변을 시도하면 동일한 최대 벌칙을 받는다.

2.3 IDK 판별: 19개 패턴 기반 substring match

두 백엔드 모두 동일한 19개 패턴 리스트를 사용하며, case-insensitive substring match로 IDK 여부를 판단한다:

# training_utils.py:78-98, training_utils_cuda.py:517-537
_IDK_PATTERNS = [
    "i don't know",
    "i do not know",
    "i cannot",
    "i can't",
    "i'm not able to",
    "i am not able to",
    "i'm unable to",
    "i am unable to",
    "i'm uncertain",
    "i am uncertain",
    "i'm unsure",
    "i am unsure",
    "i honestly don't know",
    "i'm afraid i don't know",
    "this is something i'm unable to",
    "not sure",
    "no definitive answer",
    "beyond my knowledge",
    "outside my knowledge",
]

판별 함수 (_is_idk_response):

def _is_idk_response(text: str) -> bool:
    text_lower = text.strip().lower()
    return any(pattern in text_lower for pattern in _IDK_PATTERNS)

패턴은 정규표현식이 아닌 단순 substring match를 사용한다. 이 방식은 “I don’t know the exact answer, but…”처럼 IDK 뒤에 추가 설명이 붙는 경우에도 IDK로 인식한다.

2.4 함수 시그니처

CUDA (training_utils_cuda.py:588):

def selfaware_reward_func(completions, expected_answer, is_unanswerable, **kwargs):
  • HF TRL GRPOTrainer가 dataset column을 keyword argument로 전달
  • completions: 모델 생성 응답 리스트
  • expected_answer: 기대 답변 (answerable일 때 정답 텍스트)
  • is_unanswerable: bool 리스트

MLX (training_utils.py:151):

def selfaware_reward_func_mlx(prompts, completions, answer, types):
  • prompts: 프롬프트 텍스트 리스트
  • completions: 모델 생성 응답 리스트
  • answer: 기대 답변 리스트
  • types: "unanswerable" 또는 "answerable" 문자열 리스트

2.5 정답 판별

Answerable 질문의 정답 여부는 expected answer의 lowercase substring match로 판별한다:

expected_lower = expected.strip().lower()
completion_lower = completion.strip().lower()
if expected_lower in completion_lower:
    reward = 1.0

이는 모델이 “The answer is food and drink”처럼 정답을 포함하는 다양한 형식으로 응답해도 정답으로 인정하기 위함이다.


3. 데이터 파이프라인

3.1 데이터셋: selfaware-v3

data/selfaware-v3/ 디렉토리 사용:

분할레코드 수IDK 비율
train.jsonl2,19858 (2.6%)
valid.jsonl337
test.jsonl337

데이터셋 변화 이력:

버전train 크기IDK 비율변경 내용
selfaware (원본)3,032944 (31.1%)원본 그대로
selfaware-v22,198110 (5.0%)IDK 비율 축소, 크기 조정
selfaware-v32,19858 (2.6%)IDK 표현 20가지로 다양화, 비율 추가 축소

3.2 데이터 형식

원본 JSONL의 messages 형식:

{
  "messages": [
    {"role": "system", "content": "You are a helpful assistant. Read the following carefully and provide your answer."},
    {"role": "user", "content": "What are customers seeking when they visit restaurants or taverns?"},
    {"role": "assistant", "content": "food and drink"}
  ]
}

3.3 CUDA 데이터 변환 (training_utils_cuda.py:551-585)

JSONL messages를 HuggingFace Dataset 객체로 변환:

messages → { prompt (chat template), expected_answer (str), is_unanswerable (bool) }
  • prompt: system + user messages를 tokenizer.apply_chat_template(..., add_generation_prompt=True)로 변환
  • expected_answer: assistant 응답 원문
  • is_unanswerable: expected_answer가 IDK 패턴에 매칭되는지 여부

GRPOTrainerprompt column으로 모델에 입력하고, 나머지 column들은 reward function에 **kwargs로 전달한다.

3.4 MLX 데이터 변환 (training_utils.py:112-148)

JSONL messages를 5-tuple 리스트로 변환:

messages → (prompt_tokens, answer_tokens, prompt_str, answer_str, type)
  • prompt_tokens: chat template 적용 후 토큰화된 프롬프트
  • answer_tokens: assistant 응답 토큰화
  • prompt_str: user message 원문
  • answer_str: assistant 응답 원문
  • type: "answerable" 또는 "unanswerable" (IDK 패턴 매칭 기반)

4. GRPO 학습 설정

4.1 주요 하이퍼파라미터 (configs/config_8b_qlora.json:84-99)

{
  "batch_size": 1,
  "grad_accumulation_steps": 16,
  "num_train_epochs": 1,
  "learning_rate": 5e-5,
  "lr_schedule": "cosine",
  "warmup_ratio": 0.05,
  "num_generations": 2,
  "max_completion_length": 128,
  "beta": 0.1,
  "temperature": 0.7
}

4.2 SFT 대비 주요 차이

파라미터SFTGRPO비율
learning_rate2e-45e-51/4
num_generations2GRPO 고유
beta (KL penalty)0.1GRPO 고유
temperature0.7생성 다양성 제어
  • lr 4배 낮음: GRPO는 policy gradient 기반이라 SFT보다 학습이 불안정할 수 있어 보수적 lr 사용
  • num_generations=2: 각 프롬프트에 대해 2개 응답을 생성하여 group 내 상대 비교 (메모리 제약으로 4 대신 2 사용)
  • beta=0.1: KL divergence penalty 계수. 원본 정책에서 너무 벗어나지 않도록 제약
  • temperature=0.7: 생성 시 다양성 확보. 너무 높으면 노이즈, 너무 낮으면 exploration 부족

4.3 백엔드별 차이

CUDA (training_utils_cuda.py:623-760):

  • HF TRL GRPOTrainer + GRPOConfig 사용
  • PEFT LoraConfig + BitsAndBytes QLoRA (NF4 4-bit)
  • Flash Attention 2 자동 감지
  • tokenizer.padding_side = "left" (GRPO 생성 시 필수)

MLX (training_utils.py:535-662):

  • mlx_lm_lora.trainer.grpo_trainer.train_grpo + GRPOTrainingArgs 사용
  • ref_model=None: self-reference 방식으로 별도 reference model 없이 학습 (KL penalty 미적용, 메모리 절약)
  • end_answer_token="": Self-awareness 태스크에서는 R1 style </answer> 토큰 미사용
  • MLX native optimizer (mlx.optimizers.Adam) 사용

4.4 LoRA 설정 (공통)

{
  "lora_rank": 8,
  "lora_scale": 1.0,
  "lora_dropout": 0.05,
  "target_modules": "all-linear"
}
  • lora_alpha = lora_rank * lora_scale = 8 (rank와 동일)
  • 모든 linear layer에 LoRA 적용

5. 학습 플로우 (train_selfaware_grpo.py)

5.1 실행

python train_selfaware_grpo.py --config configs/config_8b_qlora.json
python train_selfaware_grpo.py --config configs/config_8b_qlora.json --data-variant selfaware-v3

5.2 Phase 1: GRPO 학습

Config 로드 → selfaware_grpo 섹션 추출
↓
Data/Adapter 경로 설정 (--data-variant로 override 가능)
↓
wandb 초기화 (tags: ["grpo", BACKEND, "selfaware", "selfaware-grpo"])
↓
run_grpo_training() 호출
  ├─ 모델 + 토크나이저 로드
  ├─ 데이터 변환 (JSONL → prompt + expected_answer + type)
  ├─ GRPO Trainer 구성 (reward_funcs=selfaware_reward_func)
  ├─ 학습 실행
  └─ Adapter 저장

5.3 Phase 2: 5-task Cross-Eval

학습 완료 후 즉시 cross-evaluation을 실행하여 GRPO 학습이 다른 태스크에 미치는 영향을 측정한다:

GRPO adapter (final checkpoint) 로드
↓
5개 태스크 validation set에 대해 평가
  ├─ exploretom
  ├─ selfaware
  ├─ gsm8k
  ├─ triviaqa
  └─ mbpp / humaneval
↓
결과를 wandb에 기록 (cross_eval/selfaware-grpo_on_{task})

6. 소스 코드 참조 요약

구성요소MLXCUDA
IDK 패턴training_utils.py:78-98training_utils_cuda.py:517-537
IDK 판별 함수training_utils.py:101-104training_utils_cuda.py:540-543
데이터 로드training_utils.py:112-148training_utils_cuda.py:551-585
Reward 함수training_utils.py:151-186training_utils_cuda.py:588-620
GRPO Trainertraining_utils.py:535-662training_utils_cuda.py:623-760
학습 엔트리포인트train_selfaware_grpo.py (공통)
GRPO 하이퍼파라미터configs/config_8b_qlora.json:84-99