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

Probleemstelling

Proteïne-analyse taken in de gezondheidszorg, zoals het interpreteren van ziekterelateerde varianten en het analyseren van proteïne-functies, vereisen nauwkeurig redeneren onder strikte beperkingen van proteïne-sequenties. Bestaande zoekagenten (zoals BioMedSearch en MindSearch) hebben twee belangrijke beperkingen:

1. Beperkte modaliteit: Ze opereren voornamelijk op tekst en behandelen proteïne-sequenties als gewone tekst. Hierdoor kunnen ze de structurele en functionele semantiek van sequenties (zoals domeinen en mutatieplaatsen) niet effectief integreren in het zoekproces.
2. Eén ronde en gebrek aan procescontrole: Ze voeren vaak slechts één zoekronde uit en vertrouwen op beloningen (rewards) die alleen gebaseerd zijn op het eindantwoord. Dit leidt tot een gebrek aan controle over het zoekproces zelf, waardoor fouten in de selectie van trefwoorden of de richting van het redeneren vroeg in het proces niet worden gecorrigeerd en zich kunnen versterken.

Methodologie: ProtRLSearch

De auteurs stellen ProtRLSearch voor, een multi-rondige zoekagent die is getraind met Reinforcement Learning (RL) en gebruikmaakt van multi-dimensionale beloningssignalen. Het systeem combineert proteïne-sequenties en tekst als multimodale invoer.

1. Multimodale Representatie:

• Het model gebruikt een bevroren proteïne-pretrainmodel (bijv. ESM-2) om contextuele representaties op residu-niveau te extraheren (domeingrenzen, behouden motieven).
• Deze embeddings worden via een leerbare lineaire projectie naar de inputruimte van de Large Language Model (LLM) gemapt.
• De proteïne-embeddings en de tekstuele query worden samengevoegd en gezamenlijk verwerkt binnen een enkele transformer-architectuur, waardoor het model zowel de sequentie-informatie als de tekstuele context in één aandachtssfeer kan analyseren.

2. Multi-Rondige Zoekarchitectuur:
De agent bestaat uit drie modules die samenwerken in een cyclus:

• Planner: Transformeert de query in een zoekplan. Het extrahet proteïne-gerelateerde trefwoorden en wijst zoektools toe (Web, Literatuur, UniProt). Het outputt een gerichte acyclische graaf (DAG) als gestructureerd zoekplan.
• Retriever: Voert parallelle zoekopdrachten uit op basis van de DAG. Het scant resultaten op relevantie met behulp van zowel vector-scores (PubMedBERT) als LLM-gebaseerde scores.
• Executor: Analyseert de geretrieveerde resultaten, trekt een tussentijdse conclusie en beslist of er verdere zoekrondes nodig zijn (<decide>). Als het antwoord onvoldoende is, genereert het een nieuwe query (<next_query>) voor de volgende ronde.

3. Multi-Dimensionale Beloningsfunctie (RL):
Om de agent te trainen zonder expliciete supervisie van elke tussenstap, wordt een totale beloning ($R_{total}$) berekend als een gewogen som van vier componenten:

• Answer Reward: Beoordeelt de semantische overeenkomst tussen het gegenereerde antwoord en het ground truth (GT) antwoord.
• Keywords Reward: Belooft het correct extraheren van proteïne-gerelateerde trefwoorden (namen, domeinen) in de eerste ronde.
• Tool Reward: Belooft de juiste toewijzing van zoektools aan de gegenereerde trefwoorden.
• Format Reward: Zorgt ervoor dat de output voldoet aan de vereiste gestructureerde tags (DAG, zoekresultaten, redenering, antwoord).

Belangrijkste Bijdragen

1. ProtRLSearch Agent: Een nieuwe agent die proteïne-sequenties en tekst als gelijke modaliteiten behandelt, wat leidt tot sequentie-bewuste zoekbeslissingen en verbeterd redeneren.
2. Multi-Dimensionaal RL-ontwerp: Een innovatief beloningssysteem dat niet alleen het eindantwoord, maar ook het zoekproces (trefwoorden, tools, formaat) optimaliseert.
3. ProtMCQs Benchmark: Een nieuw dataset van 3.000 meerkeuzevragen (MCQ's) in drie moeilijkheidsgraden. Deze dataset test specifiek het vermogen van modellen om sequentie-informatie te integreren met zoekresultaten voor redenering over functies, fenotypes en signaalwegen.

Resultaten

De prestaties van ProtRLSearch werden getest op de bestaande dataset BioMedMCQs en de nieuwe ProtMCQs.

• Verbetering op ProtMCQs: ProtRLSearch toonde aanzienlijke verbeteringen ten opzichte van baselines en eerdere zoekagenten:
  • Niveau 1: Van 35,7% (baseline) naar 86,9%.
  • Niveau 2: Van 30,5% naar 77,4%.
  • Niveau 3: Van 26,1% naar 72,5%.
• Efficiëntie: In tegenstelling tot andere multi-rondige methoden (zoals Search-R1) die langere inferentietijden hebben, convergeert ProtRLSearch sneller naar effectieve zoekpaden, wat resulteert in acceptabele inferentietijden (bijv. 15,9s voor Niveau 1) zonder exponentiële toename bij hogere moeilijkheidsgraden.
• Ablatiestudies: Experimenten toonden aan dat zowel de RL-component als de multimodale proteïne-modelling essentieel zijn. Het verwijderen van de proteïne-modality leidde tot een sterke daling in prestaties op ProtMCQs (van 72,5% naar 55,4% op Niveau 3), terwijl het verwijderen van RL de efficiëntie en stabiliteit van het zoekproces aantastte.
• Generalisatie: Het model presteerde ook sterk op algemene medische datasets (MedMCQA en MedQA), wat aantoont dat de methode robuust is en goed generaliseert.

Significantie

Dit paper is significant omdat het een brug slaat tussen de complexe wereld van proteïne-sequenties en de kracht van zoekagenten op basis van LLM's. Door proteïne-sequenties expliciet als een multimodale invoer te behandelen en het zoekproces te sturen via multi-dimensionale beloningen, lost ProtRLSearch het probleem van "hallucinaties" en sequentie-ongerechtigdheid op in bestaande systemen. Het biedt een robuust kader voor klinisch onderzoek en ziekte-analyse, waar nauwkeurige interpretatie van mutaties en proteïne-functies cruciaal is. De introductie van ProtMCQs stelt bovendien een nieuwe standaard voor het evalueren van multimodale redeneervermogens in de bio-informatica.