기술 요약: POISE (내부 상태 가치 추정을 통한 정책 최적화)

문제 제기

검증 가능한 보상을 활용한 강화 학습 (RLVR) 은 수학적 추론과 같은 복잡한 작업에서 대규모 추론 모델 (LRM) 을 개선하는 핵심 동력입니다. 그러나 안정적이고 효율적인 RLVR 은 정책 기울기의 분산을 줄이기 위해 정확한 기준선 (baseline) 추정에 크게 의존합니다. 기존 접근 방식은 상당한 계산 비용을 수반합니다:

• PPO (근접 정책 최적화): 매 업데이트마다 생성된 전체 시퀀스를 처리하는 별도의 LLM 규모 비평가 (critic) 모델의 훈련이 필요하여, 메모리 소비와 최적화 복잡성이 약 2 배 증가합니다.
• GRPO (그룹 상대 정책 최적화): 별도의 비평가를 피하기 위해 프롬프트당 샘플링된 롤아웃 그룹의 평균 보상을 기반으로 기준선을 추정합니다. 이는 "샘플을 위해 매개변수를 교환"하는 방식으로, 안정적인 기준선을 유지하기 위해 프롬프트당 여러 번의 롤아웃이 필요합니다. 고정된 계산 예산 하에서는 배치 내 프롬프트의 다양성이 줄어들어 기울기 분산이 증가합니다. 또한, 모든 롤아웃이 동일한 보상 (영의 이득) 을 받는 '퇴화 (degenerate)' 프롬프트에 상당한 계산 자원이 낭비됩니다.

이 연구가 다루는 핵심 질문은 다음과 같습니다: 정책 훈련 중에 이미 수행된 계산으로부터 효과적인 기준선을 추출하여 추가 자원의 필요성을 제거할 수 있는가?

방법론: POISE

저자들은 POISE(내부 상태 가치 추정을 통한 정책 최적화) 를 제안합니다. 이는 정책 모델의 자체 내부 표현을 활용하여 거의 비용 없이 가치 함수 $V^\pi(x)$를 추정하는 방법입니다.

1. 내부 상태 가치 추정

별도의 비평가를 훈련하는 대신, POISE 는 경량 프로브 (probe) 를 훈련하여 정책 모델의 숨겨진 상태 (hidden states) 와 순전파 (forward pass) 중 수집된 토큰 통계를 기반으로 검증 가능한 기대 보상을 직접 예측합니다.

• 입력 특징: 프로브는 세 가지 상호 보완적인 신호를 활용합니다:
  1. 프롬프트 수준 특징: 마지막 프롬프트 토큰의 숨겨진 상태로, 프롬프트에 대한 모델의 표현과 예상 난이도를 포착합니다.
  2. 궤적 수준 특징: 추론 궤적 끝의 숨겨진 상태로, 프롬프트만으로는 존재하지 않는 가치 관련 정보를 노출합니다.
  3. 토큰 수준 엔트로피 통계: 경량 불확실성 특징입니다.
• 훈련 목적: 프로브는 프롬프트 수준의 가치 $V^\pi(x)$를 예측하도록 훈련됩니다. 평가 중인 특정 롤아웃에 대한 정보가 목표에 유출되는 것을 방지하기 위해, 저자들은 leave-one-out 몬테카를로 목표를 사용합니다. $K$개의 샘플링된 롤아웃이 있는 프롬프트 $x$에 대해, $i$번째 롤아웃의 목표는 나머지 $K-1$개 롤아웃의 평균 보상입니다.

2. 편향 제거를 위한 교차 롤아웃 구성

중요한 이론적 과제는 정책 기울기의 편향 없는 성질을 유지하기 위해 기준선이 행동 (샘플링된 응답) 에 조건부 독립적으로 유지되도록 보장하는 것입니다. 기준선이 평가하는 동일한 롤아웃에 의존하면 편향이 발생합니다.

• 해결책: POISE 는 교차 롤아웃 구성을 사용합니다. 각 프롬프트에 대해 두 개의 독립적인 롤아웃 $y^{(1)}$과 $y^{(2)}$가 샘플링됩니다.
  • $y^{(1)}$의 기준선은 $y^{(2)}$의 내부 상태를 사용하여 예측됩니다.
  • $y^{(2)}$의 기준선은 $y^{(1)}$의 내부 상태를 사용하여 예측됩니다.
• 이를 통해 특정 롤아웃을 업데이트하는 데 사용되는 기준선은 독립적인 샘플에서 파생되어 $E[b(x) \nabla \log \pi(y|x)] = 0$ 조건을 만족합니다.

3. 정책 최적화 통합

POISE 는 이 추정기를 PPO 스타일의 최적화 루프에 통합합니다:

• 온라인 훈련: 가치 추정기는 최근 궤적의 슬라이딩 버퍼에서 정책과 함께 훈련됩니다. 특징은 정책 업데이트에 이미 필요한 순전파에서 추출되므로, 계산 오버헤드는 무시할 수준입니다.
• 이득 계산: 이득은 $A^{(i)}(x) = R(x, y^{(i)}) - b^{(i)}(x)$로 계산되며, 여기서 $b^{(i)}(x)$는 교차 롤아웃 기준선입니다.
• 효율성: 기준선 추정을 위해 대규모 그룹 (예: 8 개 롤아웃) 이 필요한 GRPO 와 달리, POISE 는 롤아웃 쌍 하나만 필요합니다. 이는 고정된 계산 예산을 배치 내 더 많은 고유한 프롬프트에 할당할 수 있게 하여, 이론적으로 기울기 분산을 줄입니다.

주요 기여

1. 새로운 기준선 추정: 별도의 LLM 규모 비평가나 대규모 롤아웃 그룹의 필요성을 제거하고, 정책 모델의 내부 상태에서 기준선을 유도하는 방법을 제시합니다.
2. 이론적 편향 없음: 궤적 조건부 특징을 활용하면서도 편향 없는 정책 기울기에 필요한 조건부 독립성을 보장하는 교차 롤아웃 아키텍처를 제안합니다.
3. 계산 효율성: 경량 프로브 (선형 회귀) 가 별도의 비평가와 비교할 때 동등하거나 더 나은 가치 추정 성능을 달성하면서도 훨씬 적은 계산 자원을 필요로 함을 입증합니다.
4. 일반화 가능성: 내부 상태 신호가 수학적, 코딩, 도구 호출, 지시 따르기 등 다양한 도메인과 모델 규모에서 일반화됨을 보여줍니다.

실험 결과

이 방법은 오림피아드 수준의 수학적 추론 벤치마크 (AMC, AIME, HMMT, BRUMO) 에서 Qwen3-4B 와 DeepSeek-R1-Distill-Qwen-1.5B 를 대상으로 평가되었습니다.

• 성능: POISE 는 최첨단 GRPO 기반 알고리즘인 DAPO와 비교 가능한 성능을 달성합니다.
  • Qwen3-4B 에서 POISE 는 평균 Avg@32 점수 0.500을 기록하여 DAPO 의 0.508과 근접했으며, AMC23 및 HMMT25 와 같은 특정 벤치마크에서는 DAPO 를 능가했습니다.
  • DeepSeek-R1-Distill-Qwen-1.5B 에서 POISE 는 평균 점수를 DAPO 의 0.296 대비 0.303으로 향상시켰습니다.
• 훈련 효율성: POISE 는 현저히 적은 월클록 (wall-clock) 시간이 필요합니다. DeepSeek-R1-Distill-Qwen-1.5B 에서 단일 B200 GPU 로 비교 가능한 성능에 도달하는 데 POISE 는 18 시간이 소요된 반면 DAPO 는 24 시간이 걸렸습니다. Qwen3-4B 의 경우 36 시간 대 49 시간이었습니다.
• 안정성: POISE 의 기울기 노름은 DAPO 보다 일관되게 낮아, 감소된 기울기 분산과 더 안정적인 최적화를 나타냅니다.
• 가치 추정 품질: 내부 상태 프로브는 경험적 보상과 0.870의 피어슨 상관관계를 달성하여, 동일한 작업에서 별도로 훈련된 LLM 규모 비평가 (0.676) 보다 우수한 성능을 보였습니다.
• 일반화: 추정기는 코딩, 도구 호출, 지시 따르기 작업 전반에 걸쳐 높은 성능을 유지했으며, 종종 이러한 도메인에서 별도의 비평가보다 우수한 결과를 보였습니다.

중요성과 주장

이 논문은 POISE 가 대규모 추론 모델을 위한 안정적이고 확장 가능한 RLVR 을 위한 계산 효율적인 경로를 제공한다고 주장합니다. 모델의 자체 내부 표현을 활용함으로써 이 방법은 다음과 같은 이점을 제공합니다:

• 별도의 비평가 모델 훈련에 따른 메모리 및 계산 오버헤드를 제거합니다.
• 그룹 상대 기준선 (GRPO) 의 샘플링 비효율성을 제거하여 고정된 예산 내에서 더 높은 프롬프트 다양성을 가능하게 합니다.
• 영의 이득을 산출하는 퇴화 프롬프트에 대한 계산 낭비를 방지합니다.
• 주로 진단 도구로만 사용되던 내부 상태가 강화 학습을 위한 실용적이고 고충실도의 최적화 신호로 기능할 수 있음을 입증합니다.

저자들은 POISE 가 그룹 상대 기준선이나 보조 가치 모델 없이도 더 안정적이고 효율적인 정책 최적화를 가능하게 하여, 자원 제약 하에서 대규모 추론 모델을 훈련하기 위한 실현 가능한 대안이 된다고 결론지었습니다.