기술 요약: RREDCoT – 추론 모델을 위한 세그먼트 수준 보상 재분배

문제 정의

최근 추론 언어 모델의 발전은 특히 그룹 상대 정책 최적화(GRPO) 알고리즘과 검증 가능한 보상(RLVR) 패러다임을 결합한 강화 학습(RL) 미세 조정에 크게 의존하고 있습니다. 이러한 설정에서 모델은 최종 답변을 생성하기 전에 사고 사슬(Chain-of-Thought, CoT) 트레이스를 생성합니다. 그러나 여기서 중요한 병목 현상이 존재합니다: 보상 신호가 일반적으로 전체 CoT 트레이스가 완료되고 최종 답변이 검증된 후에만 할당된다는 점입니다. 이는 **지연된 보상 문제(delayed reward problem)**를 야기합니다.

표준적인 GRPO 접근 방식은 RL의 몬테카를로 방법(Monte Carlo methods)에 해당하며, 보상이 전체 궤적에 균등하게 분배되어 개별 추론 단계에 대한 직접적인 감독을 제공하지 못하기 때문에 높은 분산을 겪습니다. 몬테카로 샘플링(MCS)은 중간 상태의 가치를 추정할 수 있지만, 모든 세그먼트에 대해 여러 완결본(completions)을 생성해야 하므로 긴 문맥에서 고밀도 신용 할당(credit assignment)을 수행할 때 과도한 계산 오버헤드를 초래합니다. 단계별 분석을 위해 판사 모델(judge models)을 사용하거나 모델 통계로부터 중간 유틸리티를 추출하는 기존의 대안들은 종종 노이즈를 유발하거나 추가적인 계산 자원을 요구합니다.

방법론: RREDCoT

저자들은 RUDDER 알고리즘의 원리를 CoT 생성의 특정 구조에 적응시킨 RREDCoT(Reward REDistribution for Chain of Thoughts)를 제안합니다. 핵심 목표는 추가적인 모델이나 추가적인 생성 단계 없이도 더 조밀한 학습 신호를 제공하기 위해 최종 보상을 CoT 트레이스의 중간 세그먼트로 재분배하는 것입니다.

1. MDP 정식화

CoT 생성 과정은 상태 $s_t$가 원래의 쿼리와 $t-1$까지 생성된 세그먼트 시퀀스로 구성되는 마르코프 결정 과정(MDP)으로 모델링됩니다. 행동 공간은 다음과 같이 분할됩니다:

• $Y$ (Answer): 최종 출력을 트리거하는 행동.
• $U$ (Continuation): 추가적인 CoT 세그먼트를 생성하는 행동.
  생각을 생성하는 것에 대한 즉각적인 보상은 0이며, 보상은 올바른 최종 답변을 생성할 때만 실현됩니다.

2. 최적 보상 재분배

RREDCoT는 연속적인 상태 간의 기대 수익 차이를 기반으로 최적의 보상 재분배 $\sigma$를 도출합니다. 이 방법은 (원래의 RUDDER에서 사용된 것과 같은) 보조 LSTM이나 광범한 샘플링을 피하기 위해 생성 언어 모델 자체를 사용하여 가치 함수를 근사합니다.

재분배는 다음과 같이 계산됩니다:
$$\sigma_t \propto \hat{R}_{t+1} - \hat{R}_t + \hat{v}_{t+1} - \hat{v}_t$$
여기서:

• $\hat{R}_t$는 즉각적인 보상 추정치(현재 트레이스에 대한 참조 정답의 확률)입니다.
• $\hat{v}_t$는 상태의 추정 가치(미래 완결본의 기대 유틸리티)입니다.

이러한 값들을 효율적으로 추정하기 위해, 저자들은 참조 솔루션 경로($Y^\star, U^\star$)에 대한 중요도 샘플링(importance sampling)을 사용하는 확률 보상(Probability Reward, PR) 추정기를 채택합니다. 이를 통해 모델은 현재 세그먼트로부터 높은 유틸리티 상태에 도달할 가능성을 추정하는 역학 함수(dynamics function, world model) 역할을 수행할 수 있습니다.

3. 하이브리드 세그멘테이션 전략

신용 할당을 실행 가능하게 만들기 위해, RREDCoT는 하이브리드 키워드-엔트로피 세그멘테이션 전략을 도입합니다.

1. 초기 분할: 트레이스는 일반적인 키워드(예: 줄바꿈)에 의해 분할됩니다.
2. 반복적 병합: 목표 세그먼트 수에 도달할 때까지 결합된 엔트로피가 가장 낮은 연속 세그먼트들을 병합합니다.
3. 근거: 높은 엔트로피 토큰은 높은 다수성(plurality)을 나타내는 지점(결정 지점)을 의미하며, 중간 기여도를 추정하기 위한 자연스러운 "출구 지점(exit points)" 역할을 합니다. 이 전략은 세그먼트들이 최대한 탈동조화(decoupled)되도록 보장하여 가치 함수 추정의 편향을 최소ize합니다.

4. RL 목적 함수와의 통합

재분배된 보상은 표준 정책 경사(policy gradient) 목적 함수에 통합됩니다. 이 방법은 **수익 동등성(return-equivalence)**을 보장하는데, 이는 재분배된 보상의 합이 원래 에피소드 수익과 같음을 의미하며, 최적의 정책을 보존하면서 수렴을 가속화합니다.

주요 기여

1. RREDCoT 알고리즘: 가치 추정을 위해 생성 모델 자체를 활용함으로써, 추가 모델이나 추가 생성 단계 없이 CoT 트레이스에 대한 실행 가능한 신용 할당 및 보상 재분배를 수행하는 알고리즘입니다.
2. 엔트로피 기반 세그멘테이션: 중간 기여도 추정을 최적화하기 위해 키워드 기반 분할과 엔트로피 기반 병합의 균형을 맞춘, RREDCoT에 특화된 새로운 세그멘테이션 전략입니다.
3. 이론적 및 경험적 분석: 제안된 추정기를 몬테카로 샘플링 및 다양한 속성 추출 방법(예: Shapley values, Leave-One-Out)과 비교하는 상세한 분석을 통해, 정답에 조건화되지 않은 속성 추출 방법이 보상 재분배와는 다른 측면의 생성을 포착함을 입증했습니다.

실험 결과

저자들은 AIME, MATH, Minerva 데이터셋을 포함한 여러 벤치마크에서 RREDCoT를 평가했습니다.

• 긴 문맥 성능: 생성 길이가 최대 25k 토큰에 달하는 Numina-CoT 데이터셋에서, RREDCoT는 표준 GRPO에 비해 우수한 성능을 보였습니다. 예를 들어, Qwen3-4B Instruct 모델에서 RREDCoT는 AIME24 점수 0.908을 달성하여 GRPO의 0.850보다 높은 성적을 기록했습니다.
• 소규모 정교화: 온라인 정교화를 위해 소규모 추론 모델(1.5B 파라미터)에 적용했을 때, RREDCoT는 하이퍼파라미터 튜닝을 거친 GRPO 모델과 비교하여 AMC23 및 MATH500 벤치마크에서 경쟁력 있거나 더 우수한 결과를 보여주었습니다.
• 분산 및 편향: 연구 결과, 절단된 몬테카로 샘플링(truncated Monte Carlo sampling)은 유효한 완결본의 상당한 손실을 초래하여 편향을 유발하는 반면, RREDCoT의 추정기는 매우 적은 계산 비용으로 몬테카로 샘플링 결과와 잘 일치함을 보여주었습니다.
• 상관관계 분석: 논문은 그래디언트 기반 방법과 Leave-One-Out(LOO) 속성 추출 방법이 서로 높은 상관관계를 갖지만, RREDCoT와는 상관관계가 낮다는 것을 발견했습니다. 이는 속성 추출 방법이 보상 재분배와는 다른 신호 특성을 포착함을 시사하며, 속성 추출만을 사용하여 학습 가중치를 사용하는 것이 중복되거나 최적이 아닐 수 있음을 나타냅니다.

의의 및 주장

본 논문은 RREDCoT가 추론 모델의 RL 미세 조정에서 발생하는 높은 분산과 지연된 보상 문제를 해결한다고 주장합니다. 모델 자신의 예측 능력을 기반으로 보상을 중간 세그먼트로 재분배함으로써, 이 방법은 몬테카로 샘플링의 계산적 제약을 피하면서 수렴을 가속화하고 최종 성능을 향상시킵니다.

저자들은 이 접근 방식이 최적 보상 재분배의 **실행 가능한 근사치(tractable approximation)**임을 강조합니다. 또한 이 방법이 참조 솔루션의 가용성과 일부 하위 목표 구조(예: 수학 문제의 고유한 해답)에 의존한다는 점을 명시합니다. RREDCoT는 단순한 GRPO에 비해 1.5~2배의 계산 오버헤드가 발생하지만, 전체 몬테카로 샘플링보다는 훨씬 효율적입니다. 이 연구는 RREDCoT를 훈련 신호 증폭기이자 분석 도구로 위치시키며, 전통적인 속성 추출 기술과는 구별되는 신용 할당에 대한 새로운 관점을 제시합니다.