ProtRLSearch: A Multi-Round Multimodal Protein Search Agent with Large Language Models Trained via Reinforcement Learning

1. Problem Statement

Current Large Language Models (LLMs) applied to protein analysis face two critical limitations in healthcare and clinical research settings:

• Modality Gap: Existing search agents (e.g., BioMedSearch, MindSearch) operate primarily on text-only modalities. They fail to incorporate raw protein sequences as direct inputs during the search decision-making process. Since protein conclusions are strictly constrained by sequence-level characteristics (domains, mutation sites, amino acid structures), text-only agents often generate hallucinations or search paths that drift from critical sequence features.
• Search Process Constraints: Existing multi-round search agents (e.g., Search-R1) rely on Reinforcement Learning (RL) with rewards based solely on the final answer. This lack of intermediate supervision leads to:
  • Poor keyword selection in early rounds.
  • Inefficient tool selection.
  • The amplification of early errors in subsequent search rounds, as deviations are not corrected until the final output is evaluated.

2. Methodology: ProtRLSearch

The authors propose ProtRLSearch, a multi-round, multimodal search agent trained via RL with multi-dimensional reward signals.

A. Multimodal Architecture

The system integrates protein sequences and text into a unified transformer architecture:

• Protein Encoder: Uses a frozen protein foundation model (e.g., ESM-2) to extract residue-level contextual representations (domain boundaries, conserved motifs).
• Projection & Fusion: A learnable linear projection maps high-dimensional protein embeddings into the LLM's input space. These are concatenated with query text at the token level, allowing the LLM to attend to both sequence regions and textual descriptions simultaneously.
• Training Corpus: The model is pre-trained on the InterPT dataset to establish stable mappings between protein regions and their descriptions.

B. Multi-Round Search Agent

The agent consists of three modules operating in a loop:

1. Planner: Generates a Directed Acyclic Graph (DAG) of search plans. It extracts protein-centered keywords and assigns specific search tools (Web, Literature, UniProt) to each keyword.
2. Retriever: Executes parallel searches across multiple sources. It ranks results using a hybrid scoring system:
   • Vector Relevance: Semantic similarity via PubMedBERT.
   • LLM Relevance: Contextual relevance to the specific protein query.
3. Executor: Analyzes retrieved results to generate intermediate reasoning (<reason>), conclusions (<answer>), and a decision on whether to continue searching (<decide>). If insufficient, it generates a next_query for the next round.

C. Multi-Dimensional Reward Design

To address the limitations of "final-answer-only" RL, ProtRLSearch employs a composite reward function ($R_{total}$) with four distinct components:
$$R_{total} = \lambda_{Ans}R_{Ans} + \lambda_{KW}R_{KW} + \lambda_{Tool}R_{Tool} + \lambda_{Fmt}R_{Fmt}$$

• Answer Reward ($R_{Ans}$): Evaluates semantic similarity between the generated answer and the ground truth (GT).
• Keyword Reward ($R_{KW}$): Ensures the Planner extracts correct protein-related keywords (e.g., specific domains, variants) in the first round.
• Tool Reward ($R_{Tool}$): Verifies that the correct search tools (e.g., UniProt for sequence data) are assigned to specific keywords.
• Format Reward ($R_{Fmt}$): Enforces strict adherence to structured output tags (e.g., <DAG>, <search_results>) to ensure parseability.

D. Training Data

The model is trained on 3,000 multimodal samples constructed from verified UniProt protein-species pairs and PubMed literature. These samples cover six protein categories and four reasoning tasks (variant-to-phenotype, structure-to-function, etc.).

3. Key Contributions

1. ProtRLSearch Agent: A novel multi-round search agent that treats protein sequences as a parallel modality to text, enabling sequence-aware search planning and decision-making.
2. Multi-Dimensional RL: A novel reward scheme that constrains the search process at the keyword, tool, format, and answer levels, preventing early-stage drift and improving search path efficiency.
3. ProtMCQs Benchmark: A new benchmark dataset of 3,000 multiple-choice questions organized into three difficulty levels:
   • Level 1: Sequence-constrained functional/phenotypic reasoning.
   • Level 2: Reasoning with distractors based on sequence features and biological context.
   • Level 3: Complex integration of sequence features with signaling pathways and regulatory networks.

4. Experimental Results

The model was evaluated on BioMedMCQs (general biomedical) and ProtMCQs (protein-specific), comparing against baselines like BioMedSearch, Search-R1, and ProtLLM.

• Performance on ProtMCQs:

  • Level 1: Accuracy improved from 35.7% (Baseline) to 86.9%.
  • Level 2: Accuracy improved from 30.5% to 77.4%.
  • Level 3: Accuracy improved from 26.1% to 72.5%.
  • Efficiency: Despite multi-round search, inference time increased only marginally (e.g., ~15.9s vs. 8.9s for baseline) compared to other multi-round agents like Search-R1, which suffered from exponential time growth.

• Performance on BioMedMCQs:

  • Achieved 89.2% accuracy on Level 1, outperforming Search-R1 (76.1%) and BioMedSearch (73.3%).

• Generalization:

  • On general medical datasets (MedMCQA, MedQA), ProtRLSearch achieved 90.4% and 87.4% accuracy respectively, demonstrating strong transferability beyond protein-specific tasks.

• Ablation Studies:

  • Removing RL caused a significant drop in accuracy (e.g., -12.5% on ProtMCQs Level 1) and increased token usage.
  • Removing Protein Modality caused the largest drop on ProtMCQs (Level 3 accuracy dropped from 72.5% to 55.4%), confirming the necessity of sequence input for complex protein reasoning.

5. Significance

• Bridging the Modality Gap: This work demonstrates that integrating raw protein sequences directly into the LLM's search loop significantly outperforms text-only approaches, particularly for tasks requiring structural or mutational reasoning.
• Process-Oriented RL: By moving beyond final-answer rewards to include process constraints (keywords, tools), the model learns to "think" more like a domain expert, reducing hallucinations and search drift.
• Clinical Applicability: The ability to perform accurate, sequence-constrained reasoning is crucial for interpreting disease-related variants and supporting clinical decision-making, addressing a critical gap in current AI-driven healthcare tools.
• Benchmarking: The introduction of ProtMCQs provides a rigorous standard for evaluating future multimodal protein reasoning models.