Building Reward Functions for Chemical AI: A Tale of Reward Hacking

Die Entwicklung robuster und korrekter Belohnungsfunktionen für KI-Modelle im Bereich der Chemie stellt eine erhebliche Herausforderung dar. In diesem Artikel berichten die Forscher von FutureHouse über ihre monatelangen Bemühungen, zwei spezifische Belohnungsfunktionen für ihr chemisches Denkmodell ether0 zu entwerfen. Die erste Funktion befasst sich mit der Retrosynthese eines Zielmoleküls, während die zweite darauf abzielt, ein Molekül mit einer bestimmten Anzahl von Atomen zu generieren. Diese Funktionen sind entscheidend für die Leistungsfähigkeit des Modells und erfordern ein tiefes Verständnis der chemischen Domäne.

Einführung in die Belohnungsfunktionen

Die Herausforderung bei der Gestaltung von Belohnungsfunktionen liegt darin, genau zu definieren, was das Modell erreichen soll. In der Chemie bedeutet dies, dass die Modelle nicht nur die richtigen Antworten liefern müssen, sondern auch in der Lage sein sollten, mit spezifischer chemischer Sprache zu kommunizieren. Reinforcement Learning ist der Prozess, durch den ein Denkmodell trainiert wird, um hohe Punktzahlen auf der Belohnungsfunktion zu erzielen. Dies ist sowohl faszinierend als auch riskant, da es alle Möglichkeiten aufzeigt, wie die Belohnungsfunktion falsch spezifiziert sein kann, und die Modelle Wege finden, um diese zu umgehen, was als Reward Hacking bezeichnet wird.

Retrosynthese: Ein zentraler Prozess in der Chemie

Die Retrosynthese ist eine der grundlegendsten Aufgaben in der Chemie. Sie besteht darin, ein Zielmolekül zu nehmen und vorzuschlagen, wie man es aus käuflichen Molekülen herstellen kann. Der Input ist ein Zielmolekül, und das Ergebnis ist eine Reaktion, deren Produkt das Zielmolekül ist. Die Herausforderung besteht darin, dass es viele Lösungen für diese Aufgabe gibt, was die Erstellung von Modellen mit traditionellen Ansätzen erschwert.

Um die Belohnungsfunktion für die Retrosynthese zu entwickeln, begannen die Forscher mit einem der leistungsstarken Vorwärts-Synthese-Modelle aus der Arbeit von Philippe Schwaller und setzten es als GPU-serverlose Funktion ein. Dies ermöglichte es, eine vorgeschlagene Reaktion von ether0 zu nehmen, das Ergebnis über eine Netzwerkabfrage vorherzusagen und das vorhergesagte Ergebnis mit dem Zielmolekül zu vergleichen. Diese Methode wurde als Oracle Verifier bezeichnet.

Herausforderungen und Lösungen

Die Forscher stießen auf mehrere Herausforderungen während des Prozesses. Zunächst war die Genauigkeit von ether0 zu Beginn bei 0%. Die ersten Nachrichten, die das Modell sendete, waren oft nicht korrekt formatiert. Zudem war der Oracle nicht immer genau, was bedeutete, dass das Modell plausible Reaktionen vorschlug, die jedoch nicht korrekt vorhergesagt wurden. Um diese Probleme zu beheben, wurde das Oracle-Modell mit den gesammelten Daten neu trainiert.

Ein weiteres Problem war die Herausforderung, sicherzustellen, dass die vorgeschlagenen Reaktionen nur aus käuflichen Verbindungen bestanden. Dies wurde durch die Implementierung von Bloom-Filtern erreicht, die eine effiziente Möglichkeit bieten, die Zugehörigkeit zu einer großen Menge von Verbindungen zu überprüfen, ohne eine Datenbank verwenden zu müssen.

Molekülgenerierung: Eine weitere Belohnungsfunktion

Die zweite Belohnungsfunktion, die ether0 lernen sollte, bestand darin, ein Molekül mit einer spezifischen Anzahl von Atomen zu generieren, wie beispielsweise drei Kohlenstoffen, fünf Wasserstoffen und zwei Sauerstoffen (C3H5O2). Diese Aufgabe stellte sich als herausfordernd heraus, da das Modell Schwierigkeiten hatte, die Anzahl der Atome korrekt zu verfolgen, was sich auf nachgelagerte Aufgaben wie die Vorhersage von Reaktionsprodukten auswirkte.

Um sicherzustellen, dass die generierten Moleküle sinnvoll waren, wurden Strukturwarnungen verwendet, die Muster enthalten, die in der Forschung entwickelt wurden, um problematische Verbindungen zu identifizieren. Diese Regeln halfen, Verbindungen auszuschließen, die in der Arzneimittelentdeckung problematisch sein könnten.

Reflexionen und Ausblick

Das Projekt dauerte etwa vier Monate, in denen viel Zeit mit Experimenten und der Anpassung der Belohnungsfunktionen verbracht wurde. Die Forscher lernten, dass die Spezifizierung guter Belohnungsfunktionen möglicherweise der größte Teil der Arbeit ist. Dies ist sowohl für Fachexperten von Vorteil als auch eine Herausforderung für Modelle, die versuchen, domänenunabhängige Ansätze zu verfolgen.

Die entwickelten Belohnungsfunktionen sind Open Source und könnten für die Ausbildung generativer Modelle in der Chemie von großem Wert sein. Die Implementierung von Kaufbarkeit und die Verwendung von Bloom-Filtern sind besonders nützlich für die Validierung von Molekülen.

Quellenliste:

• Quelle: Building Reward Functions for Chemical AI: A Tale of Reward Hacking
• Training a Scientific Reasoning Model for Chemistry
• Molecular transformer: a model for uncertainty-calibrated chemical reaction prediction
• Bloom filters for molecules
• Preparation of a neutral nitrogen allotrope hexanitrogen C2h-N6
• Learning chemical intuition from humans in the loop