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

Questo articolo valuta sette strategie di fine-tuning per modelli fondativi di potenziali interatomici appresi tramite machine learning (MLIP) attraverso diversi benchmark chimici, rivelando che mentre prerequisiti come la qualità del modello fondativo e la corretta inizializzazione dell'energia sono fondamentali, il fine-tuning ingenuo è ottimale per l'accuratezza su sistemi singoli, mentre il multihead replay preserva in modo unico la robustezza fuori distribuzione per un impiego più ampio.

L'essenziale

Immagina di avere un maestro chef che ha passato anni imparando a cucinare pasti perfetti usando solo ingredienti inorganici come rocce, metalli e sali. Questo chef è un "Modello Fondazionale". Ora, vuoi che questo chef cucini un nuovo piatto specifico, come una delicata zuppa organica o uno stufato biologico, usando una quantità molto piccola di nuove ricette.

La grande domanda è: Come si insegna a questo chef il nuovo piatto senza fargli dimenticare i vecchi o senza rovinare le sue abilità esistenti?

Questo articolo è un enorme esperimento in cucina che testa sette modi diversi per "affinare" (fine-tuning) questo maestro chef. I ricercatori hanno scoperto che il metodo di insegnamento conta meno di tre passaggi critici di "pre-cottura": scegliere lo chef giusto, impostare la giusta linea di base e regolare il calore.

Ecco la suddivisione delle loro scoperte in termini semplici:

1. I tre controlli "Pre-Volo" (La parte più importante)

Prima ancora di iniziare a insegnare la nuova ricetta, devi fare tre cose correttamente. Se sbagli questi passaggi, nessun metodo di insegnamento potrà salvarti.

• Scegli lo Chef Giusto (Qualità del Modello Fondazionale):

  • L'Analogia: Non assumeresti uno chef che sa solo far bollire l'acqua per insegnarti come preparare un soufflé.
  • La Scoperta: La qualità del modello originale è più importante della strategia di fine-tuning. Un modello addestrato su un dataset enorme e diversificato di materiali inorganici (come il modello "OMat24") è molto più bravo a imparare nuova e strana chimica rispetto a un modello più vecchio e piccolo. Anche se usi lo stesso metodo di insegnamento, un modello fondazionale "migliore" produrrà sempre un piatto migliore.

• Imposta lo Zero (Energia di Riferimento Atomico / $E_0$):

  • L'Analogia: Immagina di misurare l'altezza di un edificio. Se inizi a misurare dal seminterrato invece che dal piano terra, i tuoi numeri saranno sbagliati e l'edificio potrebbe sembrare sospeso o interrato. In chimica, devi sottrarre il "peso" dei singoli atomi in modo che il modello impari solo come essi interagiscono.
  • La Scoperta: I ricercatori hanno scoperto che usare un modo intelligente e "consapevole del modello" per impostare questo punto zero è fondamentale. Se usi una stima media pigra, il modello diventa instabile. Potrebbe sembrare buono sulla carta (bassi punteggi di errore), ma crollerà quando proverai a simulare la fisica reale (come un edificio che crolla durante un test in galleria del vento).

• Abbassa il Calore (Iperparametri):

  • L'Analogia: Quando impari una nuova abilità, non vuoi muoverti così velocemente da inciampare, ma non vuoi nemmeno muoverti così lentamente da non finire mai.
  • La Scopia: Diversi metodi di insegnamento richiedono diversi "tassi di apprendimento" (learning rates). Ad esempio, un metodo chiamato LoRA (che cambia solo una piccola parte del modello) può gestire un tasso di apprendimento molto veloce, mentre un metodo che insegna due cose contemporaneamente ha bisogno di un ritmo molto lento e delicato.

2. Le Sette Strategie di Insegnamento

Una volta superati i tre controlli sopra citati, i ricercatori hanno testato sette modi per insegnare la nuova ricetta:

1. Naive Fine-Tuning: "Continua solo a cucinare." Prendi l'intero chef e continui ad addestrarlo sui nuovi dati.
   • Risultato: Ottimo per imparare un piatto specifico alla perfezione. Ma se provi a usare questo chef per un altro tipo di cibo in seguito, potrebbe aver dimenticato le sue vecchie abilità (un problema chiamato "oblio catastrofico").
2. Layer Freezing: "Non toccare le basi." Blocchi la conoscenza dello chef sulle abilità base del coltello e gli permetti solo di imparare la nuova salsa.
   • Risultato: Buono, ma a volte troppo rigido. Limita quanto bene lo chef può adattarsi ai nuovi ingredienti.
3. LoRA (Low-Rank Adaptation): "Aggiungi un foglietto illustrativo." Invece di riscrivere l'intero ricettario, aggiungi un piccolo e l'efficiente taccuino sul grembiule dello chef che copre solo le nuove regole.
   • Risultato: Molto efficiente e accurato per compiti specifici, simile al Naive tuning.
4. Multihead Replay: "Lo Chef a Due Cappelli." Dai allo chef due cappelli. Un cappello per il nuovo piatto e l'altro per i vecchi piatti familiari. Si esercita su entrambi contemporaneamente.
   • Risultato: Questo è il vincitore per la sicurezza. È l'unico metodo che previene costantemente l'oblio delle vecchie abilità dello chef. Mantiene lo chef bravo sia nel nuovo piatto che in quelli vecchi.
5. Pseudolabel Replay: "Lo Chef Sintetico." Invece di usare vecchie ricette reali, usi le predizioni dello chef sulle vecchie ricette per fare pratica.
   • Risultato: Funziona bene ed è flessibile perché non hai bisogno dei dati originali vecchi, basta la memoria dello chef.
6. Replay + LoRA: Combinare il foglietto illustrativo con i due cappelli.
   • Risultato: Buono, ma il solo "Dual Head" era spesso sufficiente.

3. Le Grandi Conclusioni

• Non reinventare la ruota: Se hai bisogno di un modello per un compito specifico e ristretto (come simulare solo l'acqua salata), il Naive Fine-Tuning è il modo più veloce e semplice per ottenere un ottimo risultato.
• Non dimenticare il passato: Se hai bisogno di un modello che possa gestire situazioni nuove e strane (come un nuovo tipo di batteria o una molecola biologica complessa) senza dimenticare il suo addestramento originale, devi usare il Multihead Replay. È l'unica strategia che ha mantenuto il modello robusto e sicuro dall' "oblio".
• Qualità sopra i Trucchi: L'articolo sottolinea che dedicare tempo alla scelta di un modello fondazionale di alta qualità e all'impostazione corretta dei riferimenti energetici è più importante che scegliere la strategia di fine-tuning perfetta. Se la base è debole o la matematica è impostata male, la migliore strategia di insegnamento del mondo non servirà a nulla.

In breve: Per ottenere la migliore IA per la chimica, parti da una base intelligente, imposta correttamente le tue regole matematiche e, se vuoi che l'IA sia versatile e non dimentichi, insegnale usando il metodo "Dual Head" (Multihead Replay).

Sommario tecnico

Sintesi Tecnica: Fine-tuning di Modelli Foundation per MLIP

Definizione del Problema
I modelli foundation per potenziali interatomici appresi tramite machine learning (MLIP) hanno dimostrato la capacità di trasferirsi attraverso diversi sistemi chimici, offrendo un flusso di lavoro che evita il processo dispendioso di risorse dell'addestramento di potenziali specifici per ogni compito partendo da zero. Tuttavia, la comunità manca di una guida sistematica su come e quando effettuare il fine-tuning di questi modelli. Rapporti preliminari hanno suggerito che un fine-tuning ingenuo porti spesso al "dimenticanza catastrofica" (catastrophic forgetting), promuovendo l'adozione di tecniche vincolate (ad es., congelamento dei livelli, Low-Rank Adaptation) originariamente sviluppate per i grandi modelli linguistici. Questo articolo investiga se tali vincoli siano necessari o se i fallimenti iniziali siano dovuti ad altri fattori, come modelli foundation più deboli, un'inizializzazione errata dell'energia di riferimento atomico ($E_0$) o procedure di addestramento instabili. Lo studio mira a caratterizzare i principali fattori che determinano gli esiti del fine-tuning, specificamente l'accuratezza sul compito target e la robustezza fuori distribuzione (OOD).

Metodologia
Gli autori valutano sette diverse strategie di fine-tuning su cinque benchmark chimicamente diversificati, tre generazioni di modelli foundation e set di addestramento che spaziano per cinque ordini di grandezza in termini di dimensione.

• Strategie di Fine-tuning Valutate:

  1. Naive (Ingenua: Aggiornamenti completi dei parametri tramite gradiente discendente continuato.
  2. Congelamento dei Livelli (Varianti): Congelamento dei livelli di embedding/passaggio di messaggi mentre si addestrano i readout; o congelamento dell'embedding e del primo livello di passaggio di messaggi.
  3. Low-Rank Adaptation (LoRA): Iniezione di decomposizioni a basso rango addestrabili sia nei livelli scalari che in quelli equivarianti, mantenendo congelati i pesi pre-addestrati.
  4. Multihead Replay: Ottimizzazione simultanea su dati target e un dataset di replay (dal pre-addestramento o con etichette pseudo-generate) utilizzando teste di readout separate.
  5. Pseudolabel Replay: Una variante del multihead replay in cui le etichette di replay sono generate dallo stesso modello foundation, scindendo la fonte del replay dal corpus originale di pre-addestramento.
  6. Replay + LoRA: Combinazione di multihead replay con LoRA.

• Benchmark: Lo studio abbraccia sistemi con crescente divergenza dal dominio di pre-addestramento OMat24 (solido inorganico periodico):

  • Elettroliti argiroditi di litio (solido periodico inorganico).
  • NaCl acquoso (soluzione ionica).
  • Polimorfi del ghiaccio (solido molecolare).
  • Reazioni SN2 (chimica reattiva in fase gassosa).
  • Biomolecole SPICE (conformeri organici/biomolecolari).

• Implementazioni Tecniche: Gli autori hanno implementato tre nuove capacità nel codebase MACE:

  • LoRA adattato per architetture di passaggio di messaggi equivarianti (coprendo sia i livelli scalari che quelli equivarianti).
  • Pseudolabelled replay per scindere le fonti dei dati di replay.
  • Riestima del modello-consapevole dell'energia di riferimento atomico ($E_0$) per allineare i baseline pre-addestrati ai dati target.

• Metriche di Valutazione: Oltre agli errori standard puntuali di energia e forza, lo studio sonda i comportamenti dinamici ed estrapolativi, inclusi i profili di distribuzione radiale (RDF) da dinamica molecolare (MD), profili di reazione Nudged Elastic Band (NEB), test di stabilità MD e la ricerca di strutture casuali (RSS) per rilevare fallimenti della repulsione a corto raggio.

Risultati Chiave

1. I Prerequisiti Dominano la Scelta della Strategia: Lo studio rileva che la qualità del modello foundation, la corretta inizializzazione di $E_0$ e la scelta accurata degli iperparametri sono prerequisiti il cui impatto supera regolarmente quello della specifica strategia di fine-tuning.

   • Qualità del Foundation: I nuovi modelli foundation (ad es., basati su OMat24) superano costantemente i precedenti (basati su MPTraj) nel trasferimento OOD, anche con ricette di fine-tuning fisse.
   • Inizializzazione di $E_0$: L'uso di valori di $E_0$ "mediati" porta a errori significativamente più elevati e instabilità MD (ad es., i modelli del ghiaccio falliscono entro 50 ps). Gli $E_0$ "riestimati" (che allineano lo zero-point del modello pre-addestrato ai dati target) sono critici per la stabilità e la trasferibilità, fornendo spesso risultati migliori rispetto alla scelta dell'algoritmo di fine-tuning stesso.
   • Iperparametri: Il fine-tuning ingenuo richiede tassi di apprendimento ridotti e un maggiore decadimento EMA. LoRA tollera tassi di apprendimento più elevati. Il multihead replay richiede tassi di apprendimento sostanzialmente più bassi per evitare segnali di aggiornamento contrastanti. Il weight decay dovrebbe essere impostato a zero per evitare di allontanare i parametri dalla soluzione pre-addestrata.

2. Prestazioni per Obiettivo:

   • Specializzazione In-Distribution (Singolo Sistema): Per compiti ristretti (ad es., barriere SN2, solvatazione di NaCl acquoso), la maggior parte delle strategie (Naive, LoRA, Multihead) raggiunge un'elevata accuratezza, superando costantemente i modelli addestrati da zero. Il fine-tuning ingenuo offre la migliore convergenza per applicazioni a sistema singolo.
   • Robustezza Out-of-Distribution: Quando si valuta il trasferimento a composizioni correlate ma non viste (ad es., elettroliti non argiroditi) o chimiche diverse (ad es., biomolecole), il Multihead Replay (con dati originali o pseudolabelled) è l'unico approccio che mantiene costantemente la robustezza OOD. Esso preserva l'accuratezza sulla distribuzione di pre-addestramento mentre apprende il compito target, prevenendo efficacemente la dimenticanza catastrofica.
   • Freezing e LoRA: Sebbene efficaci per l'efficienza dei parametri, il congelamento dei livelli e LoRA hanno mostrato limitazioni nell'adattarsi a caratteristiche di solvatazione o nel mantenere una vasta robustezza chimica rispetto al multihead replay nei casi testati.

Significatività e Rivendicazioni
L'articolo sostiene che la percepita fragilità del fine-tuning ingenuo negli MLIP è in gran parte il risultato di una configurazione subottimale piuttosto che di un limite intrinseco del metodo. Gli autori sostengono che:

• Il fine-tuning ingenuo è un punto di partenza valido e spesso superiore per applicazioni a sistema singolo, a condizione che il modello foundation sia di alta qualità e che gli $E_0$ siano correttamente riestimati.
• Il multihead replay è la strategia necessaria per un deployment più ampio dove è richiesto il mantenimento del comportamento del modello foundation al di fuori della distribuzione di fine-tuning.
• Il pseudolabelled replay offre un vantaggio pratico consentendo l'uso di qualsiasi dataset strutturalmente diversificato per il replay, rimuovendo la dipendenza dall'accesso al corpus originale di pre-addestramento.

Il lavoro stabilisce che, per i professionisti, investire nel modello foundation più forte e garantire il corretto allineamento dell'energia di riferimento atomico sono scelte progettuali più critiche rispetto alla selezione di uno specifico algoritmo di fine-tuning vincolato. Lo studio fornisce un quadro sistematico per il deployment dei modelli foundation MLIP, spostando il fine-tuning da un'opzione di nicchia a un punto di partenza predefinito per lo sviluppo specifico di sistemi.