지식 집약적 추론 조향을 위한 프로세스 보상 에이전트

의료와 같이 전문 지식이 필요한 분야에서 AI가 중간 추론 단계의 오류를 스스로 감지하고 외부 근거를 찾아 수정할 수 있게 한다. 모델 자체를 재학습시키지 않고도 외부 보상 에이전트만으로 소형 모델의 성능을 대형 모델 수준으로 끌어올릴 수 있음을 입증했다.

기존의 LLM은 문장을 생성할 때 이전 토큰들을 바탕으로 다음 토큰의 확률 분포를 계산하는 Autoregressive 방식을 사용한다. 하지만 의료 지식처럼 복잡한 도메인에서는 모델 내부의 파라미터에 저장된 지식만으로는 중간 추론 단계의 논리적 타당성을 검증하기 어렵고, 한 번 발생한 작은 오류가 전체 결론을 망치는 현상이 발생한다.

PRA는 이를 해결하기 위해 추론을 수행하는 '두뇌(Frozen Policy)'와 이를 감시하고 가이드하는 '코치(Reward Agent)'를 분리한다. 코치는 모델이 내뱉은 중간 문장을 보고, 필요하다면 외부 의학 데이터베이스에서 관련 문헌을 검색하여 그 근거를 바탕으로 현재 단계에 점수를 매긴다. 이는 마치 수험생이 문제를 풀 때 옆에서 튜터가 참고서를 찾아보며 각 풀이 단계가 맞는지 실시간으로 채점해 주는 것과 같다.

결과적으로 모델은 보상이 높은 경로로만 추론을 이어가게 되며, 이 과정에서 Softmax 함수를 통해 계산된 확률값이 낮은 경로는 조기에 차단된다. 이를 통해 모델의 크기가 작더라도 외부 지식과 정교한 탐색 전략을 결합하여 대형 모델에 필적하는 정확한 논리 전개를 수행할 수 있게 된다.

전체 시스템은 고정된 추론 모델(π), 프로세스 보상 에이전트(µ), 그리고 검색기(ρ)의 세 가지 모듈로 구성된다. 추론 모델이 중간 단계를 생성하면 PRA가 이를 입력받아 '검색(Search)' 또는 '보상(Reward)' 액션을 선택한다. [현재 추론 궤적 τ_t 입력 → µ_act 네트워크 연산 → 액션 ˆa_t 결정] 과정을 거치며, 검색이 필요하다고 판단되면 외부 지식 베이스 D에서 관련 문서를 추출한다.

보상 단계에서는 검색된 문서 D_t와 추론 단계 s_t를 함께 평가하여 수치화된 보상 r_t를 산출한다. [추론 궤적과 문서 세트 입력 → µ_rwd 네트워크 연산 → 0에서 1 사이의 스칼라 값 출력] 이 값은 해당 추론 단계가 의학적으로 얼마나 타당한지를 나타내는 지표가 된다.

추론 시에는 Beam Search 알고리즘을 활용하여 여러 후보 경로를 동시에 탐색한다. 각 단계에서 누적 보상 R을 계산하고, 상위 B개의 경로만 남기고 나머지는 제거(Pruning)한다. [각 단계의 보상 r_i 합산 → 누적 점수 R 산출 → 점수 기반 경로 필터링] 또한 효율적인 처리를 위해 서로 다른 질문들을 단계별로 묶어 처리하는 Stage-Level Batching 기법을 적용하여 GPU 활용도를 극대화했다.

MedQA 벤치마크에서 Qwen3-4B 모델을 기반으로 80.8%의 정확도를 기록했다. 이는 기존의 강력한 베이스라인인 RAG + Self-Consistency(76.7%)보다 4.1%p 높은 수치이며, 4B 규모 모델 중 세계 최고 수준이다. 특히 Medbullets, MedMCQA 등 학습 시 보지 못한 외부 데이터셋에서도 평균 4.8%p의 성능 향상을 보이며 강력한 일반화 능력을 입증했다.

Ablation Study 결과, 단순히 검색 결과를 입력창에 넣어주는 RAG 방식보다 PRA를 통한 단계별 보상 가이드가 성능 향상에 더 결정적인 역할을 했음이 밝혀졌다. 또한 모델의 크기가 작을수록 PRA 도입에 따른 성능 향상 폭이 컸는데, Qwen2.5-0.5B 모델의 경우 정확도가 28.4%에서 54.1%로 약 90.5%의 상대적 향상을 기록했다.

추론 시간 스케일링 분석에서는 샘플 수를 늘릴수록 성능이 빠르게 포화되는 Self-Consistency와 달리, PRA는 연산 자원을 더 투입할수록 성능이 지속적으로 우상향하는 Pareto Frontier를 형성함을 확인했다.

PRA는 Qwen3-4B-Instruct를 기반으로 파인튜닝되었으며, 보상 예측을 위한 리드아웃과 액션 결정을 위한 리드아웃 두 개의 헤드를 공유하는 구조를 가진다. 학습 데이터는 Qwen3-235B 모델을 Teacher로 활용하여 생성되었으며, 외부 지식 유무에 따른 모델의 판단 변화량인 Margin Shift(Δm)를 기준으로 검색 라벨을 자동 생성하는 기법을 도입했다.

수학적으로는 Δm = m - m_d 식을 사용하며, 여기서 m은 검색 결과가 없을 때의 정답 확률 마진, m_d는 검색 결과가 있을 때의 마진이다. [두 조건의 로그 확률 차이 계산 → 차이값의 절대값 산출 → 임계값 비교를 통한 검색 필요성 판단] 이 방식은 모델이 스스로 판단하기 모호한 지점에서만 검색을 수행하도록 유도하여 효율성을 높인다.

추론 아키텍처 측면에서는 분산된 추론 깊이를 가진 여러 쿼리를 동기화하기 위해 전역 큐(Global Queue)를 관리한다. 각 트레이스는 REASON, REWARD, SEARCH, DONE 상태를 순환하며, 동일한 상태에 있는 트레이스들을 배치로 묶어 처리함으로써 가변 길이 추론 상황에서도 높은 처리량을 유지한다.

PRA가 모든 잘못된 중간 단계나 환각(Hallucination)을 완벽하게 제거하지는 못할 수 있음을 명시했다. 또한 본 연구는 방법론적 기여에 집중하고 있으며, 실제 의료 현장에 즉시 투입 가능한 수준의 완성된 시스템은 아니라는 점을 한계로 언급했다.