s3: You Don't Need That Much Data to Train a Search Agent via RL

단 2,400개 샘플로 17만 개 샘플 성능 압도: 70배 효율성의 비밀

일리노이 대학교 어바나 샴페인 캠퍼스와 아마존이 공동 개발한 새로운 AI 검색 프레임워크 s3(Search-Select-Serve)가 기존 방법보다 70배 적은 데이터로도 더 뛰어난 성능을 보여주는 것으로 나타났다. 해당 프레임워크는 단 2,400개의 훈련 샘플로 170,000개 샘플로 훈련된 기존 모델들을 능가하는 성과를 기록했다.

s3는 강화학습(RL) 기반의 경량화된 모델 독립적 프레임워크로, 검색 기능과 생성 기능을 분리하여 검색 성능만을 집중적으로 최적화한다. 이 방식은 기존의 검색증강생성(RAG) 시스템과 달리 생성 모델은 고정한 채 검색 에이전트만을 훈련시켜 효율성과 호환성을 극대화했다.

'Gain Beyond RAG' 보상 시스템으로 기존 한계 돌파

연구팀은 기존 RAG 시스템의 세 가지 발전 단계를 분석했다. 초기 Classic RAG는 고정된 검색 방법을 사용했고, Pre-RL-Zero 시기에는 더 적극적인 참여가 가능해졌지만 제로샷 프롬프팅에 의존했다. 현재의 RL-Zero 시대에서는 강화학습을 통한 에이전트 검색이 가능해졌지만, 대부분의 기존 방법들은 검색과 생성을 공동으로 최적화하여 실제 검색 개선 효과를 분리하기 어려웠다.

s3는 이러한 문제를 해결하기 위해 'Gain Beyond RAG(GBR)'라는 새로운 보상 신호를 도입했다. GBR은 s3가 검색한 문서를 사용했을 때의 생성 정확도가 기본 상위 k개 검색 대비 얼마나 향상되었는지를 측정한다. 이를 통해 생성 모델은 동결된 상태를 유지하면서도 검색 성능을 직접적으로 최적화할 수 있게 되었다.

일반 데이터 학습으로 의료 분야 76.6% 정확도 달성한 놀라운 전이 능력

실험 결과 s3는 6개의 일반 도메인 질의응답 벤치마크와 5개의 의료 도메인 질의응답 벤치마크에서 모두 최고 성능을 기록했다. 특히 주목할 점은 일반 데이터로만 훈련된 s3가 의료 질의응답에서도 우수한 성능을 보여준 것이다. 이는 강화학습으로 습득한 검색 기술이 생성 튜닝 방식보다 더 안정적으로 다른 도메인으로 전이된다는 것을 시사한다.

Claude-3-Haiku를 생성 모델로 사용한 실험에서 s3는 일반 도메인에서 평균 58.9%의 정확도를 달성했으며, 의료 도메인에서는 Wikipedia+PubMed+Textbook 말뭉치를 사용했을 때 76.6%의 정확도를 기록했다. 이는 기존의 DeepRetrieval과 Search-R1 모델들을 모두 능가하는 수치다.

실제 사례로 보는 s3의 똑똑한 검색: "2007년 영화" 질문 완벽 해결

s3의 실제 검색 능력을 보여주는 흥미로운 사례가 있다. "재활용 타이어 등 자연·업사이클 소재로 만든 패시브 솔라 하우스 발명가에 대한 영화가 제작된 연도는?"이라는 복잡한 질문에 대해 기존 RAG와 s3의 대응을 비교해보자. 기존 RAG 시스템은 "어스십(Earthship) 발명가에 대한 특정 영화나 그 개봉 연도에 대한 정보가 제공되지 않았다"며 답변을 포기했다. 반면 s3는 첫 번째 검색에서 어스십과 발명가 마이클 레이놀즈(Michael Reynolds)를 찾아낸 후, 두 번째 검색에서 "어스십 발명가에 대한 영화가 언제 만들어졌나?"라는 구체적 질문을 생성했다.

그 결과 s3는 "Garbage Warrior는 어스십 스타일 건축의 발명가인 건축가 마이크 레이놀즈에 대한 2007년 영화다"라는 정확한 정보를 찾아 "2007년"이라는 정답을 도출했다. 이는 s3가 단순히 키워드 매칭이 아닌 맥락을 이해하고 논리적으로 검색 전략을 세운다는 것을 보여준다.

훈련 시간 33배 단축, 인간 판단 일치율 96.4% 달성

s3의 가장 큰 장점은 훈련 효율성이다. Search-R1이 2,100단계(170,000개 예제)의 훈련이 필요한 반면, s3는 단 20단계(2,400개 예제)만으로 더 나은 성능을 달성했다. 훈련 시간 역시 약 33배 단축되어 실제 활용 가능성을 크게 높였다. 또한 s3는 7B 정책 모델을 사용하고 검색기 사전 훈련이 불필요하여 저자원 환경에서도 적용하기 용이하다.

연구팀은 보상 함수 비교 실험을 통해 의미론적으로 정렬된 평가 지표의 중요성도 입증했다. Generation Accuracy 지표는 인간 판단과 96.4%의 일치율을 보인 반면, Search-R1에서 사용하는 Exact Match는 15.8%에 불과했다. 이는 보상 선택이 검색 품질에 직접적인 영향을 미친다는 것을 보여준다.

아직 남은 과제들: 계산 비용과 편향성 문제 해결이 관건

s3의 혁신적 성과에도 불구하고 몇 가지 한계점이 존재한다. 가장 큰 과제는 훈련 중 보상 계산을 위해 LLM 추론이 필요하다는 점이다. 이는 토큰 수준이나 검색 전용 목적 함수에 비해 계산 오버헤드를 증가시켜 확장성을 제한한다. 연구팀에 따르면 s3는 단계별로는 더 느리지만, 최소한의 단계로 높은 성능을 달성하여 전체적으로는 효율적이라고 설명한다.

또한 s3는 유능한 고정 생성 모델의 존재를 전제로 한다. 저성능이나 명령어 처리가 약한 생성기의 경우 향상된 검색의 이점이 완전히 출력으로 이어지지 않을 수 있다. 편향성 문제도 주의해야 할 부분이다. s3는 모든 검색증강 시스템과 마찬가지로 검색기와 생성기 양쪽의 편향을 상속받는다. 특히 의료나 과학 분야와 같은 민감한 도메인에 적용할 때는 소스 말뭉치와 출력 결과에 대한 신중한 검증이 필요하다고 연구팀은 강조했다.

그럼에도 불구하고 s3는 저자원 커뮤니티에서도 효과적인 검색 에이전트 훈련을 가능케 하여 RAG 시스템을 더욱 접근 가능하게 만든다는 긍정적 영향이 기대된다.

FAQ

Q1: s3가 기존 RAG 시스템보다 효율적인 이유는 무엇인가요?

A1: s3는 검색 기능과 생성 기능을 분리하여 검색 에이전트만을 집중적으로 훈련시키기 때문입니다. 이를 통해 2,400개의 적은 데이터로도 170,000개 데이터로 훈련된 기존 모델보다 우수한 성능을 달성할 수 있습니다.

Q2: s3의 'Gain Beyond RAG' 보상 신호는 어떻게 작동하나요?

A2: GBR은 s3가 검색한 문서를 사용했을 때의 생성 정확도와 기본 상위 k개 검색 결과를 사용했을 때의 정확도 차이를 측정합니다. 이를 통해 실제로 답변 품질을 향상시키는 검색 성능을 직접적으로 최적화할 수 있습니다.

Q3: s3는 의료 분야와 같은 전문 도메인에서도 사용할 수 있나요?

A3: 네, 가능합니다. s3는 일반 데이터로만 훈련되었음에도 불구하고 의료 질의응답에서 76.6%의 높은 정확도를 기록했습니다. 이는 강화학습으로 습득한 검색 기술이 다양한 전문 도메인으로 안정적으로 전이된다는 것을 보여줍니다.



해당 기사에 인용된 논문 원문은 링크에서 확인 가능하다.