Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

Dit paper introduceert Reference-guided Policy Optimization (RePO), een nieuwe methode die de zwaktes van bestaande technieken oplost door het combineren van verifieerbare beloningen voor exploratie en referentiesturing voor exploitatie, waardoor Large Language Models effectiever kunnen worden getraind voor moleculaire optimalisatie zonder stap-voor-stap trainingsdata.

De kern

Stel je voor dat je een chemisch meesterwerk moet maken. Je hebt een bestaand recept (een molecule) en je wilt het een beetje aanpassen om het "beter" te maken. Misschien moet het effectiever zijn tegen een ziekte, of makkelijker oplosbaar in water. Maar er is een belangrijke regel: je mag het recept niet zomaar volledig herschrijven. Het moet nog steeds op het origineel lijken, anders werkt het niet meer als medicijn.

Dit is de uitdaging van moleculaire optimalisatie. En dit is waar de nieuwe methode RePO (Reference-guided Policy Optimization) om de hoek komt kijken, zoals beschreven in dit paper.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Gouden Eend" en de "Blinde Vlek"

Stel je voor dat je een chef-kok bent die een nieuw gerecht moet bedenken op basis van een bestaand recept. Je hebt een "voorbeeldrecept" (de referentie) dat perfect is, maar je krijgt alleen het eindresultaat te zien. Je krijgt geen stap-voor-stap uitleg van de chef over hoe hij het heeft gemaakt.

• De oude methode (SFT): De chef probeert het eindresultaat na te bootsen. Hij leert het recept uit zijn hoofd, maar hij stopt met nadenken over waarom bepaalde ingrediënten erin zitten. Hij wordt een robot die alleen het antwoord geeft, zonder de creatieve redenering. Hij durft geen nieuwe combinaties te proberen.
• De andere oude methode (RLVR/GRPO): De chef krijgt een puntensysteem. Als het gerecht lekkerder is, krijgt hij een punt. Maar omdat de keuken zo groot is, is het bijna onmogelijk om per ongeluk een perfect gerecht te maken zonder hulp. De chef blijft daarom heel conservatief: hij verandert bijna niets, want hij durft niet te experimenteren uit angst dat hij faalt. Hij blijft hangen in de "veilige zone".

Het resultaat? Ofwel een robot die niet creatief is, ofwel een chef die te bang is om iets nieuws te proberen.

2. De Oplossing: RePO (De Slimme Mentor)

De auteurs van dit paper hebben een nieuwe manier bedacht, RePO, die de beste van beide werelden combineert. Ze gebruiken een slimme mentor die twee dingen doet tegelijk:

Deel A: De "Vrije Ruimte" (Verkenning)

De chef mag eerst vrijuit experimenteren. Hij mag 10 verschillende variaties van het gerecht bedenken. Hij krijgt een score voor elke variatie:

• Is het lekkerder? (Beter voor het medicijn)
• Lijkt het nog op het origineel? (Niet te veel veranderd)
  Als een variatie goed scoort, krijgt de chef een beloning. Dit moedigt hem aan om nieuwe paden te verkennen in plaats van alleen maar veilig te spelen.

Deel B: De "Anker" (Referentie)

Hier komt het slimme deel. Terwijl de chef experimenteert, houdt de mentor het voorbeeldrecept (de referentie) in de gaten.

• De mentor zegt niet: "Doe precies wat ik deed."
• De mentor zegt wel: "Zorg dat je eindresultaat vergelijkbaar is met dit voorbeeld, maar gebruik je eigen redenering om daar te komen."

Dit is als een anker. Het houdt de chef vast aan de goede kant (zodat hij niet verdwaalt in de grote oceaan van mogelijke recepten), maar laat hem vrij om zijn eigen route te kiezen om daar te komen.

3. Waarom werkt dit zo goed?

In het paper zien ze drie grote voordelen:

1. Het voorkomt "leeg denken": De chef blijft niet steken in korte, saaie antwoorden. Hij leert weer om stap-voor-stap na te denken over de chemie (bijvoorbeeld: "Als ik deze groep vervang door een andere, wordt het molecuul stabieler").
2. Het voorkomt "angst voor verandering": Omdat de mentor het voorbeeldrecept gebruikt als leidraad, durft de chef grotere sprongen te maken dan alleen met het puntensysteem. Hij weet dat hij ergens naartoe moet, maar hij mag zelf de weg vinden.
3. Het werkt zelfs als de instructies raar zijn: Zelfs als iemand vraagt om een heel vreemd soort aanpassing (een "nieuwe instructiestijl"), kan de chef dit aan. Hij heeft geleerd hoe hij moet redeneren, niet alleen wat hij moet zeggen.

De Gouden Kooi-analogie

Stel je voor dat je een vogel (het AI-model) in een enorme kooi (de chemische ruimte) hebt.

• Oude methode 1: Je geeft de vogel een foto van een andere vogel en zegt: "Wees precies zo." De vogel stopt met vliegen en zit alleen maar te staren.
• Oude methode 2: Je zegt: "Vlieg maar rond, als je een mooie plek vindt, krijg je een zadenbeloning." De vogel durft niet weg te vliegen omdat de kooi te groot is en hij geen zaden ziet.
• RePO: Je hangt een treklijn naar een bekende, veilige plek (het voorbeeld), maar je laat de vogel zelf beslissen hoe hij daarheen vliegt. Hij mag door de lucht zwieren, maar de lijn zorgt dat hij niet verdwaalt. Zo leert hij vliegen én vindt hij de beste plekken.

Conclusie

Dit paper introduceert RePO, een methode die AI-modellen helpt om betere medicijnen te ontwerpen. Het combineert de kracht van vrij experimenteren (zodat ze nieuwe ideeën vinden) met slimme begeleiding (zodat ze niet de weg kwijtraken).

Het resultaat? AI's die niet alleen antwoorden geven, maar ook verstandig redeneren over chemie, en dat doen met een balans tussen creativiteit en veiligheid. Het is alsof we een chemische chef-kok hebben getraind die zowel een kunstenaar als een wetenschapper is.

Technische samenvatting

Probleemstelling

Het paper adresseert een fundamenteel probleem bij het toepassen van Large Language Models (LLMs) op instructiegebaseerde moleculaire optimalisatie. In deze taak moet een model een inputmolecuul ($m_0$) aanpassen om een specifiek chemisch eigenschap (bijv. QED, LogP) te verbeteren, terwijl de structurele gelijkenis met het oorspronkelijke molecuul behouden blijft.

De auteurs identificeren een supervisie-mismatch in bestaande methoden:

1. Supervised Fine-Tuning (SFT): Bestaande datasets bevatten vaak slechts één geoptimaliseerd referentiemolecuul zonder stap-voor-stap redeneertraject. SFT op deze "antwoord-only" data zorgt ervoor dat het model de noodzakelijke multi-stap redenering verliest (collapse) en direct naar een antwoord springt. Dit beperkt de exploratie van de chemische ruimte.
2. Reinforcement Learning with Verifiable Rewards (RLVR): Methoden zoals GRPO (Group Relative Policy Optimization) vertrouwen op beloningen. Echter, in de chemische ruimte zijn succesvolle moleculen (die zowel de eigenschap verbeteren als aan de gelijkenisvoorwaarde voldoen) zeldzaam. Dit leidt tot spaarzame feedback (sparse rewards), waardoor het model conservatief blijft en nauwelijks leert.

Bestaande aanpakken falen dus ofwel in het behoud van redeneercapaciteit (SFT) of in het effectief verkennen van de oplossingruimte (RLVR).

Methodologie: RePO

De auteurs stellen Reference-Guided Policy Optimization (RePO) voor, een hybride optimalisatieframework dat de voordelen van RLVR combineert met referentiegeleiding, zonder dat er gelabelde tussenstappen nodig zijn.

De kern van RePO bestaat uit een objectivelfunctie die drie componenten combineert:

1. RLVR Term (Exploratie):
   Het model genereert een groep van kandidaat-moleculen met bijbehorende redeneertrajecten. Een beloningsfunctie ($r$) beoordeelt deze op basis van:

   • Eigenschapverbetering: Is de doel-eigenschap verbeterd?
   • Structuurbehoud: Voldoet de Tanimoto-gelijkenis aan de drempelwaarde?
     De policy wordt bijgewerkt om de kans op hoge-beloning antwoorden te vergroten (via GRPO-stijl updates), wat exploratie van nieuwe moleculen stimuleert.

2. Referentiegeleiding Term (Exploitatie):
   Om de spaarzame beloning te mitigeren, wordt een supervisie-term toegevoegd. Deze term vergroot de waarschijnlijkheid van het referentiemolecuul ($m_{ref}$) uit de dataset, maar alleen geconditioneerd op het gegenereerde redeneertraject ($t_i$) van het model.

   • Cruciaal: De gradienten worden niet teruggevoerd naar de redeneer-token ($t_i$), maar alleen naar het antwoordtoken ($\hat{m}$).
   • Dit zorgt ervoor dat het model wordt "geankerd" aan een geldige oplossing (verminderend de leegte van de beloning), maar vrij blijft om diverse redeneerpaden te verkennen om daarheen te komen.

3. KL-Regularisatie:
   Een standaard KL-divergentie-term zorgt voor stabiliteit en voorkomt dat de policy te ver afwijkt van de oorspronkelijke distributie.

De totale objectief is:
$$J_{RePO} = \text{Exploratie (RLVR)} + \beta \cdot \text{Referentiegeleiding (Answer-level)} - \gamma \cdot \text{KL-Regularisatie}$$

Belangrijkste Bijdragen

1. Analyse van Supervisie-Mismatch: Het paper toont empirisch aan dat SFT leidt tot "answer-only" output (verlies van redenering) en dat pure RLVR (GRPO) vastloopt in conservatieve, marginale aanpassingen vanwege spaarzame beloningen.
2. Ontwikkeling van RePO: Een nieuw framework dat referentiegeleiding op antwoordniveau koppelt aan RL-gedreven exploratie. Dit lost het probleem op van het ontbreken van traject-data in wetenschappelijke datasets.
3. Uitgebreide Validatie: RePO wordt getest op twee benchmarks (TOMG-Bench en MuMOInstruct) voor zowel single-objective als multi-objective taken, en toont superioriteit in generalisatie naar onbekende instructiestijlen.

Resultaten

De experimenten tonen aan dat RePO consequent beter presteert dan bestaande baselines (Base Model, SFT, GRPO, en GRPO met SFT-initialisatie):

• Prestatieverbetering: Op TOMG-Bench behaalt RePO de beste resultaten op 4 van de 6 taken, gemeten aan de hand van de gecombineerde metric Success Rate × Similarity. Er wordt een verbetering van tot 17,4% in success rate gezien ten opzichte van GRPO.
• Balans: RePO slaagt erin om een betere balans te vinden tussen het verbeteren van de eigenschap en het behouden van de structuur, terwijl SFT vaak de structuur opoffert en GRPO te conservatief is.
• Generalisatie: RePO generaliseert beter naar onbekende instructiestijlen (unseen instructions) in multi-objective taken (MuMOInstruct).
• Redeneerkwaliteit: Kwalitatieve analyses en "LLM-as-a-judge" evaluaties tonen aan dat RePO chemisch valide redeneringen produceert (bijv. correcte substitutie van atomen op basis van sterische hindernis), terwijl GRPO vaak chemisch onjuiste redeneringen of ongeldige SMILES-strings genereert.
• Inference Scaling: RePO profiteert van meer inferentie-berekeningen (Best-of-k sampling), waarbij zowel het succespercentage als de gelijkenis verbetert naarmate het aantal trials toeneemt.

Betekenis en Impact

Dit werk is significant voor de toepassing van AI in de wetenschap, specifiek in de drug discovery en materiaalwetenschap.

• Het oplost een kritiek probleem: hoe LLMs effectief te trainen voor complexe, beperkte optimalisatietaken waar stap-voor-stap data schaars is.
• Het demonstreert dat het combineren van "demonstratie" (referentie) en "exploratie" (RL) leidt tot robuustere wetenschappelijke redeneerders dan het gebruik van één van deze methoden alleen.
• De methode is breed toepasbaar op andere wetenschappelijke domeinen waar de oplossing makkelijk te verifiëren is maar moeilijk te specificeren (bijv. retrosynthese of voorspelling van geneesmiddeleninteracties), zoals aangegeven in de toekomstperspectieven van het paper.

Samenvattend biedt RePO een nieuwe standaard voor het trainen van LLMs in domeinen met strikte constraints en schaarse supervisie, door slim gebruik te maken van beschikbare referentieoplossingen zonder de creatieve exploratie van het model te onderdrukken.