Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

Il paper introduce la Compositional Probe Decomposition (CPD) per dimostrare che l'allineamento tra il compito di addestramento e la proprietà target, insieme all'architettura equivariante, determina il grado di disaccoppiamento lineare tra informazioni geometriche e composizionali nei modelli fondazionali atomistici, rivelando che canali vettoriali e scalari codificano selettivamente diverse proprietà fisiche.

L'essenziale

Ecco una spiegazione semplice e creativa di questo articolo scientifico, pensata per chiunque voglia capire come funzionano i "cervelli" artificiali che studiano le molecole.

🧪 Il Problema: Il "Rumore" della Ricetta vs. la Forma del Pasticcio

Immagina di avere un cuoco robot (un'intelligenza artificiale) che deve prevedere le proprietà di un dolce. Per farlo, il robot guarda due cose:

1. Gli ingredienti (la ricetta: quanta farina, zucchero, uova).
2. La forma (come sono impastati: è una torta alta, un biscotto croccante o una crema liscia).

Il problema è che spesso gli ingredienti e la forma sono legati. Se hai molte uova, probabilmente farai una torta alta. Se vuoi prevedere quanto è dolce il dolce, il robot potrebbe essere pigro e dire: "Ah, vedo che ci sono molte uova, quindi sarà dolce", senza nemmeno guardare la forma.

Gli scienziati volevano sapere: i modelli di intelligenza artificiale più avanzati riescono a separare chiaramente gli ingredienti dalla forma? O mescolano tutto insieme, costringendo chi usa il modello a fare un lavoro extra per capire cosa sta succedendo?

🔍 La Soluzione: Il "Filtro Magico" (CPD)

Gli autori hanno creato un nuovo metodo chiamato CPD (Decomposizione della Sonda Compositiva). Immagina questo metodo come un filtro magico o un setaccio molto preciso.

1. Il Filtro: Prendono le informazioni che il robot ha memorizzato e usano il filtro per rimuovere esattamente la parte relativa agli ingredienti (la composizione).
2. Il Test: Poi, guardano cosa rimane nel secchio (il "residuo"). Se il robot ha imparato bene la forma, nel secchio dovrebbe esserci ancora molta informazione sulla geometria. Se il robot era pigro e ha solo memorizzato gli ingredienti, nel secchio non rimane nulla.

Attenzione all'inganno: Hanno scoperto che se usi un "testatore" troppo intelligente (come un albero decisionale complesso) per guardare il residuo, questo riesce a ricostruire gli ingredienti che avevi appena buttato via, fingendo di aver trovato informazioni sulla forma. È come se un detective, dopo aver pulito la scena del crimine, usasse la sua immaginazione per "inventare" prove che non esistono. Per questo, l'articolo consiglia di usare solo testatori semplici e lineari (come una riga retta) per essere sicuri di vedere la verità.

🏆 I Risultati: Chi vince la gara?

Hanno testato 10 diversi "cervelli" artificiali. Ecco cosa hanno scoperto, usando tre fattori chiave:

1. L'Allenamento è tutto (Task Alignment)

Questo è il fattore più importante.

• L'analogia: Immagina due studenti. Uno studia per un esame di Geometria (prevede la forma delle molecole), l'altro studia per un esame di Chimica di base (prevede l'energia totale, che dipende molto dagli ingredienti).
• Il risultato: Lo studente che ha studiato Geometria è molto meglio nel separare la forma dagli ingredienti. Anche se lo studente di Chimica ha un cervello più potente (architettura più complessa), se non ha mai esercitato a pensare alla forma, non la imparerà bene.
• In sintesi: Se vuoi che un modello capisca la forma, allenalo su compiti che dipendono dalla forma, non solo su quelli che dipendono dagli ingredienti.

2. La Diversità dei Dati aiuta (ma non basta)

• L'analogia: Uno studente che ha viaggiato in tutto il mondo e visto milioni di case diverse (dati diversificati) impara meglio la struttura degli edifici rispetto a uno che ha visto solo 100 case dello stesso quartiere.
• Il risultato: I modelli addestrati su enormi quantità di dati diversi sono migliori di quelli addestrati su dati piccoli, ma non riescono a superare lo studente che si è specializzato proprio nel compito specifico. La diversità aiuta, ma l'allenamento mirato vince.

3. L'Architettura conta (ma solo se combinata)

• L'analogia: Avere un'auto sportiva (architettura avanzata) è utile, ma se guidi su una strada sterrata sbagliata (obiettivo di allenamento sbagliato), non arriverai mai in tempo.
• Il risultato: Le architetture moderne (che rispettano le leggi della fisica, come la rotazione) sono potenti, ma solo se combinate con un obiettivo di allenamento corretto. Se combini l'architettura giusta con l'allenamento sbagliato, il modello performa peggio di modelli più semplici!

🧬 La Scoperta Sorprendente: Le "Autostrade" dei Dati

Hanno scoperto che in alcuni modelli (come MACE), le informazioni viaggiano su "corsie" diverse, proprio come le autostrade:

• Le informazioni scalari (come la differenza di energia tra orbitali) viaggiano su corsie "piatte" (canali scalari).
• Le informazioni vettoriali (come la direzione di un campo magnetico) viaggiano su corsie "curve" (canali vettoriali).

È come se il modello avesse organizzato il suo magazzino: gli oggetti rotondi vanno nello scaffale rotondo, gli oggetti allungati in quello allungato. Questo rende tutto molto ordinato e facile da trovare. Altri modelli, invece, buttano tutto in un unico grande mucchio, rendendo difficile trovare l'informazione specifica.

💡 Cosa significa per noi?

1. Scegliere il modello giusto: Se vuoi analizzare le proprietà geometriche di una molecola (ad esempio per progettare nuovi farmaci), non scegliere il modello più "famoso" o complesso. Scegli quello che è stato addestrato su compiti simili a quello che vuoi fare.
2. Attenzione alle false certezze: Non fidarti ciecamente di test complessi che dicono "il modello sa tutto". A volte stanno solo "indovinando" o ricostruendo informazioni che avevamo rimosso. Usa metodi semplici e lineari per essere sicuro.
3. L'ordine è potere: I modelli che organizzano le informazioni in modo pulito (separando ingredienti da forma) sono molto più efficienti. Hanno bisogno di meno esempi per imparare e sono più facili da capire.

In parole povere: non importa quanto è potente il motore della tua auto (l'architettura), importa soprattutto dove hai imparato a guidare (l'allenamento) e quanto bene hai ordinato il tuo bagagliaio (la struttura delle rappresentazioni).

Sommario tecnico

Ecco un riassunto tecnico dettagliato del paper "Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement" di Joshua Steier.

1. Il Problema

I modelli di base atomistici (come MACE, SchNet, PaiNN, ViSNet) hanno raggiunto livelli di accuratezza vicini alla Teoria del Funzionale Densità (DFT) nella previsione di proprietà molecolari. Tuttavia, rimane una domanda fondamentale aperta: come sono organizzate le loro rappresentazioni interne?

Nello specifico, le rappresentazioni separano in modo pulito due fattori critici:

1. Composizione: Di quali elementi è fatta la molecola e in quali proporzioni.
2. Geometria: Come sono disposti gli atomi nello spazio.

Le proprietà molecolari sono spesso correlate a entrambi i fattori. Le tecniche di probing (sondaggio) tradizionali falliscono perché non riescono a distinguere se un modello sta codificando realmente la geometria o se sta semplicemente "barando" sfruttando la composizione (che è un forte predittore per molte proprietà). Inoltre, l'uso di sonde non lineari (come Gradient Boosted Trees) su rappresentazioni residuali può portare a risultati fuorvianti, ricostruendo artificialmente il segnale composto rimosso.

2. Metodologia: Compositional Probe Decomposition (CPD)

L'autore introduce CPD, un nuovo protocollo di sondaggio progettato per isolare l'informazione geometrica rimuovendo linearmente il segnale di composizione.

• Definizione delle Feature di Composizione: Per ogni molecola, viene costruito un vettore $z$ contenente le frazioni degli elementi (C, H, N, O, F) e il numero di atomi standardizzato.
• Proiezione OLS (Ordinary Least Squares): Le rappresentazioni del modello ($X$) vengono proiettate sul sottospazio definito dalla composizione ($Z$) per ottenere i residui:
  $$X_{geom} = X - Z\hat{\beta}$$
  Questo residuo $X_{geom}$ contiene l'informazione linearmente ortogonale alla composizione (topologia, connettività, conformazione).
• Probing Lineare: Viene utilizzata una regressione Ridge sui residui $X_{geom}$ per misurare quanto è accessibile l'informazione geometrica ($R^2_{geom}$).
• Validazione Critica: Il paper dimostra che le sonde non lineari (es. Gradient Boosted Trees - GBT) falliscono su questo compito: riescono a ricostruire il segnale di composizione rimosso dai residui, producendo $R^2$ artificialmente alti (0.68–0.95) su target puramente composizionali. Di conseguenza, solo le sonde lineari forniscono una misura fedele dell'accessibilità lineare.
• Benchmark di Isomeri Strutturali: Per validare il metodo, viene utilizzato un set di isomeri (stessa composizione, diversa geometria). La componente di composizione deve avere un'accuratezza al caso (50%), mentre il residuo geometrico deve distinguere gli isomeri.

3. Contributi Chiave

1. Metodologia CPD Validata: Un nuovo standard per il probing che separa composizione e geometria, validato da 12 controlli di robustezza (incluso il benchmark sugli isomeri e il confronto con LEACE).
2. Il Gradiente di Accessibilità Lineare: Identificazione di una vasta variazione (fattore di 6.6x) nella quantità di informazione geometrica accessibile linearmente tra diversi modelli, anche a parità di architettura.
3. Dominanza dell'Allineamento del Task: Dimostrazione che l'obiettivo di addestramento (training objective) è il fattore dominante, superando l'architettura e la diversità dei dati.
4. Routing dell'Informazione per Simmetria: Scoperta che in architetture specifiche (MACE), l'informazione viene instradata selettivamente attraverso canali di simmetria diversi (scalari vs vettoriali) in base alla natura fisica della proprietà.

4. Risultati Principali

A. Il Gradiente di Accessibilità Lineare

Su 10 modelli (5 famiglie architetturali) addestrati su QM9, l'accessibilità geometrica ($R^2_{geom}$ per il gap HOMO-LUMO) varia da 0.081 a 0.533.
I tre fattori che spiegano questo gradiente sono:

1. Allineamento del Task (Dominante):

   • I modelli addestrati direttamente sul gap HOMO-LUMO (una proprietà sensibile alla geometria) ottengono $R^2_{geom}$ molto più alti (0.44–0.53) rispetto a quelli addestrati sull'Energia Totale (dove la composizione domina).
   • L'effetto è enorme: un modello PaiNN addestrato su HOMO-LUMO ha un $R^2_{geom}$ di 0.533, mentre la stessa architettura addestrata sull'energia scende a 0.310 ($\Delta \approx 0.22$).
   • Questo vale indipendentemente dall'architettura (confermato su PaiNN e MACE).

2. Equivarianza (Condizionale):

   • Contrariamente all'ipotesi comune, l'uso di architetture equivarianti (che rispettano le simmetrie rotazionali) non garantisce di per sé un'alta accessibilità geometrica se il task non è allineato.
   • MACE (equivariante) addestrato solo su energia e QM9 ha un $R^2_{geom}$ di 0.081, peggio di modelli invarianti più semplici come SchNet (0.262).
   • L'equivalenza è utile solo se combinata con un obiettivo di addestramento sensibile alla geometria.

3. Diversità dei Dati (Compensativa):

   • La pre-addestramento su dataset vasti e diversificati (es. MPTraj per MACE) compensa parzialmente la mancanza di allineamento del task. MACE pre-addestrato su MPTraj raggiunge 0.364, superando i modelli addestrati solo su QM9, ma non riesce a raggiungere il livello dei modelli task-aligned (~0.52).

B. Routing dell'Informazione in MACE

Analizzando i canali di rappresentazione di MACE (che decompongono le rappresentazioni in irriducibili di SO(3)):

• I canali L=0 (scalari) codificano prevalentemente il gap HOMO-LUMO ($R^2 = 0.76$).
• I canali L=1 (vettoriali) codificano prevalentemente il momento di dipolo ($R^2 = 0.59$).
• Questo pattern di "instradamento" per tipo di simmetria è assente in ViSNet, che concentra quasi tutta l'informazione nei canali scalari, suggerendo che l'architettura di MACE (basata su prodotti tensoriali) mantiene una struttura leggibile esternamente che ViSNet non possiede.

C. Efficienza del Campione

I modelli con rappresentazioni linearmente disaccoppiate (come PaiNN addestrato su HOMO-LUMO) richiedono drasticamente meno dati per estrarre segnali geometrici. A soli 50 molecole, PaiNN supera le prestazioni di SchNet addestrato su 2000 molecole.

5. Significato e Implicazioni

• Per la Selezione dei Modelli: Quando si sceglie un modello pre-addestrato per un downstream task, l'obiettivo di addestramento è più importante dell'architettura. Un modello equivariante addestrato sull'energia non è necessariamente un punto di partenza migliore di un modello invariante se il task target è sensibile alla geometria.
• Interpretabilità: Le rappresentazioni non sono "scatole nere" uniformi; l'organizzazione interna dipende criticamente da come il modello è stato addestrato. La disaccoppiamento lineare è una proprietà acquisita, non intrinseca.
• Metodologia di Probing: Il paper avverte che l'uso di sonde non lineari su rappresentazioni "pulite" (residualizzate) è metodologicamente errato, poiché può ricostruire artificialmente il segnale rimosso. Le sonde lineari sono necessarie per valutare l'accessibilità lineare reale.
• Generalizzazione: I risultati si estendono anche ai cristalli inorganici (Materials Project), confermando che la separazione composizione-geometria e il ruolo dell'allineamento del task sono principi generali nell'apprendimento automatico molecolare.

In sintesi, il lavoro dimostra che la "qualità" di una rappresentazione per compiti geometrici non è data dalla complessità architetturale, ma da quanto l'obiettivo di addestramento ha forzato il modello a organizzare le informazioni geometriche in modo linearmente accessibile.