Fine-tuning MLIP foundation models: strategies for accuracy and transferability

Diese Arbeit evaluiert sieben Strategien zur Feinabstimmung von auf maschinellem Lernen basierenden interatomaren Potenzial-Fundamentmodellen (MLIP) über diverse chemische Benchmarks hinweg und zeigt auf, dass, während Voraussetzungen wie die Qualität des Fundamentmodells und die korrekte Energieinitialisierung von entscheidender Bedeutung sind, eine naive Feinabstimmung für die Genauigkeit bei Einzelsystemen optimal ist, wohingegen Multihead-Replay die Robustheit gegenüber Out-of-Distribution-Daten für einen breiteren Einsatz einzigartig bewahrt.

Das Wesentliche

Stellen Sie sich vor, Sie hätten einen Meisterkoch, der Jahre damit verbracht hat, perfekte Mahlzeiten zu kochen, indem er ausschließlich anorganische Zutaten wie Gesteine, Metalle und Salze verwendet. Dieser Koch ist ein „Foundation Model“. Nun möchten Sie diesen Koch ein ganz bestimmtes neues Gericht kochen lassen, wie zum Beispiel eine delikate organische Suppe oder einen biologischen Eintopf, wobei Sie nur eine sehr kleine Menge neuer Rezepte zur Verfügung haben.

Die große Frage lautet: Wie bringt man diesem Koch das neue Gericht bei, ohne dass er vergisst, wie man die alten zubereitet, oder ohne seine bestehenden Fähigkeiten zu ruinieren?

Dieser Artikel ist ein massives Küchenexperiment, bei dem sieben verschiedene Wege getestet werden, diesen Meisterkoch zu „feintunen“ (nachzutrainieren). Die Forscher fanden heraus, dass die Methode des Lehrens weniger entscheidend ist als drei kritische „Vorkoch-Schritte“: die Wahl des richtigen Kochs, das Festlegen der richtigen Basis und das Einstellen der Hitze.

Hier ist die Aufschlüsselung ihrer Ergebnisse in einfachen Worten:

1. Die drei „Pre-Flight“-Checks (Der wichtigste Teil)

Bevor Sie überhaupt mit dem Erlernen des neuen Rezepts beginnen, müssen Sie drei Dinge richtig machen. Wenn Sie hier Fehler machen, wird keine Lehrmethode Sie retten.

• Wählen Sie den richtigen Koch (Qualität des Foundation Models):

  • Die Analogie: Man würde keinen Koch einstellen, der nur Wasser kochen kann, um ihm beizubringen, wie man ein Soufflé backt.
  • Das Ergebnis: Die Qualität des ursprünglichen Modells ist wichtiger als die Strategie des Feintunings. Ein Modell, das auf einem riesigen, vielfältigen Datensatz anorganischer Materialien trainiert wurde (wie das „OMat24“-Modell), ist viel besser darin, neue, seltsame Chemie zu lernen, als ein älteres, kleineres Modell. Selbst wenn Sie dieselbe Lehrmethode verwenden, wird ein „besseres“ Foundation Model immer ein besseres Endergebnis liefern.

• Setzen Sie den Nullpunkt (Atomare Referenzenergie / $E_0$):

  • Die Analogie: Stellen Sie sich vor, Sie messen die Höhe eines Gebäudes. Wenn Sie nicht vom Erdgeschoss, sondern aus dem Keller heraus messen, werden Ihre Zahlen falsch sein, und das Gebäude könnte so aussehen, als würde es schweben oder im Boden versunken sein. In der Chemie müssen Sie das „Gewicht“ der einzelnen Atome abziehen, damit das Modell nur lernt, wie sie miteinander interagieren.
  • Das Ergebnis: Die Forscher fanden heraus, dass die Verwendung einer intelligenten, „modellbewussten“ Methode zur Festlegung dieses Nullpunkts entscheidend ist. Wenn Sie eine faule, durchschnittliche Schätzung verwenden, wird das Modell instabil. Es mag auf dem Papier gut aussehen (niedrige Fehlerraten), aber es wird in sich zusammenbrechen, wenn Sie versuchen, reale physikalische Prozesse zu simulieren (wie etwa ein Gebäude in einem Windkanaltest).

• Drehen Sie die Hitze herunter (Hyperparameter):

  • Die Analogie: Beim Erlernen einer neuen Fähigkeit wollen Sie nicht so schnell vorgehen, dass Sie stolpern, aber Sie wollen auch nicht so langsam vorgehen, dass Sie nie fertig werden.
  • Das Ergebnis: Verschiedene Lehrmethoden benötigen unterschiedliche „Lernraten“. Zum Beispiel kann eine Methode namens LoRA (die nur einen winzigen Teil des Modells verändert) eine sehr schnelle Lernrate vertragen, während eine Methode, die zwei Dinge gleichzeitig lehrt, ein sehr langsames, sanftes Tempo benötigt.

2. Die sieben Lehrstrategien

Sobald die drei Checks oben bestanden sind, testeten die Forscher sieben Wege, das neue Rezept zu lehren:

1. Naives Fine-Tuning: „Einfach weiterkochen.“ Sie nehmen den gesamten Koch und trainieren ihn weiter an den neuen Daten.
   • Ergebnis: Hervorragend geeignet, um ein spezifisches Gericht perfekt zu lernen. Aber wenn Sie diesen Koch später für eine andere Art von Essen nutzen wollen, könnte er seine alten Fähigkeiten vergessen haben (ein Problem, das als „katastrophales Vergessen“ bezeichnet wird).
2. Layer Freezing: „Die Grundlagen nicht anfassen.“ Sie sperren das Wissen des Kochs über grundlegende Messerschnitte und lassen ihn nur die neue Sauce lernen.
   • Ergebnis: Gut, aber manchmal zu starr. Es schränkt ein, wie gut der Koch sich an die neuen Zutaten anpassen kann.
3. LoRA (Low-Rank Adaptation): „Ein Spickzettel.“ Anstatt das ganze Kochbuch umzuschreiben, fügen Sie dem Koch eine kleine, effiziente Notiz am Revers hinzu, die nur die neuen Regeln abdeckt.
   • Ergebnis: Sehr effizient und genau für spezifische Aufgaben, ähnlich wie Naives Tuning.
4. Multihead Replay: „Der Dual-Head-Koch.“ Sie geben dem Koch zwei Hüte. Einen Hut für das neue Gericht und den anderen für die alten, vertrauten Gerichte. Er übt beide gleichzeitig.
   • Ergebnis: Dies ist der Gewinner in Sachen Sicherheit. Es ist die einzige Methode, die konsistent verhindert, dass der Koch seine alten Fähigkeiten vergisst. Sie hält den Koch sowohl für das neue Gericht als auch für die alten Gerichte kompetent.
5. Pseudolabel Replay: „Der synthetische Koch.“ Anstatt echte alte Rezepte zu verwenden, nutzt der Koch seine eigenen Vorhersagen alter Rezepte, um zu üben.
   • Ergebnis: Funktioniert gut und ist flexibel, da Sie nicht die ursprünglichen alten Daten benötigen, sondern nur das Gedächtnis des Kochs.
6. Replay + LoRA: Die Kombination aus dem Spickzettel und den zwei Hüten.
   • Ergebnis: Gut, aber der „Dual-Head“ allein war oft schon ausreichend.

3. Die wichtigsten Erkenntnisse

• Erfinden Sie das Rad nicht neu: Wenn Sie ein Modell für eine spezifische, eng gefasste Aufgabe benötigen (wie die bloße Simulation von Salzwasser), ist Naives Fine-Tuning der schnellste und einfachste Weg, um ein großartiges Ergebnis zu erzielen.
• Vergessen Sie die Vergangenheit nicht: Wenn Sie ein Modell benötigen, das auch seltsame, neue Situationen bewältigen kann (wie eine neue Art von Batterie oder ein komplexes biologisches Molekül), ohne seine ursprüngliche Ausbildung zu vergessen, müssen Sie Multihead Replay verwenden. Es ist die einzige Strategie, die das Modell robust und sicher gegen das „Vergessen“ gemacht hat.
• Qualität vor Tricks: Der Artikel betont, dass es wichtiger ist, Zeit in die Auswahl eines hochwertigen Foundation Models und die korrekte Einstellung der Energiereferenzen zu investieren, als die perfekte Fine-Tuning-Algorithmus zu wählen. Wenn das Fundament schwach ist oder die Mathematik falsch aufgesetzt wurde, wird die beste Lehrstrategie der Welt nichts ausrichten können.

Kurz gesagt: Um das beste KI-Modell für die Chemie zu erhalten, beginnen Sie mit einem klugen Fundament, legen Sie Ihre mathematischen Regeln korrekt fest, und wenn Sie möchten, dass die KI vielseitig und nicht vergesslich ist, bringen Sie ihr das Essen mit der „Dual-Head“-Methode (Multihead Replay) bei.

Technische Zusammenfassung

Technische Zusammenfassung: Fine-Tuning von MLIP-Foundation-Modellen

Problemstellung
Maschinengestützte interatomare Potentialmodelle (MLIP) mit Foundation-Modell-Charakter haben die Fähigkeit demonstriert, über diverse chemische Systeme hinweg zu transferieren, was einen Workflow ermöglicht, der den ressourcenintensiven Prozess des Trainings aufgabenspezifischer Potentiale von Grund auf vermeidet. Die Fachwelt mangelt es jedoch an systematischer Anleitung dazu, wie und wann diese Modelle feinabzustimmen (fine-tuning). Frühe Berichte deuteten darauf hin, dass naives Fine-Tuning oft zu „katastrophalem Vergessen“ führt, was die Einführung von einschränkenden Techniken (z. B. Layer Freezing, Low-Rank Adaptation) zur Folge hatte, die ursprünglich für Large Language Models entwickelt wurden. Diese Arbeit untersucht, ob diese Einschränkungen notwendig sind oder ob frühe Misserfolge auf andere Faktoren zurückzuführen waren, wie etwa schwächere Foundation-Modelle, eine ungeeignete Initialisierung der atomaren Referenzenergie ($E_0$) oder instabile Trainingsverfahren. Das Ziel der Studie ist es, die wesentlichen Faktoren zu charakterisieren, welche die Ergebnisse des Fine-Tunings prägen, insbesondere die Genauigkeit für die Zielaufgabe und die Out-of-Distribution (OOD)-Robustheit.

Methodik
Die Autoren evaluieren sieben verschiedene Fine-Tuning-Strategien über fünf chemisch diverse Benchmarks, drei Generationen von Foundation-Modellen und Trainingsdatensätze, die sich über fünf Größenordnungen in der Größe erstrecken.

• Evaluierte Fine-Tuning-Strategien:

  1. Nativ (Naive): Vollständige Parameteraktualisierungen via fortgesetztem Gradientenabstieg.
  2. Layer Freezing (Varianten): Einfrieren von Embedding- oder Message-Passing-Layern beim Training der Readouts; oder das Einfrieren des Embeddings und des ersten Message-Passing-Layers.
  3. Low-Rank Adaptation (LoRA): Injektion trainierbarer Low-Rank-Zerlegungen in sowohl skalare als auch äquivariante lineare Schichten bei gleichzeitigem Einfrieren der vortrainierten Gewichte.
  4. Multihead Replay: Gleichzeitige Optimierung auf Ziel- und Replay-Daten (aus dem Pretraining oder durch Pseudolabels) unter Verwendung separater Readout-Heads.
  5. Pseudolabel Replay: Eine Variante des Multihead Replay, bei der die Replay-Labels vom Foundation-Modell selbst generiert werden, wodurch die Quelle der Replay-Daten vom ursprünglichen Pretraining-Korpus entkoppelt wird.
  6. Replay + LoRA: Kombination von Multihead Replay mit LoRA.

• Benchmarks: Die Studie umfasst Systeme mit zunehmender Abweichung von der OMat24-Pretraining-Domäne (periodische anorganische Bulk-Strukturen):

  • Lithium-Argyrodit-Elektrolyte (anorganischer periodischer Feststoff).
  • Aqueous NaCl (ionische Lösung).
  • Eis-Polymorphe (molekularer Feststoff).
  • SN2-Reaktionen (gasphasige reaktive Chemie).
  • SPICE-Biomoleküle (organische/biomolekulare Konformere).

• Technische Implementierungen: Die Autoren haben drei neue Funktionen im MACE-Codebase implementiert:

  • LoRA, angepasst für äquivariante Message-Passing-Architekturen (Abdeckung sowohl skalarer als auch äquivarianter linearer Schichten).
  • Pseudolabelled Replay zur Entkopplung der Replay-Datenquellen.
  • Modell-bewusste Neuschätzung der atomaren Referenzenergie ($E_0$), um die vortrainierten Baselines mit den Ziel-Daten abzugleichen.

• Evaluationsmetriken: Über Standardfehler für Energie und Kräfte (pointwise) hinaus untersucht die Studie dynamische und extrapolative Verhaltensweisen, einschließlich radialer Verteilungsfunktionen (RDFs) aus Molekulardynamik-Simulationen (MD), Nudged Elastic Band (NEB) Reaktionsprofilen, MD-Stabilitätstests und der Random Structure Search (RSS), um Versagen bei der Nahbereichsrepulsion zu detektieren.

Kernergebnisse

1. Voraussetzungen dominieren die Strategiewahl: Die Studie stellt fest, dass die Qualität des Foundation-Modells, eine korrekte $E_0$-Initialisierung und gut gewählte Hyperparameter Voraussetzungen sind, deren Einfluss die spezifische Fine-Tuning-Strategie routinemäßig übersteigt.

   • Foundation-Qualität: Neuere Foundation-Modelle (z. B. auf OMat24 basierend) übertreffen ältere Modelle (MPTraj-basiert) konsistent in der OOD-Transferleistung, selbst bei fixen Fine-Tuning-Rezepten.
   • $E_0$-Initialisierung: Die Verwendung von „gemittelten“ $E_0$-Werten führt zu signifikant höheren Fehlern und MD-Instabilitäten (z. B. Versagen von Eis-Modellen innerhalb von 50 ps). „Neugeschätzte“ $E_0$-Werte (die den Nullpunkt des vortrainierten Modells an die Ziel-Daten anpassen) sind entscheidend für Stabilität und Transferierbarkeit und liefern oft bessere Ergebnisse als die Wahl des eigentlichen Fine-Tuning-Algorithmus.
   • Hyperparameter: Natives Fine-Tuning erfordert reduzierte Lernraten und eine erhöhte EMA-Decay. LoRA toleriert höhere Lernraten. Multihead Replay erfordert wesentlich niedrigere Lernraten, um konkurrierende Update-Signale zu vermeiden. Die Gewichtsbegrenzung (Weight Decay) sollte auf Null gesetzt werden, um zu verhindern, dass Parameter von der vortrainierten Lösung weggezogen werden.

2. Performance nach Zielsetzung:

   • In-Distribution Spezialisierung (Einzelsystem): Für enge Aufgaben (z. B. SN2-Barrieren, wässriges NaCl-Solvatisierung) erreichen die meisten Strategien (Nativ, LoRA, Multihead) eine starke Genauigkeit und übertreffen konsistent Modelle, die von Grund auf trainiert wurden. Natives Fine-Tuning bietet die beste Konvergenz für Single-System-Anwendungen.
   • Out-of-Distribution Robustheit: Bei der Evaluierung des Transfers auf verwandte, aber ungesehene Zusammensetzungen (z. B. Nicht-Argyrodit-Elektrolyte) oder andere Chemikalien (z. B. Biomoleküle) ist Multihead Replay (mit Original- oder Pseudolabel-Daten) der einzige Ansatz, der konsistent die OOD-Robustheit bewahrt. Er erhält die Genauigkeit auf der Pretraining-Verteilung, während er die Zielaufgabe lernt, und verhindert effektiv das katastrophale Vergessen.
   • Freezing und LoRA: Obwohl effektiv für die Parametereffizienz, zeigten Layer Freezing und LoRA Limitationen bei der Anpassung an Solvatisierungsmerkmale oder der Aufrechterhaltung einer breiten chemischen Robustheit im Vergleich zu Multihead Replay in den getesteten Szenarien.

Bedeutung und Behauptungen
Das Paper behauptet, dass die wahrgenommene Fragilität des nativen Fine-Tunings in MLIPs weitgehend ein Resultat suboptimaler Setups ist und nicht eine intrinsische Limitation der Methode. Die Autoren argumentieren:

• Natives Fine-Tuning ist eine lebensfähige und oft überlegene Ausgangsbasis für Single-System-Anwendungen, vorausgesetzt, das Foundation-Modell ist hochwertig und die $E_0$-Werte sind korrekt neu geschätzt.
• Multihead Replay ist die notwendige Strategie für einen breiteren Einsatz, wenn die Bewahrung des Verhaltens des Foundation-Modells außerhalb der Fine-Tuning-Verteilung erforderlich ist.
• Pseudolabelled Replay bietet einen praktischen Vorteil, da es die Verwendung jeglicher strukturell diverser Datensätze für das Replay ermöglicht und somit die Abhängigkeit vom Zugriff auf den ursprünglichen Pretraining-Korpus aufhebt.

Die Arbeit etabliert, dass es für Praktiker wichtiger ist, in das stärkste verfügbare Foundation-Modell zu investieren und eine korrekte Ausrichtung der atomaren Referenzenergie sicherzustellen, als einen spezifischen eingeschränkten Fine-Tuning-Algorithmus auszuwählen. Die Studie liefert einen systematischen Rahmen für den Einsatz von MLIP-Foundation-Modellen und macht das Fine-Tuning von einer Nischenoption zum Standard-Ausgangspunkt für die systemspezifische Entwicklung.