Process Reward Agents for Steering Knowledge-Intensive Reasoning

Motivation

• Known: 기존 연구는 process reward models (PRMs), 특히 retrieval-augmented PRM(Med-PRM 등)을 통해 의료 추론에서 중간 단계를 평가해왔으며, Med-S3처럼 policy와 reward model을 공동 학습하는 self-evolving 프레임워크도 존재한다.
• Gap: 기존 PRM들은 완성된 reasoning trajectory에 대해 post hoc으로만 점수를 매기기 때문에, 오류가 누적된 이후에나 감지되며 동적인 inference 절차(search-based decoding)에 통합될 수 없다는 근본적 한계가 있다.
• Why: 의료와 같은 고위험 도메인에서는 최종 답만이 아니라 추론 과정 전체가 방어 가능해야 하며, 새로운 LRM backbone마다 의학 지식에 맞춰 재학습하는 것은 비용이 매우 크기 때문에, frozen policy와 도메인 특화 보상 모듈을 분리하는 패러다임은 실용적으로 중요하다.
• Approach: PRA는 frozen reasoning policy와 상호작용하는 별도의 에이전트로서, 매 reasoning step마다 현재 추론 흔적을 관찰하고 필요시 외부 의학 근거를 검색한 뒤 로컬 보상 신호를 부여하여 search-based decoding(후보 trajectory의 랭킹 및 pruning)을 실시간으로 유도한다.

Achievement

Figure 2. Performance on MedQA under inference time scaling.

1. MedQA 신기록: Qwen3-4B-Instruct와 결합해 MedQA에서 81.9% 정확도를 달성, 4B 규모에서 state-of-the-art 성능을 기록했다.
2. 정책 모델 간 일반화: PRA는 학습에 사용되지 않은 0.5B~8B 규모의 다양한 frozen policy 모델에도 적용되어 최대 25.7%의 정확도 향상을 정책 모델 업데이트 없이 달성했다.
3. 추론 시점 스케일링 우위: 동일한 policy sampling budget 하에서 PRA는 inference-time scaling에 따라 지속적으로 성능이 향상되는 반면, self-consistency는 조기에 성능이 포화된다.
4. 보상 위치의 중요성 검증: reward granularity와 timing에 대한 ablation을 통해 성능 향상이 완료 시점이 아닌 중간 단계에서 보상을 적용하는 데서 비롯됨을 입증했다.

How

Figure 1. Overview of our approach. A Process Reward Agent (PRA) observes the reasoning trace generated by the frozen po

• 문제를 (q, π, D) → ŷq를 출력하는 parameterized inference procedure Gϕ로 형식화하고, 정답 정확성 함수 C를 이용해 목표를 정의함.
• π는 파라미터가 고정된(frozen) reasoning model로, reasoning trace τ = (s1,...,sK)를 autoregressive하게 생성함.
• PRA는 각 step에서 현재 trace τt를 관찰하고, 필요시 fixed knowledge base D에서 검색을 수행할지 결정한 뒤 step-level reward를 산출함.
• 이 reward 신호는 search-based decoding(후보 경로의 ranking 및 pruning)에 사용되어 policy의 생성 과정을 실시간으로 유도함.
• 여러 의료 추론 벤치마크(MedQA 포함)에서 matched policy sampling budget 하에 강력한 decoding baseline들과 비교 평가함.

Originality

• 기존 retrieval-augmented PRM(Med-PRM 등)이 post hoc으로만 완성된 trajectory를 평가하던 것과 달리, PRA는 online, step-wise 보상을 제공하는 최초의 방법론임.
• retrieval과 verification을 policy로부터 완전히 분리(decouple)하여 별도의 agent가 담당하게 함으로써, frozen reasoner를 domain-specific reward module과 결합하는 새로운 패러다임을 제시함.
• 검색이 완료 후가 아니라 생성 도중 매 스텝마다 발생하여 search-based decoding(ranking/pruning)에 직접 통합될 수 있다는 점이 수학적 추론 영역의 retrieval-augmented PRM들과도 차별화됨.

Limitation & Further Study

• 매 스텝마다 검색 및 보상 계산이 필요하므로 inference-time 계산 비용과 지연(latency)이 상당히 증가할 가능성이 있으나 이에 대한 상세한 비용-효율 분석이 제한적일 수 있음.
• 평가가 의료 분야(MedQA 등)에 집중되어 있어, 다른 지식집약적 도메인(법률, 과학 등)으로의 일반화 가능성은 추가 검증이 필요함.
• knowledge base D의 품질과 최신성에 성능이 크게 의존할 것으로 보이며, 저품질 또는 편향된 검색 코퍼스에서의 견고성에 대한 논의가 부족함.
• PRA 자체(reward agent)의 학습 방법과 그 일반화 한계(예: 매우 큰 policy 모델이나 전혀 다른 도메인)에 대한 후속 연구가 필요함.

Evaluation

총평: 지식집약적 추론에서 process reward를 post hoc이 아닌 online·step-wise 방식으로 제공한다는 아이디어는 참신하고 실용적 파급력이 크며, 다양한 frozen policy 모델에 대한 일반화 실험 결과가 이를 잘 뒷받침한다. 다만 계산 비용과 타 도메인으로의 확장성에 대한 추가 검증이 향후 연구로 남아 있다.