LLM 에이전트를 위한 검증된 핵심 단계 최적화

초록

본 논문은 LLM 기반 에이전트의 장기 과제 수행 시, 실패한 정책 궤적에서 핵심 의사결정 지점을 찾아 검증된 대안을 통해 선호 학습(DPO)을 진행하는 CSO(Critical Step Optimization) 방법을 제안한다. PRM으로 후보 단계를 선정하고, 전문가 모델이 제시한 대안을 정책이 직접 이어 실행해 성공 여부를 검증함으로써, 단계별 보상 추정의 노이즈를 없애고 전체 궤적의 16%만 감독하면서도 GAIA‑Text‑103과 XBench‑DeepSearch에서 각각 37%·26%의 상대적 성능 향상을 달성한다.

상세 분석

CSO는 기존의 두 축, 즉 결과 기반 전체 궤적 보상과 단계별 추정 보상의 한계를 동시에 극복한다. 첫째, 결과만을 이용하는 트래젝터리 레벨 최적화는 모든 단계에 동일한 신호를 부여해 실제로 성공에 기여한 작은 결정들을 흐리게 만든다. 둘째, 단계별 프로세스 보상 모델은 인간이 설계한 루브릭에 의존해 점수를 매기지만, 모델의 편향과 평가 노이즈가 누적돼 학습 효율을 저하시킨다. CSO는 이러한 문제를 “검증된 핵심 단계”라는 개념으로 재정의한다. 핵심 단계는 (1) 정책이 낮은 PRM 점수를 받고, (2) 전문가 모델이 높은 점수를 부여한 지점이며, (3) 해당 대안을 적용했을 때 정책 자체가 이후 과정을 성공적으로 수행할 경우에만 선택된다. 이 과정에서 PRM은 후보 탐색을 가볍게 수행하고, 실제 성공 여부는 정책이 직접 롤아웃함으로써 확인한다. 따라서 보상의 근거가 실제 결과에 기반하므로 추정 오차가 사라지고, 정책이 도달 가능한 행동 공간 내에서만 학습이 이루어진다.

핵심 단계 수집 파이프라인은 다음과 같다. 1) 현재 정책으로 작업을 실행해 실패 궤적을 모은다. 2) 각 단계마다 전문가 모델(Claude‑3.7‑Sonnet)로 K=5개의 대안을 생성한다. 3) PRM이 정책 행동과 대안들을 점수화하고, 저점수·고점수 기준(γ_low=0.45, γ_high=0.65)으로 후보를 추린다. 4) 후보 단계마다 고점수 대안을 적용해 정책이 이어서 실행하도록 하고, 최종 성공 여부를 검증한다. 성공한 경우 (상태, 성공 대안, 원래 실패 행동) 형태의 선호 쌍을 만든다. 이렇게 만든 D_pref를 DPO 손실에 넣어 KL 정규화와 함께 학습한다. 학습은 반복적으로 수행되며, 매 라운드마다 개선된 정책으로 새로운 실패 궤적을 수집해 핵심 단계가 점진적으로 어려워지는 형태의 커리큘럼을 형성한다.

실험에서는 8B 파라미터 CK‑Pro‑8B 모델을 사용했으며, GAIA‑Text‑103(복합 웹 검색·문서 작성)과 XBench‑DeepSearch(다단계 탐색) 두 벤치마크에서 SFT 대비 각각 37%·26%의 상대적 정확도 향상을 기록했다. 특히 전체 궤적 중 16%만을 감독했음에도 불구하고, 기존의 전체 궤적 DPO, 단계별 보상 기반 방법, 그리고 Monte‑Carlo 기반 IPR보다 월등히 높은 성능을 보였다. 이는 핵심 단계가 전체 과제 성공에 미치는 영향이 크다는 가설을 실증적으로 뒷받침한다. 또한 정책 자체가 후속 행동을 수행하도록 함으로써, 전문가 대안이 정책에 “도달 불가능”한 경우를 방지하고, 학습 안정성을 크게 높였다. 비용 측면에서도 PRM 기반 후보 탐색과 제한된 브랜치 롤아웃만으로 충분히 효과적인 데이터를 생성할 수 있어, 대규모 Monte‑Carlo 샘플링에 비해 계산량이 크게 절감된다.

전반적으로 CSO는 (1) 실패 중심 데이터 수집, (2) 고품질 전문가 대안 활용, (3) 실제 성공 검증을 통한 신뢰성 있는 선호 데이터 생성이라는 세 축을 결합함으로써, LLM 에이전트의 장기 의사결정 학습에 필요한 “정확한 신호·적은 비용·정책 친화성”을 동시에 만족한다는 점에서 의미가 크다.