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

jNO is een geïntegreerde, JAX-native bibliotheek die het trainen van neurale operatoren en fundamentele modellen stroomlijnt door datagedreven en op natuurkunde gebaseerde benaderingen te integreren in één systeem voor symbolische traceerbaarheid, waardoor naadloze overgangen mogelijk zijn tussen operatorregressie, mesh-bewuste residu-evaluatie en door partiële differentiaalvergelijkingen beperkte optimalisatie zonder code te herschikken.

De kern

Stel je voor dat je probeert een computer te leren de wetten van de natuurkunde te begrijpen, zoals hoe warmte door een metalen plaat stroomt of hoe water om een rots draait. In het verleden was het doen van dit soort dingen met kunstmatige intelligentie alsof je een huis probeerde te bouwen waarbij de architect, de loodgieter, de elektricien en de timmerman allemaal verschillende talen spraken en verschillende blauwdrukken gebruikten. Je moest één set code schrijven voor de vorm van de kamer (geometrie), een andere voor de wiskundige vergelijkingen (natuurkunde) en een derde voor het daadwerkelijke leerproces. Als je wilde overstappen van het ene type wiskunde naar het andere, moest je vaak het hele huis afbreken en opnieuw beginnen.

jNO (jax Neural Operators) is een nieuw hulpmiddel dat fungeert als een "universele vertaler" en een "hoofdbouwer" in één. Het is een softwarebibliotheek die is ontworpen om het trainen van deze natuurkundige AI-modellen veel eenvoudiger, sneller en flexibeler te maken, specifiek voor de JAX-programmeertaal (een populair hulpmiddel voor hoogwaardig wetenschappelijk rekenen).

Hier is hoe het werkt, met behulp van enkele eenvoudige analogieën:

1. De "Eén-Script"-magie (Traceersysteem)

Stel je jNO voor als een enkel, magisch script dat de hele bouwplaats beheert.

• Voorheen: Je moest een script schrijven voor de blauwdruk, een apart script voor de wiskunde en een ander voor de leerrules. Als je de wiskunde wilde veranderen, moest je ook het blauwdrukscript herschrijven.
• Met jNO: Je schrijft alles in één taal. Je definieert de vorm van de kamer, de natuurkundige vergelijkingen en de leerdoelen allemaal in één keer. De software "traceert" (of registreert) je instructies alsof een filmregisseur een scène filmde. Later compileert het deze film tot een super-efficiënt, hoogwaardig programma. Dit betekent dat je kunt wisselen tussen verschillende soorten wiskundige problemen of nieuwe natuurkunderegels kunt toevoegen zonder je code te herschrijven.

2. De "Lego"-fundamentmodellen

Op dit moment zijn er veel verschillende "fundamentmodellen" (voorgetrainde AI-hersenen) voor natuurkunde, maar ze zijn als Lego-setjes van verschillende fabrikanten die niet bij elkaar passen. Het ene merk gebruikt rode bakstenen, het andere blauwe, en ze kunnen niet op elkaar worden gestapeld.

• De rol van jNO: Het fungeert als een universele adapter. Het neemt deze verschillende AI-modellen (zoals Poseidon, Walrus en Morph) en vertaalt ze zodat ze allemaal in hetzelfde JAX-ecosysteem passen. Nu kan een onderzoeker een voorgetraind "brein" nemen, het lichtjes aanpassen en combineren met hun eigen aangepaste natuurkunderegels, allemaal zonder dat ze van softwaretools hoeven te wisselen.

3. Het "Slimme Net" (Het omgaan met vormen)

Bij het simuleren van natuurkunde moeten computers vormen (zoals een gebogen pijp of een complex gebouw) opsplitsen in tiny roosterstukjes die een "mesh" worden genoemd.

• De innovatie: jNO heeft een ingebouwd "slim net"-systeem. Het is alsof je een robot hebt die direct een rooster over elke vorm kan tekenen die je beschrijft, of het nu een simpele vierkant is of een complex 3D-voorwerp met gaten. Het houdt bij welk deel van het rooster de "binnenkant" is, welke de "wand" en welke de "grens", zodat de AI precies weet waar de natuurkunderegels moeten worden toegepast.

4. De "Fine-Tuning"-knop

Soms wil je een voorgetrainde AI nemen en hem een specifieke nieuwe taak leren.

• Het bedieningspaneel: jNO geeft je een zeer gedetailleerd bedieningspaneel. Je kunt de AI vertellen: "Vries deze delen van je brein in zodat ze niet veranderen", of "Leer alleen van deze specifieke verbindingen", of "Gebruik een specifieke leersnelheid". Je kunt dit doen voor individuele onderdelen van het model zonder het hele ding te moeten herbouwen. Het is alsof je het volume alleen op de drums in een nummer kunt aanpassen zonder de gitaar of zang te veranderen.

5. De "Dubbele-modus"-motor (FEM en PINNs)

Het artikel benadrukt dat jNO twee verschillende manieren kan hanteren om natuurkundeproblemen op te lossen:

• Punt-voor-punt: Het controleren van de natuurkunde op specifieke punten (zoals het controleren van de temperatuur op specifieke plekken op een kaart).
• Hele-vorm (Finite Element): Het bekijken van de natuurkunde als een continue stroom over een hele vorm (zoals het berekenen van de totale spanning op een brug).
• Het voordeel: jNO stelt je in staat om tussen deze twee modi te wisselen met dezelfde code. Het is alsof je een auto hebt die zowel op grindwegen als op snelwegen kan rijden zonder dat je de motor of het stuurwiel hoeft te veranderen.

Waarom is dit belangrijk?

Het hoofddoel van jNO is het stoppen van de "fragmentatie" van wetenschappelijke software. In plaats dat onderzoekers vijf verschillende tools moeten hanteren om één AI-model te trainen, brengt jNO alles op één plek samen.

• Snelheid: Omdat het gebruikmaakt van de speciale compilatiefuncties van JAX, draait het sneller op moderne computerchips.
• Eenvoud: Je hoeft geen software-architect te zijn om tussen verschillende soorten natuurkundeproblemen te wisselen.
• Opnieuw te gebruiken: Zodra je een programma in jNO hebt geschreven, kun je het opslaan, delen en later opnieuw uitvoeren, zelfs op verschillende computers, met het vertrouwen dat het op dezelfde manier zal werken.

Kortom, jNO probeert de complexe wereld van "wetenschappelijk machine learning" zo eenvoudig en verenigd te laten voelen als het schrijven van één samenhangend verhaal, in plaats van het naaien van een patchworkdeken van verschillende codefragmenten.

Technische samenvatting

Technische Samenvatting: jNO – Een JAX-bibliotheek voor het trainen van Neuronale Operator- en Fundamentele Modellen

Probleemstelling

Het huidige landschap van wetenschappelijk machine learning (SciML), met name wat betreft neuronale operatoren en Physics-Informed Neural Networks (PINNs), lijdt onder aanzienlijke softwarefragmentatie. State-of-the-art fundamentele modellen (bijv. Poseidon, Walrus, PDEformer2) en hulpmiddelen voor operatorlering worden voornamelijk onderhouden in afzonderlijke repositories binnen het PyTorch-ecosysteem. Deze scheiding creëert barrières voor systematische vergelijking, transfer learning en reproduceerbare fijne afstemming (fine-tuning) over verschillende modelfamilies heen. Bovendien zijn praktische workflows vaak losgekoppeld, waarbij gebruikers aparte abstracties moeten beheren voor geometrie, verliesconstructie, modeldefinitie en trainingslogica. Deze fragmentatie wordt verergerd bij pogingen om toezicht op operatorlering te combineren met physics-informed residuals of bij het overdragen van vooraf getrainde ruggegraten naar nieuwe settingen voor Partiële Differentiaalvergelijkingen (PDE's).

Methodologie

jNO (jax Neural Operators) adresseert deze uitdagingen door een geünificeerde, JAX-native bibliotheek te bieden die de volledige trainingsworkflow behandelt als één enkel getraceerd programma. De kernmethodologie steunt op een symbolische Domain-Specific Language (DSL) die is gecentreerd rondom trace en uitgestelde uitvoering.

1. Trace-First Ontwerp

In tegenstelling tot modellen met dringende uitvoering (eager execution), bouwt jNO een uitgestelde berekeningstabel op waarin domeinen, modeloproepen, residuals, toezichtverliezen en diagnostiek worden geschreven in één symbolische taal.

• Symbolische Knopen: Het systeem gebruikt Placeholder-knopen om variabelen, tensors, modellen en differentiaaloperatoren te representeren. Aritmetische en differentiaaloperaties overschrijven standaard JAX-primitieven om expressiebomen te bouwen in plaats van direct uit te voeren.
• Geünificeerde Pijplijn: Gebruikers definiëren domeinen en variabelen, stellen modellen samen met differentiaaluitdrukkingen en lossen alles op via één enkele interface. Dit zorgt voor naadloze overgangen tussen operatorregressie, mesh-bewuste residual-evaluatie en PDE-beperkte training zonder dat de code moet worden herstructureerd.
• Optimalisatie: Voor uitvoering voert het systeem gemeenschappelijke subexpressie-eliminatie (CSE) uit op de getraceerde boom, waarbij structureel identieke subbomen worden genormaliseerd naar één enkel XLA-programma. Dit verenigt doelstellingen, operatordefinities en fysica-beperkingen.

2. Domein- en Mesh-workflow

De domain-klasse fungeert als een mesh-bewiste data-organiseringslaag die geometrie-definities en getraceerde uitdrukkingen met elkaar verbindt.

• Integratie: Het ondersteunt mesh-generatie via PyGmsh/Gmsh en het laden via meshio.
• Structuur: Data is georganiseerd in een mesh-pool (voor gelabelde puntsets) en een runtime-contextwoordenboek. Het ondersteunt batch-training waarbij elk batch-element een andere PDE-parametrisatie of krachtfeld kan vertegenwoordigen.
• Adaptiviteit: Het systeem bevat haken voor adaptief puntbeheer, waardoor training puntsets kan bijwerken zonder symbolische vergelijkingen te wijzigen.

3. Integratie van de Finite Elementen Methode (FEM)

jNO breidt zijn getraceerde programmeermodel uit tot zwakke vormassemblage via de FEAX-backend.

• Geünificeerde DSL: In plaats van een afzonderlijk FEM-front-end, gebruikt jNO dezelfde symbolische laag voor mesh-gebaseerde zwakke vormen, variationale PINN's en operatorlering.
• Assemblage: Symbolische zwakke vormen worden verlaagd naar backend-neutrale representaties. Afhankelijk van het doel kan het systeem doorsturen naar vpinn (variational PINN), fem_system (lineair stationair), fem_residual (niet-lineair stationair) of fem_time (transiënt).
• Differentieerbaarheid: De integratie houdt geassembleerde finite element-operatoren binnen het JAX-ecosysteem, waardoor volledig differentieerbare workflows mogelijk zijn (assemblage, residual-evaluatie, Jacobiaan-construktie) die compatibel zijn met JAX-transformaties.

4. Modelinterface en Fundamentele Modellen

jNO behandelt architecturen voor neuronale operatoren als gewone operaties binnen het getraceerde programma.

• Fundamentele Modellen: Het ondersteunt JAX-native vertalingen van PDE-fundamentele modellen (Poseidon, Walrus, Morph, enz.) die worden onderhouden in de externe foundax-repository. Deze modellen kunnen worden gewikkeld via jno.nn.wrap() en worden samengesteld met willekeurige getraceerde uitdrukkingen.
• Gedetailleerde Controle: Gebruikers kunnen model-, optimalisator- en leersnelheidsgedrag op parameterniveau configureren. Functies omvatten selectieve parametermaskering, bevriezen/ontvriezen en ondersteuning voor LoRA (Low-Rank Adaptation) voor zowel Equinox- als Flax-modellen.

5. Uitvoering en Training

De jno.core-module compileert de getraceerde beperkingen tot één enkel JIT-gecompileerd trainingsprogramma.

• Prestaties: Het maakt gebruik van XLA voor agressieve subexpressie-eliminatie, geheugenoptimalisatie en multi-apparaat parallelisme (via een 2-assen apparaatnetwerk).
• Hyperparameter-tuning: De ArchSpace-interface ondersteunt grid search en optimalisatie zonder gradienten (via Nevergrad) voor architectuur- en trainingsparameters, direct geïntegreerd in de trainingslus.
• Persistentie: Experimenten kunnen worden gesequencerd met jno.save (met ondersteuning voor RSA-handtekeningen voor integriteit) om solver-statussen, domeinobjecten en gecompileerde inferentie-wikkelaars te behouden.

Belangrijkste Bijdragen

1. Geünificeerde Getraceerde Taal: jNO is het eerste geconsolideerde kader om neuronale operatoren, physics-informed residuals en fine-tuning van fundamentele modellen tot expressie te brengen binnen één enkele symbolische interface in JAX.
2. Ecosysteem-consolidatie: Het fungeert als een geünificeerde backend voor het vertalen en integreren van diverse PDE-fundamentele modelfamilies (bijv. Poseidon, Walrus) in JAX, waardoor zij-aan-zij evaluatie en compositie mogelijk zijn zonder afzonderlijke codebases te onderhouden.
3. Unificatie van FEM en Operator: Door de getraceerde DSL uit te breiden tot assemblage van zwakke vormen, stelt jNO mesh-gebaseerde zwakke vormen, variationale PINN's en operatormodellen in staat om te worden samengesteld zonder het programmeermodel te wijzigen.
4. Controle op Parameterniveau: De bibliotheek biedt gedetailleerde controle over trainingsdynamiek, inclusief selectieve maskering en LoRA-adaptatie, direct binnen de symbolische flow.
5. Reproduceerbaarheid en Portabiliteit: De bibliotheek is verpakt voor PyPI, bevat geautomatiseerde tests en ondersteunt het delen van gesigneerde artefacten voor verifieerbare workflows.

Resultaten en Beweringen

Het artikel presenteert geen nieuwe benchmarkresultaten op specifieke PDE-datasets, maar richt zich op de architecturale en workflow-bijdragen:

• Workflow-efficiëntie: Het artikel beweert dat het getraceerde ontwerp de noodzaak voor code-herstructurering elimineert bij het wisselen tussen operatorlering, residual-evaluatie en PDE-beperkte training.
• Prestaties: Door te compileren naar één enkel XLA-grafiek, beweert jNO de compilatie-overhead te verminderen en de numerieke stabiliteit te verbeteren ten opzichte van PyTorch-gecentreerde alternatieven, terwijl de eenvoud van prototyping behouden blijft.
• Interoperabiliteit: De integratie van FEAX maakt differentieerbare finite element-workflows mogelijk die compatibel blijven met JAX-solvers (bijv. Diffrax, Optimistix) of externe bibliotheken (bijv. SciPy).

Betekenis

De auteurs positioneren jNO als een cruciale stap richting het verminderen van fragmentatie in het SciML-ecosysteem. De betekenis ligt in:

• Verlagen van Barrières: Het maakt geavanceerde workflows voor neuronale operatoren en fundamentele modellen toegankelijker voor onderzoekers door een geünificeerde, JAX-native stack te bieden.
• Aanbieden van Hybride Workflows: Het faciliteert workflows die eerder moeilijk waren, zoals ensemble-methoden, progressieve transfer learning en multi-model-beperkingen, door trainingscontroles en optimalisatie op parameterniveau te standaardiseren.
• Fundament voor Toekomstig Onderzoek: Door fundamentele modellen te consolideren in één enkel kader, maakt het systematische vergelijking en compositie-workflows mogelijk die eerder werden belemmerd door uiteenlopende repositories.

Het artikel concludeert dat jNO een geloofwaardige onderzoeksrelevantie biedt door deze geavanceerde workflows binnen het JAX-ecosysteem toegankelijker en efficiënter te maken, en dient als een herbruikbare onderzoekssoftwarestack voor de gemeenschap.