jNO: A JAX Library for Neural Operator and Foundation Model Training

jNO è una libreria unificata e nativa di JAX che semplifica l'addestramento di operatori neurali e modelli fondazionali integrando approcci basati sui dati e informati dalla fisica in un unico sistema di tracciamento simbolico, consentendo transizioni fluide tra regressione degli operatori, valutazione dei residui consapevole della mesh e ottimizzazione vincolata da equazioni differenziali alle derivate parziali senza ristrutturazione del codice.

L'essenziale

Immagina di dover insegnare a un computer a comprendere le leggi della fisica, come il flusso del calore attraverso una lastra metallica o il vortice dell'acqua attorno a una roccia. In passato, farlo con l'intelligenza artificiale era come tentare di costruire una casa in cui l'architetto, l'idraulico, l'elettricista e il falegname parlavano tutte lingue diverse e utilizzavano progetti differenti. Dovevi scrivere un set di codice per la forma della stanza (geometria), un altro per le equazioni matematiche (fisica) e un terzo per il processo effettivo di apprendimento. Se volevi passare da un tipo di matematica a un altro, spesso dovevi smantellare l'intera casa e ricominciare da capo.

jNO (jax Neural Operators) è un nuovo strumento che agisce come un "traduttore universale" e un "capomastro" tutto in uno. È una libreria software progettata per rendere molto più semplice, veloce e flessibile l'addestramento di questi modelli di IA consapevoli della fisica, specificamente per il linguaggio di programmazione JAX (uno strumento popolare per il calcolo scientifico ad alta velocità).

Ecco come funziona, utilizzando alcune semplici analogie:

1. La magia dello "Script Unico" (Sistema di Tracciamento)

Pensa a jNO come a un singolo script magico che controlla l'intero cantiere.

• Prima: Dovevi scrivere uno script per il progetto, uno script separato per la matematica e un altro per le regole di apprendimento. Se volevi cambiare la matematica, dovevi riscrivere anche lo script del progetto.
• Con jNO: Scrivi tutto in un'unica lingua. Definisci la forma della stanza, le equazioni della fisica e gli obiettivi di apprendimento tutto in una volta. Il software "traccia" (o registra) le tue istruzioni come un regista che gira una scena. In seguito, compila questo filmato in un programma super-efficiente e ad alta velocità. Questo significa che puoi passare tra diversi tipi di problemi matematici o aggiungere nuove regole fisiche senza riscrivere il tuo codice.

2. I modelli di base "Lego"

Attualmente, esistono molti diversi "modelli di base" (cervelli di IA pre-addestrati) per la fisica, ma sono come set di Lego di produttori diversi che non si adattano tra loro. Un marchio usa mattoncini rossi, un altro usa blu, e non possono essere impilati.

• Il ruolo di jNO: Agisce come un adattatore universale. Prende questi diversi modelli di IA (come Poseidon, Walrus e Morph) e li traduce in modo che tutti si inseriscano nello stesso ecosistema JAX. Ora, un ricercatore può prendere un "cervello" pre-addestrato, modificarlo leggermente e combinarlo con le proprie regole fisiche personalizzate, tutto senza dover cambiare strumenti software.

3. La "Mesh Intelligente" (Gestione delle Forme)

Quando si simulano fenomeni fisici, i computer devono scomporre le forme (come un tubo curvo o un edificio complesso) in piccoli pezzi di griglia chiamati "mesh".

• L'innovazione: jNO possiede un sistema di "mesh intelligente" integrato. È come avere un robot che può disegnare istantaneamente una griglia su qualsiasi forma tu descriva, sia essa un semplice quadrato o un oggetto 3D complesso con fori. Tiene traccia di quale parte della griglia è "interna", quale è "muro" e quale è "bordo", in modo che l'IA sappia esattamente dove applicare le regole fisiche.

4. La manopola "Fine-Tuning"

A volte si vuole prendere un'IA pre-addestrata e insegnarle un nuovo compito specifico.

• Il pannello di controllo: jNO ti offre un pannello di controllo molto dettagliato. Puoi dire all'IA: "Congela queste parti del tuo cervello in modo che non cambino", oppure "Impara solo da queste connessioni specifiche", o ancora "Usa una velocità di apprendimento specifica". Puoi farlo per singole parti del modello senza dover ricostruire l'intero sistema. È come poter regolare il volume solo dei tamburi in una canzone senza cambiare la chitarra o la voce.

5. Il motore "Dual-Mode" (FEM e PINNs)

Il documento evidenzia che jNO può gestire due modi diversi di risolvere problemi fisici:

• Punto per punto: Controllare la fisica in punti specifici (come controllare la temperatura in punti specifici su una mappa).
• Forma intera (Elementi Finiti): Considerare la fisica come un flusso continuo su un'intera forma (come calcolare lo stress totale su un ponte).
• Il vantaggio: jNO ti permette di passare tra queste due modalità utilizzando lo stesso codice. È come avere un'auto che può guidare sia su strade sterrate che su autostrade senza che tu debba cambiare il motore o lo sterzo.

Perché è importante?

L'obiettivo principale di jNO è fermare la "frammentazione" del software scientifico. Invece che i ricercatori giochino con cinque strumenti diversi per addestrare un singolo modello di IA, jNO porta tutto in un unico luogo.

• Velocità: Poiché utilizza le funzioni di compilazione speciali di JAX, esegue più velocemente sui moderni chip informatici.
• Semplicità: Non devi essere un architetto software per passare tra diversi tipi di problemi fisici.
• Riutilizzabilità: Una volta scritto un programma in jNO, puoi salvarlo, condividerlo ed eseguirlo di nuovo in seguito, anche su computer diversi, con la certezza che funzionerà allo stesso modo.

In breve, jNO sta cercando di rendere il mondo complesso dell'"apprendimento automatico scientifico" semplice e unificato quanto scrivere una singola storia coerente, piuttosto che cucire insieme un patchwork di diversi frammenti di codice.

Sommario tecnico

Riepilogo Tecnico: jNO – Una Libreria JAX per l'Addestramento di Operatori Neurali e Modelli Fondamentali

Enunciato del Problema

L'attuale panorama dell'apprendimento automatico scientifico (SciML), in particolare per quanto riguarda gli operatori neurali e le Reti Neurali Informate dalla Fisica (PINN), soffre di una significativa frammentazione software. I modelli fondamentali all'avanguardia (ad esempio, Poseidon, Walrus, PDEformer2) e gli strumenti per l'apprendimento degli operatori sono mantenuti prevalentemente in repository separati all'interno dell'ecosistema PyTorch. Questa separazione crea ostacoli al confronto sistematico, al trasferimento di apprendimento e al fine-tuning riproducibile tra diverse famiglie di modelli. Inoltre, i flussi di lavoro pratici sono spesso disgiunti, richiedendo agli utenti di gestire astrazioni separate per la geometria, la costruzione della funzione di perdita, la definizione del modello e la logica di addestramento. Tale frammentazione si aggrava quando si tenta di combinare l'apprendimento supervisionato degli operatori con residui informati dalla fisica o quando si trasferiscono backbones pre-addestrati in nuovi contesti di Equazioni Differenziali Parziali (PDE).

Metodologia

jNO (jax Neural Operators) affronta queste sfide fornendo una libreria unificata e nativa JAX che tratta l'intero flusso di lavoro di addestramento come un singolo programma tracciato. La metodologia di base si basa su un Linguaggio Specifico di Dominio (DSL) simbolico incentrato sul tracciamento e sull'esecuzione differita.

1. Design basato sul Tracciamento

A differenza dei modelli di esecuzione eager, jNO costruisce un grafo di calcolo differito in cui domini, chiamate al modello, residui, funzioni di perdita supervisionate e diagnostica sono scritti in un unico linguaggio simbolico.

• Nodi Simbolici: Il sistema utilizza nodi Placeholder per rappresentare variabili, tensori, modelli e operatori differenziali. Le operazioni aritmetiche e differenziali sovraccaricano le primitive JAX standard per costruire alberi di espressione invece di eseguire immediatamente.
• Pipeline Unificata: Gli utenti definiscono domini e variabili, compongono modelli con espressioni differenziali e risolvono tutto attraverso un'unica interfaccia. Ciò consente transizioni fluide tra regressione degli operatori, valutazione dei residui consapevoli della mesh e addestramento vincolato alle PDE senza ristrutturare il codice.
• Ottimizzazione: Prima dell'esecuzione, il sistema esegue l'eliminazione delle sotto-espressioni comuni (CSE) sull'albero tracciato, canonizzando sotto-alberi strutturalmente identici in un singolo programma XLA. Ciò unifica obiettivi, definizioni degli operatori e vincoli fisici.

2. Flusso di Lavoro di Dominio e Mesh

La classe domain agisce come un livello di organizzazione dei dati consapevole della mesh, fungendo da ponte tra le definizioni geometriche e le espressioni tracciate.

• Integrazione: Supporta la generazione di mesh tramite PyGmsh/Gmsh e il caricamento tramite meshio.
• Struttura: I dati sono organizzati in un pool di mesh (per insiemi di punti etichettati) e un dizionario di contesto runtime. Supporta l'addestramento in batch in cui ogni elemento del batch può rappresentare una diversa parametrizzazione PDE o campo di forzatura.
• Adattività: Il sistema include hook per la gestione adattiva dei punti, consentendo all'addestramento di aggiornare gli insiemi di punti senza alterare le equazioni simboliche.

3. Integrazione del Metodo agli Elementi Finiti (FEM)

jNO estende il suo modello di programmazione tracciata all'assemblaggio della forma debole tramite il backend FEAX.

• DSL Unificato: Invece di un front-end FEM separato, jNO utilizza lo stesso livello simbolico per forme deboli basate su mesh, PINN variazionali e apprendimento degli operatori.
• Assemblaggio: Le forme deboli simboliche sono abbassate a rappresentazioni neutrali rispetto al backend. A seconda dell'obiettivo, il sistema può instradare verso vpinn (PINN variazionale), fem_system (stazionario lineare), fem_residual (stazionario non lineare) o fem_time (transitorio).
• Differenziabilità: L'integrazione mantiene gli operatori agli elementi finiti assemblati all'interno dell'ecosistema JAX, consentendo flussi di lavoro differenziabili end-to-end (assemblaggio, valutazione del residuo, costruzione dello Jacobiano) compatibili con le trasformazioni JAX.

4. Interfaccia del Modello e Modelli Fondamentali

jNO tratta le architetture degli operatori neurali come operazioni ordinarie all'interno del programma tracciato.

• Modelli Fondamentali: Supporta traduzioni native JAX di modelli fondamentali per le PDE (Poseidon, Walrus, Morph, ecc.) mantenuti nel repository esterno foundax. Questi modelli possono essere incapsulati tramite jno.nn.wrap() e composti con espressioni tracciate arbitrarie.
• Controllo Granulare: Gli utenti possono configurare il comportamento del modello, dell'ottimizzatore e del tasso di apprendimento a livello di parametro. Le funzionalità includono mascheramento selettivo dei parametri, congelamento/scongelamento e supporto per LoRA (Low-Rank Adaptation) sia per modelli Equinox che Flax.

5. Esecuzione e Addestramento

Il modulo jno.core compila i vincoli tracciati in un singolo programma di addestramento compilato JIT.

• Prestazioni: Sfrutta XLA per l'eliminazione aggressiva delle sotto-espressioni, l'ottimizzazione della memoria e il parallelismo su più dispositivi (tramite una mesh di dispositivi a 2 assi).
• Ottimizzazione degli Iperparametri: L'interfaccia ArchSpace supporta la ricerca a griglia e l'ottimizzazione senza gradienti (tramite Nevergrad) per parametri architetturali e di addestramento, integrati direttamente nel ciclo di addestramento.
• Persistenza: Gli esperimenti possono essere serializzati utilizzando jno.save (che supporta firme RSA per l'integrità) per preservare gli stati del solver, gli oggetti di dominio e i wrapper di inferenza compilati.

Contributi Chiave

1. Linguaggio Tracciato Unificato: jNO è il primo framework consolidato per esprimere operatori neurali, residui informati dalla fisica e fine-tuning di modelli fondamentali all'interno di un'unica interfaccia simbolica in JAX.
2. Consolidamento dell'Ecosistema: Agisce come backend unificato per tradurre e integrare diverse famiglie di modelli fondamentali per le PDE (ad esempio, Poseidon, Walrus) in JAX, consentendo valutazione e composizione affiancate senza mantenere basi di codice separate.
3. Unificazione FEM e Operatori: Estendendo il DSL tracciato all'assemblaggio della forma debole, jNO permette di comporre forme deboli basate su mesh, PINN variazionali e modelli di operatori senza cambiare il modello di programmazione.
4. Controllo a Livello di Parametro: La libreria fornisce un controllo granulare sulla dinamica di addestramento, inclusi mascheramento selettivo e adattamento LoRA, direttamente all'interno del flusso simbolico.
5. Riproducibilità e Portabilità: La libreria è distribuita per PyPI, include test automatizzati e supporta la condivisione di artefatti firmati per flussi di lavoro verificabili.

Risultati e Affermazioni

Il documento non presenta nuovi risultati di benchmark su specifici dataset PDE, ma si concentra sui contributi architettonici e di flusso di lavoro:

• Efficienza del Flusso di Lavoro: Il documento afferma che il design tracciato elimina la necessità di ristrutturare il codice quando si passa tra apprendimento degli operatori, valutazione dei residui e addestramento vincolato alle PDE.
• Prestazioni: Compilando in un singolo grafo XLA, jNO afferma di ridurre l'overhead di compilazione e migliorare la stabilità numerica rispetto alle alternative incentrate su PyTorch, mantenendo al contempo la semplicità di prototipazione.
• Interoperabilità: L'integrazione di FEAX consente flussi di lavoro agli elementi finiti differenziabili che rimangono compatibili con solver JAX (ad esempio, Diffrax, Optimistix) o librerie esterne (ad esempio, SciPy).

Significato

Gli autori posizionano jNO come un passo critico verso la riduzione della frammentazione nell'ecosistema SciML. Il suo significato risiede nel:

• Abbassare le Barriere: Rende i flussi di lavoro avanzati per operatori neurali e modelli fondamentali più accessibili ai ricercatori fornendo uno stack unificato e nativo JAX.
• Abilitare Flussi di Lavoro Ibridi: Facilita flussi di lavoro precedentemente difficili, come metodi di ensemble, trasferimento di apprendimento progressivo e vincoli multi-modello, standardizzando i controlli di addestramento e l'ottimizzazione a livello di parametro.
• Fondamento per la Ricerca Futura: Consolidando i modelli fondamentali in un unico framework, abilita confronti sistematici e flussi di lavoro compositivi che erano precedentemente ostacolati da repository disparati.

Il documento conclude che jNO fornisce un significativo contributo alla ricerca rendendo questi flussi di lavoro avanzati più accessibili ed efficienti all'interno dell'ecosistema JAX, fungendo da stack di software di ricerca riutilizzabile per la comunità.