Riepilogo Tecnico: Modelli Fondamentali Atomistici Equivarianti SO(3) Compatti tramite Potatura Strutturale

Enunciato del Problema

Le reti neurali grafo (GNN) equivarianti SO(3) si sono affermate come il paradigma dominante per i modelli fondamentali atomistici, raggiungendo alta accuratezza ed efficienza nei dati codificando esplicitamente la simmetria rotazionale. Tuttavia, questi modelli si basano su operazioni di prodotto tensoriale di ordine superiore per accoppiare le caratteristiche geometriche, risultando in costi computazionali che scalano fino a $O(L_{max}^6)$, dove $L_{max}$ è l'ordine massimo delle rappresentazioni irriducibili. Ciò crea un compromesso significativo tra accuratezza del modello ed efficienza di inferenza.

Le strategie esistenti di compressione del modello—nello specifico Distillazione della Conoscenza (KD) e Quantizzazione—presentano limitazioni in questo dominio. La KD richiede spesso l'addestramento di uno studente da zero e non riutilizza direttamente i pesi dell'insegnante, mentre la quantizzazione riduce la precisione in bit ma lascia invariato il numero di parametri e i percorsi del prodotto tensoriale. Inoltre, i metodi standard di potatura strutturale sviluppati per i Large Language Models (LLM) non possono essere applicati direttamente ai modelli equivarianti SO(3). Rimuovere in modo ingenuo le dimensioni delle caratteristiche rompe la simmetria rotazionale del percorso del prodotto tensoriale, violando la proprietà fondamentale di equivarianza richiesta per la coerenza fisica. I precedenti tentativi di potatura di modelli atomistici si sono limitati alla rimozione di blocchi a grana grossa (ad esempio, interi strati di interazione), il che riduce il campo ricettivo del modello ed è particolarmente dannoso per le architetture superficiali.

Metodologia

Gli autori propongono un framework di potatura strutturale specificamente progettato per i modelli fondamentali atomistici equivarianti SO(3). Il metodo opera su checkpoint pre-addestrati per produrre modelli più piccoli e densi che mantengono l'equivarianza SO(3) e l'alta accuratezza.

1. Granularità della Potatura Equivariante

L'innovazione principale risiede nella definizione dell'unità di potatura. Nei modelli equivarianti SO(3), le caratteristiche sono decomposte in rappresentazioni irriducibili di ordine $l \in [0, L_{max}]$. Un vettore di caratteristiche per un canale specifico $k$ e ordine $l$ consiste in $2l+1$ componenti indicizzate da $m \in [-l, l]$.

• Vincolo: L'operazione di rotazione in SO(3) accoppia tutte le componenti $m$ all'interno di un ordine specifico $l$. Rimuovere singole componenti (potatura elemento per elemento) rompe l'equivarianza.
• Soluzione: Il metodo tratta lo spazio vettoriale generato dalle $2l+1$ componenti di un ordine specifico $l$ e canale $k$ come un'unità atomica. La potatura viene eseguita sulla terna $(k, l)$, rappresentata da una maschera binaria $z^{(t)}_{k,l}$. Ciò garantisce che, se un blocco di caratteristiche viene mantenuto, tutte le sue componenti rotazionali sono preservate, mantenendo rigorosamente l'equivarianza SO(3).

2. Criterio di Importanza Informato dalla Fisica

Per determinare quali blocchi $(k, l)$ mantenere, gli autori propongono un punteggio di importanza basato sulla sensibilità della superficie di energia potenziale.

• Obiettivo Duale: Gli MLIP devono prevedere sia l'energia potenziale ($E$) che le forze ($F_i = -\nabla_{r_i} E$).
• Metrica: Il punteggio di importanza $I^{(t)}_{k,l}$ è derivato da un'espansione di Taylor del primo ordine della variazione di energia al momento della rimozione di un blocco di caratteristiche. Aggrega il prodotto della magnitudine della caratteristica e del gradiente dell'energia rispetto a tale caratteristica su tutti gli atomi e le dimensioni di rappresentazione:
  $$I^{(t)}_{k,l} = \mathbb{E}_{R,Z \sim \mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h^{(t)}_{j,klm}} \cdot h^{(t)}_{j,klm} \right|$$
• Vantaggio: Questo criterio è invariante SO(3) e comporta un sovraccarico trascurabile perché i gradienti $\frac{\partial E}{\partial h}$ sono già calcolati durante l'inferenza standard delle forze.

3. Allineamento Strutturale e Pipeline

Il framework segue una pipeline a quattro stadi:

1. Calibrazione: Un piccolo sottoinsieme dei dati di pre-addestramento viene passato attraverso il modello grande per calcolare i punteggi di importanza.
2. Potatura Guidata dall'Obiettivo: Viene definita un'architettura target. I blocchi $(k, l)$ meglio classificati vengono mantenuti in base ai punteggi di importanza.
3. Allineamento Strutturale: La struttura del modello viene fisicamente alterata per adattarsi alle dimensioni ridotte. Ciò comporta:
   • Riconfigurazione del Prodotto Tensoriale: Taglio dei pesi radiali e dei percorsi di Clebsch-Gordan per allinearsi ai canali mantenuti.
   • Taglio dei Pesi: Rimozione fisica di righe e colonne dai pesi delle trasformazioni lineari per creare un tensore di pesi più piccolo e denso.
4. Riaddestramento e Fine-tuning: Il modello potato subisce un riaddestramento leggero su un sottoinsieme dei dati di pre-addestramento per recuperare dagli errori di approssimazione, seguito da un fine-tuning su specifici compiti a valle.

Contributi Chiave

1. Metodo di Potatura Equivariante: Introduzione di un metodo di potatura strutturale per tensori equivarianti SO(3) di ordine superiore attraverso le dimensioni del canale e dell'ordine che preserva rigorosamente l'equivarianza SO(3) trattando i blocchi $(k, l)$ come unità atomiche.
2. Framework a Quattro Stadi: Una pipeline completa (Calibrazione, Potatura, Riaddestramento, Fine-tuning) per la compressione dei modelli fondamentali.
3. Generalizzabilità: Dimostrazione che il metodo si applica a multiple architetture, inclusi MACE, SevenNet ed eSCN.
4. Combinabilità: Validazione che la potatura strutturale può essere combinata con quantizzazione e distillazione della conoscenza per ottenere guadagni di efficienza cumulativi.

Risultati Sperimentali

Il metodo è stato valutato principalmente su MACE-MP (materiali inorganici) e MACE-OFF (molecole organiche), con test di generalizzabilità su SevenNet ed eSCN.

Accuratezza vs. Addestramento da Zero

• Superiorità: I modelli potati superano costantemente i modelli della stessa dimensione addestrati da zero.
  • Su MACE-MP, i modelli potati hanno ottenuto un MAE delle forze inferiore e prestazioni migliori fuori distribuzione sul dataset sAlex rispetto ai baseline addestrati da zero.
  • Su Matbench Discovery, un modello MACE-MP potato (utilizzando solo il 25% dei dati di pre-addestramento) ha superato il modello piccolo ufficiale addestrato da zero in 7 metriche su 9.
  • Su MACE-OFF, i modelli potati hanno eguagliato o superato i baseline addestrati da zero su diversi sottoinsiemi molecolari, anche utilizzando solo il 50% dei dati di pre-addestramento.
• Prestazioni a Valle: Il fine-tuning di modelli potati su otto dataset inorganici a valle ha ridotto gli errori di energia e forza del 70,1% e del 34,4% rispettivamente, rispetto all'addestramento di modelli specifici per il compito da zero.

Guadagni di Efficienza

• Riduzione dei Parametri: I modelli compressi contengono da 1,5× a 4× parametri in meno.
• Costo di Addestramento: L'intera pipeline richiede da 2,5× a 4× meno calcolo di pre-addestramento rispetto all'addestramento di un modello piccolo da zero (ad esempio, 181 ore GPU contro 750 ore GPU per una configurazione specifica).
• Velocità di Inferenza: I modelli potati raggiungono un accelerazione di inferenza da 2,7× a 3,0× e una riduzione fino a 5,7× nell'uso di picco della memoria.

Combinazione con Altri Metodi

• Quantizzazione: La combinazione della potatura con la quantizzazione FP32 ha prodotto un miglioramento del throughput del 9,6× (rispetto al modello base FP64).
• Distillazione della Conoscenza: La combinazione della potatura con KD basata sull'Hessian ha ulteriormente affinato le prestazioni del modello, dimostrando che i due metodi sono complementari.

Significato e Affermazioni

Il documento afferma che la potatura strutturale offre un percorso naturale ed efficace per colmare il divario tra accuratezza ed efficienza nei modelli fondamentali atomistici. Ereditando direttamente i pesi da checkpoint grandi e pre-addestrati e comprimendoli preservando le simmetrie fisiche, il metodo produce modelli compatti che sono:

1. Più accurati dei modelli della stessa dimensione addestrati da zero.
2. Significativamente più economici da addestrare e distribuire.
3. Fisicamente coerenti, mantenendo la rigorosa equivarianza SO(3) richiesta per simulazioni affidabili di dinamica molecolare (validata tramite simulazioni MD di acqua liquida che mostrano una deriva energetica stabile e funzioni di distribuzione radiale accurate).

Gli autori sottolineano che il principio fondamentale—valutare l'importanza delle rappresentazioni irriducibili accoppiate—è ampiamente applicabile a qualsiasi architettura equivariante che rappresenti le caratteristiche come somme di rappresentazioni irriducibili SO(3). Il lavoro suggerisce che la potatura permette ai praticanti di distribuire modelli fondamentali ad alta fedeltà in ambienti con risorse limitate senza sacrificare la trasferibilità appresa durante il pre-addestramento su larga scala.