Learning Inter-Atomic Potentials without Explicit Equivariance

1. Il Problema: Limiti degli attuali Modelli MLIP

Le simulazioni a livello atomico sono fondamentali per la scoperta di farmaci e la progettazione di materiali. Attualmente, la Teoria del Funzionale Densità (DFT) è lo standard per il calcolo accurato di energie e forze interatomiche, ma i suoi costi computazionali (che scalano cubicamente con la dimensione del sistema) rendono impossibili simulazioni su larga scala o su lunghi tempi.
Per ovviare a ciò, si utilizzano Potenziali Interatomici Appresi tramite Macchina (MLIP). Lo stato dell'arte attuale si basa su reti neurali equivarianti (ad esempio, GNN equivarianti) che incorporano rigidamente le simmetrie di rotazione e traslazione (SO(3) o SE(3)) nell'architettura stessa. Sebbene efficaci, questi approcci presentano limiti significativi:

• Rigidità architetturale: L'incorporazione "hard-wired" di vincoli di simmetria riduce la flessibilità del modello.
• Efficienza computazionale: Meccanismi come armoniche sferiche o messaggi equivarianti sono computazionalmente costosi.
• Scalabilità: La complessità di queste architetture può ostacolare l'addestramento su modelli molto grandi o dataset massicci.

L'obiettivo del lavoro è sviluppare un modello MLIP che sia scalabile ed efficiente (basato su Transformer generici) ma che riesca comunque a rispettare le simmetrie fisiche necessarie senza imporre vincoli architetturali espliciti.

2. Metodologia: TransIP (Transformer-based Inter-Atomic Potentials)

Gli autori propongono TransIP, un nuovo paradigma di addestramento che ottiene l'equivarianza SO(3) attraverso l'apprendimento nello spazio latente, anziché tramite vincoli architetturali.

Architettura di Base

• Backbone Transformer: Il modello utilizza un Transformer standard non equivariante. Ogni atomo è trattato come un "token" in una sequenza variabile.
• Input: Le caratteristiche atomiche includono coordinate ($r$), numeri atomici ($z$), carica totale ($q$) e molteplicità di spin ($s$). Le coordinate sono centrate e le proprietà globali ($q, s$) vengono incorporate come bias appresi aggiunti a ogni layer.
• Predizione: Un aggregatore permutazionale invariante combina gli embedding degli atomi per predire l'energia totale. Le forze sono calcolate come gradienti conservativi dell'energia rispetto alle posizioni atomiche ($F = -\nabla_r E$).

Apprendimento dell'Equivarianza Latente

Il cuore dell'innovazione risiede nella perdita di equivarianza latente appresa:

1. Rete di Trasformazione ($T_\tau$): Viene introdotta una rete neurale (MLP) che impara come le azioni del gruppo di simmetria (rotazioni $g \in SO(3)$) agiscono sugli embedding latenti. Invece di definire analiticamente la rappresentazione $\rho(g)$, il modello la impara.
2. Obiettivo Contrastivo: Viene definita una funzione di perdita che forza l'embedding di una molecola ruotata a corrispondere alla trasformazione appresa dell'embedding originale:
   $$L_{leq} = \| f(\phi(g)(m)) - T_\tau(\phi(g), f(m)) \|^2$$
   Dove $f$ è l'encoder Transformer, $\phi(g)$ è la rotazione delle coordinate di input, e $T_\tau$ è la rete che mappa l'azione del gruppo nello spazio latente.
3. Obiettivo Totale: L'addestramento combina tre termini: perdita sull'energia ($L_E$), perdita sulle forze ($L_F$) e la perdita di equivarianza latente ($L_{leq}$).

3. Contributi Chiave

1. Pipeline di Addestramento in Una Fase: TransIP raggiunge l'equivarianza SO(3) tramite un singolo processo di addestramento su un modello Transformer generico, eliminando la necessità di strati equivarianti complessi o framework di pre-training/fine-tuning separati.
2. Funzione di Perdita Agnostica all'Architettura: Viene introdotta una perdita contrastiva che promuove l'equivarianza nello spazio di embedding. Questo approccio scala meglio rispetto alle tecniche tradizionali di aumento dei dati (data augmentation) su diversi dataset e dimensioni di modello.
3. Performance Competitiva: Dimostrato su un benchmark diversificato (Open Molecules 25), TransIP supera le tecniche di aumento dei dati e raggiunge performance paragonabili ai modelli equivarianti più avanzati (SOTA), pur mantenendo un'architettura più semplice e scalabile.

4. Risultati Sperimentali

Il modello è stato valutato sul dataset Open Molecules 2025 (OMol25), che include oltre 4 milioni di molecole (organiche, biomolecole, complessi metallici, elettroliti).

• Scalabilità dei Dati:
  • In regime di dati limitati (1M campioni), TransIP supera di gran lunga il baseline con aumento dei dati (TransAug), riducendo l'errore MAE sulle forze da 600 a 255 meV/Å e migliorando la coerenza direzionale.
  • Anche su dataset più grandi (4M), TransIP mantiene un vantaggio significativo sulla predizione delle forze rispetto all'aumento dei dati, sebbene le prestazioni energetiche tendano a convergere.
• Relazione tra Errore Latente e Accuratezza: È stata osservata una forte correlazione monotona: minori errori di equivarianza latente corrispondono a migliori prestazioni complessive. La riduzione dell'errore latente del 48% (da 1M a 4M campioni) si traduce in un miglioramento delle prestazioni del 40-60%.
• Efficienza di Inferenza:
  • TransIP scala fluidamente con l'aumento dei parametri (Small, Medium, Large).
  • Al contrario, i modelli basati su aumento dei dati (TransAug) mostrano una scarsa scalabilità: modelli più grandi con aumento dati performano peggio di quelli più piccoli.
  • Velocità: TransIP è significativamente più veloce in inferenza rispetto ai baselines equivarianti SOTA (come eSEN e GemNet-OC), raggiungendo fino a 160.000 atomi/secondo per la versione Small.
• Confronto con SOTA: La versione Large di TransIP (addestrata per 80 epoche) supera le versioni "small" dei modelli equivarianti SOTA (eSEN-sm) nella predizione dell'energia totale e risulta competitiva sulle forze.

5. Significato e Conclusioni

Il lavoro dimostra che l'equivarianza può essere appresa efficacemente invece di essere imposta rigidamente.

• Flessibilità vs. Simmetria: TransIP rompe il compromesso tradizionale tra flessibilità architetturale (tipica dei Transformer) e rispetto delle simmetrie fisiche.
• Efficienza: L'approccio basato su Transformer è più efficiente in termini computazionali e di hardware rispetto alle architetture equivarianti specializzate (che usano armoniche sferiche).
• Impatto Futuro: Con la crescita esponenziale dei dati chimici e la necessità di modelli sempre più grandi, TransIP offre una via scalabile per modellare potenziali interatomici complessi. L'uso di un'architettura generica con un obiettivo di equivarianza latente apre la strada a modelli più semplici, veloci e potenti per la simulazione molecolare.

In sintesi, TransIP valida l'ipotesi che un modello generico, se guidato da un obiettivo di apprendimento appropriato nello spazio latente, possa catturare le simmetrie fisiche fondamentali con un'efficienza superiore rispetto alle soluzioni architetturalmente vincolate.