Search-R1

제목 : 강화학습을 통해서 LLM이 reasoning을 하며 검색 엔진을 사용하도록 훈련하다.

들어가며

Pasted%image%20250407195609.png
Pasted image 20250407195638.png
Pasted image 20250407195822.png
최근 링크드인 등지에서 이러한 제목의 포스트들을 보았다. Beyond RAG라던지 LLM Search's DeepSeek moment? 와 같은 자극적인 제목은 RAG를 공부해왔던 나의 관심을 끌기에 충분했다.
이 논문은 정말 RAG를 제치는 혁신적인 기술일까? 아니면 알맹이 없는 어그로일까?
논문의 내용을 정리해보고, 내 생각을 덧붙여 보겠다.

Introduction

LLM이 외부 지식을 활용하는 것은 LLM의 능력치가 한껏 올라가고, 어떤 분야든지 적용이 가능해진다는 점에서 매우 유용하다. 그래서 대표적으로 RAG가 외부 지식을 LLM이 사용하기 위해 이용된다. 더불어서, 검색 엔진을 하나의 tool로서 활용하게 하는 경우도 많다.
그런데 이러한 기존 방식은 약간의 한계점들이 존재한다.
RAG의 경우에는

Retrieval의 정확도에 의존하고
Multi-hop retrieval이 가능한가에도 많이 의존한다.
이를 해결하기 위한 multi-query retrieval 등도 여전히 LLM이 검색 엔진을 활용하는 방법을 학습하지 못한다.
다른 방법인 tool로서 활용하는 경우에는,
프롬프트 방식의 접근은 일반화(generalization)이 어렵고
훈련을 한다면 크고 높은 품질의 훈련 데이터가 필요하다.

그래서, 기존에 OpenAI o1 혹은 DeepSeek R1과 같은 LLM이 강화 학습을 통해 reasoning을 학습한 것 처럼, RL을 이용해 검색 엔진 활용을 LLM에게 학습하려 한다.
이 때, 세 가지의 도전이 존재한다.

RL이 안정적으로 최적화 될 수 있는가? (어떤 RL 프레임워크를 쓸 것인가?)
Multi-Turn 상황에서 어떻게 동적으로 retrieval 전략을 수정하고 수행할 것인가?
리워드 모델을 어떻게 디자인 할 것인가?

이 논문에서는 reasoning 과정에서 외부 검색 엔진을 활용하는 방법을 RL을 통해 학습시킨다.

Search-R1

쉬운 설명

자, Search-R1은 간단하게 보려고 하면 굉장히 간단하다.
일단 훈련 템플릿을 아래와 같이 준비한다.

Pasted image 20250407150730.png
한국어로 해석해보면 다음과 같다.

주어진 질문에 대해 답하여라. 새로운 정보가 물어오면 반드시 reasoning을 수행하되, <think>와 </think>내에서 해라. Reasoning 이후, 특정 지식이 부족한 것 같으면 <search>와 </search> 사이에 검색어를 넣어 검색 엔진을 사용할 수 있다. 그러면 <information>와 </information> 사이에 검색 결과가 반환됩니다. 원하는만큼 검색을 할 수 있습니다. 외부 지식이 더 이상 필요하지 않으면, 자세한 설명 없이 <answer>와 </answer> 사이에 대답을 직접 제공하세요. 예를 들어서, <answer> 베이징 </answer>와 같이 하면 됩니다. 질문: *질문 내용*

이제 LLM은 훈련 중, 그리고 훈련 이후 해당 템플릿을 토대로 글을 생성해 낼 것이다.
그리고 적힌 것처럼 search 토큰을 탐지하면 검색 엔진을 통해 검색 결과를 information 토큰 내로 return한다.

이제 RL을 통해서 더 정확한 답변을 내놓도록 훈련된다. 실제로 쓰이는 리워드 모델에서는 모델이 answer 토큰 내에 내놓은 대답이 실제 대답과 일치하는가를 판단하고, 정답이라면 리워드를 부여하는 형식이다.
RL의 경우에는 PPO 및 GRPO를 모두 사용할 수 있다. 쉬운 설명에서는 해당 프레임워크에 대한 설명을 생략한다.

Multi-turn Search는 어떻게 해결했을까? 실제로 LLM에게 템플릿에서 '원하는 만큼 검색하라'라고 검색을 많이 사용하도록 유도하였다. 하지만 실제로는 무한 loop에 빠질 수 있기 때문에 최대 search limit를 정하고 그 제한을 넘으면 LLM의 인퍼런스를 종료하였다.

요약하자면, 검색 엔진을 사용하는 것을 포함하여 LLM의 reasoning을 RL을 통해 훈련시켰다!

어려운 설명

Search-R1의 학습 과정은 이렇게 정리할 수 있다.

\max_{π_{θ}} E_{x \in D, y \in π_{θ} (\cdot | x; R)} [r_{ϕ} (x, y)] - β D_{K L} [π_{θ} (y | x; R) | | π_{r e f} (y | x; R)]

$π_{θ}$ - 훈련하려는 LLM (policy LLM)
$π_{r e f}$ - 레퍼런스 LLM (훈련되는 LLM이 초기 상태를 너무 벗어나지 않도록 가이드 해주는 역할)
$r_{ϕ}$ - 리워드
$R$ - Retrieval 시스템 (검색 엔진 등)
$D$ - 데이터셋

첫 항을 보면, 어떤 데이터셋에서 샘플링한 데이터와, 그 데이터를 훈련하려는 LLM에 넣었을 때 나온 output인 $y$ 를 통해 리워드를 구하고, 그 리워드의 기댓값이 바로 첫 항이 된다. 이것을 최대화 해야 하므로, 리워드를 최대화해야 한다는 것과 동일한 의미다.
두 번째 항은 같은 데이터에 대한 훈련하려는 LLM과 레퍼런스 LLM의 KL Divergence가 최소가 되도록 한다. 이는 훈련되는 LLM이 초기 상태를 너무 벗어나지 않도록 가이드 해주는 역할이다.

$R$ 의 경우에는, 위의 쉬운 설명에서 설명했다시피 LLM이 텍스트를 출력하는 과정에 수행하는 검색 과정을 포함하겠다는 것이다. 이것이 포함되었다는 것을 제외하면 위 수식에 특별한 것은 없다.

해당 목표를 위해 훈련할 때, loss masking을 수행한다. 기본적으로 PPO나 GRPO를 사용할 때에, LLM에서는 토큰 별로 loss가 수행된다고 한다. 그런데 LLM이 생성하지 않은 retrieved contents (검색 결과)에 대해 loss를 적용하면 의도되지 않게 학습이 수행될 수 있어 검색 결과에 대해서는 loss masking을 통해 loss를 계산하지 않는다.

PPO를 이용한 훈련

이 리뷰에서는 수식은 생략하겠다. 원 논문을 참고하시길..

Pasted image 20250407154927.png

주어진 쿼리에 대하여 Search Engine (Retrieval system 등도 될 수 있겠다)과 훈련하려는 LLM (Policy LLM)이 위해서 언급한 방식을 통해 답변을 생성한다.
이후, Value LLM을 통해 Value가 예측된다.
또한, Reward Model과 Reference LLM을 이용하여 어느 정도의 리워드가 주어질 것인지를 결정한다. 이 리워드와 예측한 Value를 통하여 Policy LLM의 해당 답변이 얼마나 이득적(advantage)이었는지 계산한다. 이 advantage를 이용하여 이득적인 대답을 최대한 많이 생성하려 할 것이고, 최종적으로 리워드 총합이 크도록 학습된다.

GRPO를 이용한 훈련

PPO에서는 Policy LLM 뿐 아니라 Value LLM도 훈련해 주어야 했다. 이런 Value LLM도 학습해야 하는 한계 (귀찮음?)를 해결하기 위하여 GRPO를 사용하기도 한다.

Pasted image 20250407160036.png

GRPO에서는 PPO와 달리 한 쿼리에 대해 여러 답변들을 생각한다. 이 답변 그룹들의 리워드를 모두 구하고, 이 리워드들을 통해서 예상되는 advantage를 구하게 된다.
특히, 학습 과정에 Reference LLM과 Policy LLM간의 KL Divergence를 직접 도입하여 초기 상태랑 너무 달라지는 것을 방지했다.
이렇게 여러 답변들을 통해 advantage를 계산해서 Value LLM이 없어도 안정적인 학습을 가능하게 했다고 한다.

리워드 모델

이 논문에서의 리워드 모델은 매우 심플한다.

r_{ϕ} (x, y) = E M (a_{p r e d}, a_{g o l d})

$a$ 는 답변이다. 즉, 답변이 실제 정답과 exact-match하는가. 즉 문자가 얼마나 겹치는가를 이용해 리워드를 부여한다. 이 논문에서는 factual reasoning 태스크를 수행했고 그 답이 단답형으로 떨어지기에 가능한 것으로 보인다.

템플릿과 LLM의 답변 생성 방식

템플릿은 위해서 설명했던 것과 같다.
답변 생성 방식을 한 번 더 정리하면 아래와 같다.

주어진 질문에 따라 답변을 생성한다.
만약 search 토큰이 감지되면,
1. search 토큰 내의 검색어를 추출한다.
2. 검색어를 통해 검색 엔진에 검색을 수행한다.
3. Top n개의 검색 결과를 information 토큰에 담아서 시퀀스에 추가한다.
4. 검색 횟수 카운트를 늘려준다.
만약 검색 횟수 카운트가 threshold를 넘어서면 종료한다.
만약 answer 토큰이 감지되면, 그 안의 내용을 답변으로 내놓고 종료한다.

실험 및 결과

데이터셋

총 7가지 데이터셋을 사용했다.
NQ (Natural Questions), TriviaQA, PopQA, 멀티-홉으로는 HotpotQA, 2WikiMultiHopQA, Musique, Bamboogle

비교군

그냥 인퍼런스
CoT reasoning (검색없이)
RAG
IRCoT
Search-o1
SFT
R1 (검색기능 없이 강화학습) (같은 데이터로 강화학습 수행)

Setup

상대적으로 작은 LLM들인 Qwen-2.5-3B, Qwen-2.5-7B, Llama-3.2-3B를 이용.
Retrieval로는 2018년 버전 위키피디아 전체 dump에서 검색을 수행함. 또한 E5 임베딩 모델을 retriever로 사용.
Top-n은 3으로 설정했다고 한다. (검색 후 사용하는 passage 수)

훈련에는 NQ와 HotpotQA의 훈련 데이터셋을 모두 합쳐서 사용했다고 한다.

평가 메트릭은 EM으로 보인다.

결과

Pasted image 20250407190822.png

전반적으로 Search-R1의 base와 instruct 모두 높은 성능을 보여주고 있다.
이외 논문에서 나온 분석점을 정리해보자면 다음과 같다.
GRPO의 성능이 대체로 PPO를 사용할 때보다 좋다.
Instruct 모델을 RL하는 것이 그냥 Base 모델부터 훈련하는 것보다 더 빨리 최적화가 되지만, 최종 성능은 유사했다.
학습을 시작하면 긴 대답 길이가, 리워드가 늘어나면서 급격히 짧아지고 이 과정에서 검색 엔진을 사용하는 법을 배운다. 이후에는 대답 길이가 500토큰 가량으로 안정화되며 리워드가 천천히 올라간다.
retrieved context에 loss 마스킹을 하여 훈련하는 것이 마스킹을 하지 않는 것보다 성능이 뛰어났다.

개인적인 생각

논문 전체는 흥미롭게 읽었으나, RAG의 DeepSeek Moment라던지 Beyond RAG와 같은 점들은 과장된 면이 있다고 생각한다. 더불어 머릿속에 생각나는 여러 의문점들도 있다.

실험에 사용한 데이터셋이 적절하였는가?
실제 적용 단계에서의 리워드 모델의 한계
실제 산업에서 적용이 가능할까?

1. 실험에 사용한 데이터셋이 적절하였는가?

실험에는 총 7가지 데이터셋이 사용되었는데, 이것들이 과연 적절한 데이터셋인가에 의문이 있다.
기본적으로 논문의 핵심 메시지는, reasoning 학습을 위한 RL에서 검색을 하는 능력도 같이 학습시키는 것에 있다. 즉, reasoning과 검색을 모두 요하는 데이터셋이 이상적인 데이터셋이라고 할 수 있을 것이다.

그러나 사용한 데이터셋들은 주로 단답형 위주의 Open-Domain Question Answering (ODQA)에 관련한 데이터셋들이다.

Pasted image 20250407202700.png
위 사진은 실험에서 사용한 데이터셋 중 일부인 TriviaQA의 예시 질문들이다. 물론 Multi-hop QA도 포함되어 있기에 어느 정도의 Reasoning은 포함한다고 할 수 있다. 하지만 우리가 흔히 떠올리는 복잡한 Reasoning과는 거리가 멀다.
실제로 DeepSeek-R1이 나오면서 Reasoning Task를 통해서 굉장히 어려운 문제나 복잡한 코딩 태스크도 곧잘 해낸다는 것이 주목을 받았다. 그래서 Humanity's Last Exam과 같은 어려운 Reasoning 데이터셋도 각광을 받았다.

그런데 이 실험에서 사용한 데이터셋들은 아무리 복잡해도 몇 가지 step의 multi-hop reasoning만 사용하면 그만이다. 이렇게 간단한 reasoning은 Query Decomposition (Visconde 참조) 방법 등으로 충분히 해결이 가능하다. LLM까지 안가더라도 훈련된 BERT 모델을 이용하면 Multi-hop Question을 Single-hop으로 잘 분해를 해 줄 것이다.

결국 LLM이 검색 엔진까지 동원하여 Reasoning과 함께 풀기에는 너무나 쉬운 문제들이 데이터셋으로 쓰였다는 것이다.
내 생각으로는 논문에서 제시한 방법으로 훨씬 복잡한 실생활의 문제나 학문적인 문제들의 풀이가 가능할 것으로 생각한다. 간단한 ODQA 데이터셋보다는 검색을 동반한 복잡한 Reasoning 및 검색을 요하는 데이터셋으로 실험을 진행했다면 훨씬 드라마틱한 성능 향상을 보였을 것으로 생각한다. 물론, 이러한 데이터셋이 구축되어 있지 않았을 확률이 높고, 추후 연구에서 더 보완해야 한다고 생각한다.

2. 실제 적용 단계에서의 리워드 모델의 한계

논문에서 제안하는 리워드 모델은 exact match만을 이용해 매우 simple하게 설계되었다.
물론 논문에서 사용한 데이터셋처럼 단답형 대답이 가능한 경우 매우 좋은 리워드 설계라고 생각한다.

그러나, 단답형 대답을 보기 위해서 RAG와 같은 시스템을 구성하려는 회사는 없을 것이다. 우리가 Perplexity와 같은 AI 검색 엔진 혹은 어시스턴트에게 기대하는 것은 해당 지식에 대한 자세한 설명과 배경 지식 설명, 맞춤형 안내 등이다. 이러한 경우에는 어떻게 리워드를 주어야 할까? 이러한 태스크에 있어서 검색을 활용하는 능력은 정말 중요한 것이 분명하다.
논문에서도 분명히 명시되어 있는데, 다른 태스크에 대한 탐색은 future work로 남겨둔다고 명시되어 있다.
결국 현실의 적용과는 거리가 떨어져 있는 리워드 모델링은 아래와 같은 문제를 유발한다.

3. 실제 산업에서 적용이 가능할까?

현실적이지 않은 리워드 모델은 실제 산업에서의 적용을 힘들게 할 것이다.
더군다나, LLM에 RL을 한 번 더 학습시키는 것은 큰 부담이 따르기 마련이다.

이에 더하여, 사실 LLM이 검색 엔진을 다루는 방식은 현재 당면한 문제가 아닐 확률이 높다.
현재 RAG를 적용하려는 많은 분들이 겪고 있는 대부분의 문제는 무엇일까? LLM의 성능?
아니다. 보통 Retrieval System의 정확도이다. 즉, Bottleneck은 주로 LLM보다는 Retrieval system 자체에 있다.

실제로 논문의 실험에서는 E5 임베딩 모델을 사용했다고 나온다. 상당한 양의 데이터셋 내에서 임베딩 모델만을 사용해 retrieve한다면 그것이 높은 정확도를 가지는 것은 매우 힘들다고 생각한다. 마치 부족한 실제 Retrieval system의 성능을 대변하는 것 같다.
이런 상황에서 Search-R1의 결과를 다시 보자. 평균 EM이 모두 0.3대인 것을 볼 수 있다. 즉, 10문제 중 3~4문제만을 맞추는 데에 성공했다는 것이다. 이러한 AI 시스템을 자신있게 도입할 수 있는 회사가 얼마나 될까?
결국 좋은 검색 엔진이 없다면 RL은 아무 의미 없는 것이다.

결론

나는 논문의 아이디어가 너무 마음에 든다. 사람이라면 어떠한 문제를 해결하면서 여기저기 검색을 해보고, 그 결과를 종합해서 대답을 도출하지 않겠는가? 그리고 이러한 구글링과 문제 해결력은 인간도 여러번 시행착오를 통해 학습해야 한다고 믿는다. 이를 LLM에게 학습시킨 것은 너무 좋은 아이디어고 좋은 접근이라는 생각이다.

그러나, 실험에서 사용한 데이터셋은 정말 Search-R1과 같은 시스템이 필요한 데이터셋이 아니라는 생각이 든다. 보다 양질의 데이터셋이 있다면, 검색까지 하는 아주 똑똑한 AI를 만들 수 있지 않을까?

더불어, 리워드 모델링의 어려움, RL을 위한 데이터셋 구축과 훈련 과정의 필요는 실제 적용에 장벽이 된다. 더더욱 장벽이 되는 것은 여전히 Retrieval System 자체의 성능에 의존할 수 밖에 없다는 것이다. 그리고 좋은 Retrieval System을 만드는 것은 아직도 어려운 일이다.

그렇다면 글의 시작에서 언급한 RAG를 제치는 혁신적인 기술인가? 에 대한 나의 대답은 그 정도는 아니다. 위에서 언급한 한계에 더불어서, 사실 RAG는 많은 부분이 Retrieval에 의존한다고 생각한다. 모든 LLM이 검색을 유용하게 활용하게 RL-tuning된다면 RAG는 사라질 수도 있겠지만, Retrieval은 여전히, 매우 중요할 것이다.