Mini-R1: GRPO와 카운트다운 게임을 통한 DeepSeek-R1 추론 재현

Hugging Face는 DeepSeek-R1에서 관찰된 '아하 순간'—사람의 피드백 없이도 모델이 더 많은 사고 시간을 할당하고 접근 방식을 다시 평가하는 능력—을 그룹 상대 정책 최적화(GRPO)와 수학 퍼즐인 카운트다운 게임을 사용해 재현하는 기술 튜토리얼을 발표했다.

그룹 상대 정책 최적화(GRPO) 설명

GRPO는 대규모 언어 모델(LLM)의 추론 능력을 향상시키기 위해 설계된 강화 학습(RL) 알고리즘이다. DeepSeekMath 논문에서 소개된 GRPO는 전통적인 근접 정책 최적화(PPO)에서 가치 함수 모델의 필요성을 제거한다. 대신 그룹 점수에서 기준점을 추정함으로써 메모리 사용량과 계산 부담을 크게 줄인다.

GRPO 과정은 네 가지 주요 단계로 구성된다:

  1. 샘플링: 현재 정책이 단일 프롬프트에 대해 여러 출력을 생성한다.
  2. 보상 점수 매기기: 각 생성물은 규칙 기반 또는 결과 기반의 보상 함수를 통해 점수를 매긴다.
  3. 우위 계산: 그룹의 평균 보상이 기준점이 되며, 각 개별 해결책의 우위는 이 정규화된 그룹 평균에 대해 계산된다.
  4. 정책 최적화: 계산된 우위와 KL 발산 항을 포함한 GRPO 목적 함수를 최대화하도록 정책을 최적화한다.

기술적 구현 및 설정

Mini-R1 실험은 다음과 같은 기술 스택과 구성 사용했다:

  • 기본 모델: Qwen/Qwen2.5-3B-Instruct. 3B 파라미터 모델을 선택한 이유는 일반적으로 모델이 추론 과정을 효과적으로 학습하기 위해 1.5B 이상의 파라미터가 필요하다는 관찰 때문이다.
  • 데이터셋: Jiayi-Pan/Countdown-Tasks-3to4, 3~4개의 숫자로 구성된 퍼즐로 이루어져 있다.
  • 하드웨어: 4x NVIDIA H100 80GB GPU.
  • 소프트웨어 스택: Hugging Face trl(GRPOTrainer 사용), transformers, datasets, accelerate, deepspeed, vLLM을 사용해 생성을 가속화한다.

보상 함수

인간 레이블 없이 정확성을 검증하기 위해, 다음 두 가지 규칙 기반 보상 함수를 사용했다:

  • 형식 보상: 모델이 </tool_call> [thinking] </tool_call><answer> [answer] </answer> 형식을 따르도록 보장한다.
  • 정확성 보상: <answer> 태그에서 식을 추출하고, 그 식이 목표 숫자와 일치하며 제공된 모든 숫자를 정확히 한 번씩 사용하는지 확인한다.

분산 학습 구성

학습은 DeepSpeed Zero-3과 vLLM을 사용해 실행되었다. 4개 GPU 환경에서 3개 GPU는 학습에 전용되고(num_processes 3), 마지막 GPU는 vLLM 생성에만 사용되었다. 전체 450단계 학습은 약 6시간이 소요되었으며, 각 단계는 45~60초 정도 걸렸다.

학습 결과 및 관찰

학습 성능은 TensorBoard를 통해 추적되었으며, 25단계마다 체크포인트가 저장되었다. 모델의 행동은 세 가지 명확한 단계로 진화했다:

  • 단계 50: 모델은 요구되는 </tool_call> 및 <answer> 형식을 성공적으로 학습했다.
  • 단계 100: 성공률이 약 25%에 도달했다. 모델은 시도와 오류 과정을 자연어로 설명하며 '추론'을 시작했다.
  • 단계 200: 성공률이 약 40%에 도달했다. 모델은 자연어 추론에서 '프로그래밍 실행' 스타일로 전환되었으며, 여러 조합을 나열하고 결과를 체계적으로 검토하는 방식을 사용했다.
  • 단계 450: 성공률이 50%에 도달했다. 모델은 프로그래밍 추론 형식을 유지했으며, 성능은 느리지만 지속적으로 향상되었다.

하이퍼파라미터 튜닝

초기 DeepSeekMath 하이퍼파라미터(학습률 1e-6, beta 0.04)를 사용한 시도는 150단계 이후 불안정한 학습을 초래했다. 학습률을 5e-7로 낮추고, beta(KL 계수)를 0.001로 조정함으로써 안정성을 확보했다.

추론 방식의 변화 분석

자연어 기반 추론에서 프로그래밍 실행 스타일로의 전환은 다음과 같은 몇 가지 가능한 요인에 기인한다:

  • 모델 용량: Qwen 2.5 3B는 DeepSeek에서 사용한 더 큰 모델에 비해 복잡한 자연어 추론을 유지하기에 너무 작을 수 있다.
  • 보상 명세: 보상 함수가 의도치 않게 '보상 해킹'을 유도했을 수 있다. 즉, 자연어보다 더 효율적인 수학적 약어를 찾는 방식으로 모델이 학습했다.
  • 과제 특이성: 카운트다운 게임에만 집중적으로 학습한 결과, 해당 퍼즐 유형에 가장 효과적인 해결 방식으로 자연스럽게 몰렸을 수 있다.
  • 학습 기간: 모델이 충분히 오래 학습되지 않았을 수 있다. 원본 R1 논문에서는 8,000단계 이상의 학습이 기록되어 있다.

"이 실험은 GRPO가 LLM의 추론 능력을 자율적으로 향상시킬 수 있음을 보여준다. 특히, 복잡한 문제 해결에서 자연어 추론에서 프로그래밍 스타일로의 전환은 모델이 더 효율적인 사고 방식을 스스로 발견할 수 있음을 시사한다." — @Jiayi-Pan

Sources

관련