Search-R1: Treinando LLMs para Raciocinar e Utilizar Mecanismos de Busca com Aprendizado por Reforço

Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning

March 12, 2025

Autores: Bowen Jin, Hansi Zeng, Zhenrui Yue, Dong Wang, Hamed Zamani, Jiawei Han

cs.AI

Resumo

Adquirir conhecimento externo e informações atualizadas de forma eficiente é essencial para o raciocínio eficaz e a geração de texto em modelos de linguagem de grande escala (LLMs). Abordagens de aumento de recuperação e treinamento de uso de ferramentas, onde um mecanismo de busca é tratado como uma ferramenta, carecem de flexibilidade complexa de recuperação em múltiplas interações ou exigem grandes quantidades de dados supervisionados. Solicitar LLMs avançados com capacidades de raciocínio durante a inferência para usar mecanismos de busca não é ideal, pois o LLM não aprende a interagir de forma ótima com o mecanismo de busca. Este artigo apresenta o Search-R1, uma extensão do modelo DeepSeek-R1, onde o LLM aprende — exclusivamente através de aprendizado por reforço (RL) — a gerar (múltiplas) consultas de busca de forma autônoma durante o raciocínio passo a passo com recuperação em tempo real. O Search-R1 otimiza as execuções do LLM com interações de busca em múltiplas etapas, aproveitando o mascaramento de tokens recuperados para um treinamento RL estável e uma função de recompensa simples baseada em resultados. Experimentos em sete conjuntos de dados de perguntas e respostas mostram que o Search-R1 melhora o desempenho em 26% (Qwen2.5-7B), 21% (Qwen2.5-3B) e 10% (LLaMA3.2-3B) em relação às linhas de base SOTA. Este artigo ainda fornece insights empíricos sobre métodos de otimização RL, escolhas de LLMs e dinâmicas de comprimento de resposta no raciocínio aumentado por recuperação. O código e os checkpoints do modelo estão disponíveis em https://github.com/PeterGriffinJin/Search-R1.

English

Efficiently acquiring external knowledge and up-to-date information is essential for effective reasoning and text generation in large language models (LLMs). Retrieval augmentation and tool-use training approaches where a search engine is treated as a tool lack complex multi-turn retrieval flexibility or require large-scale supervised data. Prompting advanced LLMs with reasoning capabilities during inference to use search engines is not optimal, since the LLM does not learn how to optimally interact with the search engine. This paper introduces Search-R1, an extension of the DeepSeek-R1 model where the LLM learns -- solely through reinforcement learning (RL) -- to autonomously generate (multiple) search queries during step-by-step reasoning with real-time retrieval. Search-R1 optimizes LLM rollouts with multi-turn search interactions, leveraging retrieved token masking for stable RL training and a simple outcome-based reward function. Experiments on seven question-answering datasets show that Search-R1 improves performance by 26% (Qwen2.5-7B), 21% (Qwen2.5-3B), and 10% (LLaMA3.2-3B) over SOTA baselines. This paper further provides empirical insights into RL optimization methods, LLM choices, and response length dynamics in retrieval-augmented reasoning. The code and model checkpoints are available at https://github.com/PeterGriffinJin/Search-R1.