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

jNO est une bibliothèque unifiée et native JAX qui rationalise l'entraînement des opérateurs neuronaux et des modèles de fondation en intégrant des approches basées sur les données et informées par la physique dans un seul système de traçage symbolique, permettant des transitions fluides entre la régression d'opérateurs, l'évaluation résiduelle consciente du maillage et l'optimisation contrainte par les EDP sans restructuration du code.

L'essentiel

Imaginez que vous essayez d'enseigner à un ordinateur les lois de la physique, comme la façon dont la chaleur traverse une plaque de métal ou comment l'eau tourbillonne autour d'un rocher. Autrefois, faire cela avec l'intelligence artificielle revenait à essayer de construire une maison où l'architecte, le plombier, l'électricien et le menuisier parlaient tous des langues différentes et utilisaient des plans distincts. Vous deviez écrire un ensemble de code pour la forme de la pièce (géométrie), un autre pour les équations mathématiques (physique) et un troisième pour le processus d'apprentissage lui-même. Si vous vouliez passer d'un type de mathématiques à un autre, il fallait souvent démolir toute la maison et recommencer à zéro.

jNO (jax Neural Operators) est un nouvel outil qui agit comme un « traducteur universel » et un « maître d'œuvre » réunis en un seul. C'est une bibliothèque logicielle conçue pour rendre l'entraînement de ces modèles d'IA intelligents en physique beaucoup plus facile, rapide et flexible, spécifiquement pour le langage de programmation JAX (un outil populaire pour le calcul scientifique haute performance).

Voici comment cela fonctionne, en utilisant quelques analogies simples :

1. La magie du « Script Unique » (Système de traçage)

Considérez jNO comme un script unique et magique qui contrôle l'ensemble du chantier de construction.

• Avant : Vous deviez écrire un script pour le plan, un script séparé pour les mathématiques et un autre pour les règles d'apprentissage. Si vous vouliez changer les mathématiques, vous deviez réécrire le script du plan également.
• Avec jNO : Vous écrivez tout dans un seul langage. Vous définissez la forme de la pièce, les équations de la physique et les objectifs d'apprentissage en une seule fois. Le logiciel « trace » (ou enregistre) vos instructions comme un réalisateur de cinéma filmant une scène. Plus tard, il compile ce film en un programme ultra-efficace et haute vitesse. Cela signifie que vous pouvez passer d'un type de problème mathématique à un autre ou ajouter de nouvelles règles de physique sans réécrire votre code.

2. Les modèles de base « Lego »

Actuellement, il existe de nombreux « modèles de base » (cerveaux d'IA pré-entraînés) pour la physique, mais ils sont comme des sets de Lego de différents fabricants qui ne s'assemblent pas. Une marque utilise des briques rouges, une autre des briques bleues, et elles ne peuvent pas être empilées.

• Le rôle de jNO : Il agit comme un adaptateur universel. Il prend ces différents modèles d'IA (comme Poseidon, Walrus et Morph) et les traduit afin qu'ils s'intègrent tous dans le même écosystème JAX. Désormais, un chercheur peut prendre un « cerveau » pré-entraîné, l'ajuster légèrement et le combiner avec ses propres règles de physique personnalisées, le tout sans avoir besoin de changer d'outils logiciels.

3. Le « Maillage Intelligent » (Gestion des formes)

Lors de la simulation de phénomènes physiques, les ordinateurs doivent décomposer les formes (comme un tuyau courbé ou un bâtiment complexe) en de minuscules pièces de grille appelées « maillage ».

• L'innovation : jNO dispose d'un système de « maillage intelligent » intégré. C'est comme avoir un robot capable de dessiner instantanément une grille sur n'importe quelle forme que vous décrivez, qu'il s'agisse d'un carré simple ou d'un objet 3D complexe avec des trous. Il garde une trace de quelle partie de la grille est « l'intérieur », quelle partie est le « mur » et quelle partie est la « frontière », afin que l'IA sache exactement où appliquer les règles de physique.

4. Le bouton de « Réglage Fin »

Parfois, vous souhaitez prendre une IA pré-entraînée et lui apprendre une nouvelle tâche spécifique.

• Le Panneau de contrôle : jNO vous offre un panneau de contrôle très détaillé. Vous pouvez dire à l'IA : « Gèle ces parties de ton cerveau pour qu'elles ne changent pas », ou « Apprends uniquement à partir de ces connexions spécifiques », ou « Utilise une vitesse d'apprentissage particulière ». Vous pouvez le faire pour des parties individuelles du modèle sans avoir à reconstruire l'ensemble. C'est comme pouvoir régler le volume uniquement sur la batterie d'une chanson sans modifier la guitare ni les voix.

5. Le moteur « Double Mode » (FEM et PINNs)

L'article souligne que jNO peut gérer deux manières différentes de résoudre des problèmes de physique :

• Point par point : Vérifier la physique à des points spécifiques (comme vérifier la température à des endroits précis sur une carte).
• Forme entière (Éléments finis) : Observer la physique comme un flux continu sur une forme entière (comme calculer la contrainte totale sur un pont).
• L'avantage : jNO vous permet de basculer entre ces deux modes en utilisant le même code. C'est comme avoir une voiture capable de rouler à la fois sur des routes de terre et sur des autoroutes sans que vous ayez besoin de changer le moteur ou le volant.

Pourquoi cela importe-t-il ?

L'objectif principal de jNO est de mettre fin à la « fragmentation » des logiciels scientifiques. Au lieu que les chercheurs jonglent avec cinq outils différents pour entraîner un seul modèle d'IA, jNO rassemble tout en un seul endroit.

• Vitesse : Parce qu'il utilise les fonctionnalités de compilation spéciales de JAX, il s'exécute plus rapidement sur les puces informatiques modernes.
• Simplicité : Vous n'avez pas besoin d'être architecte logiciel pour passer d'un type de problème physique à un autre.
• Réutilisabilité : Une fois que vous avez écrit un programme en jNO, vous pouvez l'enregistrer, le partager et le relancer plus tard, même sur des ordinateurs différents, avec la certitude qu'il fonctionnera de la même manière.

En résumé, jNO tente de rendre le monde complexe de « l'apprentissage automatique scientifique » aussi simple et unifié que d'écrire une histoire unique et cohérente, plutôt que de assembler un patchwork de différents fragments de code.

Résumé technique

Résumé technique : jNO – Une bibliothèque JAX pour l'entraînement d'opérateurs neuronaux et de modèles de fondation

Énoncé du problème

Le paysage actuel de l'apprentissage automatique scientifique (SciML), en particulier concernant les opérateurs neuronaux et les réseaux de neurones informés par la physique (PINNs), souffre d'une fragmentation logicielle significative. Les modèles de fondation les plus avancés (par exemple, Poseidon, Walrus, PDEformer2) et les outils d'apprentissage d'opérateurs sont principalement maintenus dans des dépôts séparés au sein de l'écosystème PyTorch. Cette séparation crée des obstacles à la comparaison systématique, à l'apprentissage par transfert et au réglage fin reproductible entre différentes familles de modèles. De plus, les flux de travail pratiques sont souvent disjoints, obligeant les utilisateurs à gérer des abstractions séparées pour la géométrie, la construction de la fonction de perte, la définition du modèle et la logique d'entraînement. Cette fragmentation est exacerbée lors de la tentative de combiner l'apprentissage supervisé d'opérateurs avec des résidus informés par la physique ou lors du transfert de backbones préentraînés vers de nouveaux contextes d'équations aux dérivées partielles (EDP).

Méthodologie

jNO (jax Neural Operators) répond à ces défis en fournissant une bibliothèque unifiée et native JAX qui traite l'ensemble du flux d'entraînement comme un seul programme tracé. La méthodologie centrale repose sur un langage spécifique à un domaine (DSL) symbolique centré sur le tracage et l'exécution différée.

1. Conception axée sur le tracage

Contrairement aux modèles d'exécution immédiate (eager execution), jNO construit un graphe de calcul différé où les domaines, les appels de modèle, les résidus, les pertes supervisées et les diagnostics sont écrits dans un seul langage symbolique.

• Nœuds symboliques : Le système utilise des nœuds Placeholder pour représenter les variables, les tenseurs, les modèles et les opérateurs différentiels. Les opérations arithmétiques et différentielles surchargent les primitives JAX standard pour construire des arbres d'expressions plutôt que d'exécuter immédiatement.
• Pipeline unifié : Les utilisateurs définissent les domaines et les variables, composent des modèles avec des expressions différentielles et résolvent tout via une interface unique. Cela permet des transitions transparentes entre la régression d'opérateurs, l'évaluation de résidus sensibles au maillage et l'entraînement contraint par les EDP sans restructurer le code.
• Optimisation : Avant l'exécution, le système effectue une élimination des sous-expressions communes (CSE) sur l'arbre tracé, canalisant les sous-arbres structurellement identiques vers un seul programme XLA. Cela unifie les objectifs, les définitions d'opérateurs et les contraintes physiques.

2. Flux de travail Domaine et Maillage

La classe domain agit comme une couche d'organisation des données sensible au maillage, faisant le lien entre les définitions géométriques et les expressions tracées.

• Intégration : Elle prend en charge la génération de maillages via PyGmsh/Gmsh et le chargement via meshio.
• Structure : Les données sont organisées en un pool de maillages (pour des ensembles de points étiquetés) et un dictionnaire de contexte d'exécution. Elle prend en charge l'entraînement par lots où chaque élément de lot peut représenter une paramétrisation d'EDP ou un champ de force différent.
• Adaptativité : Le système inclut des crochets pour la gestion adaptative des points, permettant à l'entraînement de mettre à jour les ensembles de points sans modifier les équations symboliques.

3. Intégration de la Méthode des Éléments Finis (MEF)

jNO étend son modèle de programmation tracée à l'assemblage de la forme faible via le backend FEAX.

• DSL unifié : Au lieu d'une interface MEF séparée, jNO utilise la même couche symbolique pour les formes faibles basées sur le maillage, les PINNs variationnelles et l'apprentissage d'opérateurs.
• Assemblage : Les formes faibles symboliques sont abaissées vers des représentations neutres vis-à-vis du backend. Selon la cible, le système peut dispatcher vers vpinn (PINN variationnelle), fem_system (linéaire stationnaire), fem_residual (non linéaire stationnaire) ou fem_time (transitoire).
• Différentiabilité : L'intégration maintient les opérateurs d'éléments finis assemblés dans l'écosystème JAX, permettant des flux de travail entièrement différentiables (assemblage, évaluation de résidus, construction de Jacobienne) compatibles avec les transformations JAX.

4. Interface de modèle et modèles de fondation

jNO traite les architectures d'opérateurs neuronaux comme des opérations ordinaires au sein du programme tracé.

• Modèles de fondation : Il prend en charge les traductions natives JAX des modèles de fondation d'EDP (Poseidon, Walrus, Morph, etc.) maintenus dans le dépôt externe foundax. Ces modèles peuvent être encapsulés via jno.nn.wrap() et composés avec des expressions tracées arbitraires.
• Contrôle granulaire : Les utilisateurs peuvent configurer le comportement du modèle, de l'optimiseur et du taux d'apprentissage au niveau des paramètres. Les fonctionnalités incluent le masquage sélectif des paramètres, le gel/dégel, et le support de LoRA (Low-Rank Adaptation) pour les modèles Equinox et Flax.

5. Exécution et entraînement

Le module jno.core compile les contraintes tracées en un seul programme d'entraînement compilé JIT.

• Performance : Il exploite XLA pour une élimination agressive des sous-expressions, une optimisation de la mémoire et un parallélisme multi-périphériques (via une grille de périphériques à 2 axes).
• Réglage des hyperparamètres : L'interface ArchSpace prend en charge la recherche par grille et l'optimisation sans gradient (via Nevergrad) pour les paramètres d'architecture et d'entraînement, intégrés directement dans la boucle d'entraînement.
• Persistance : Les expériences peuvent être sérialisées à l'aide de jno.save (prenant en charge les signatures RSA pour l'intégrité) afin de préserver les états du solveur, les objets de domaine et les wrappers d'inférence compilés.

Contributions clés

1. Langage tracé unifié : jNO est le premier cadre consolidé permettant d'exprimer les opérateurs neuronaux, les résidus informés par la physique et le réglage fin de modèles de fondation au sein d'une interface symbolique unique dans JAX.
2. Consolidation de l'écosystème : Il agit comme un backend unifié pour traduire et intégrer diverses familles de modèles de fondation d'EDP (par exemple, Poseidon, Walrus) dans JAX, permettant une évaluation côte à côte et une composition sans maintenir des bases de code séparées.
3. Unification MEF et Opérateurs : En étendant le DSL tracé à l'assemblage de forme faible, jNO permet de composer des formes faibles basées sur le maillage, des PINNs variationnelles et des modèles d'opérateurs sans changer le modèle de programmation.
4. Contrôle au niveau des paramètres : La bibliothèque offre un contrôle granulaire sur la dynamique d'entraînement, y compris le masquage sélectif et l'adaptation LoRA, directement au sein du flux symbolique.
5. Reproductibilité et portabilité : La bibliothèque est emballée pour PyPI, inclut des tests automatisés et prend en charge le partage d'artefacts signés pour des flux de travail vérifiables.

Résultats et affirmations

L'article ne présente pas de nouveaux résultats de référence sur des jeux de données d'EDP spécifiques, mais se concentre sur les contributions architecturales et de flux de travail :

• Efficacité du flux de travail : L'article affirme que la conception tracée élimine le besoin de restructuration du code lors du passage de l'apprentissage d'opérateurs à l'évaluation de résidus et à l'entraînement contraint par les EDP.
• Performance : En compilant vers un seul graphe XLA, jNO affirme réduire la surcharge de compilation et améliorer la stabilité numérique par rapport aux alternatives centrées sur PyTorch, tout en maintenant la simplicité du prototypage.
• Interopérabilité : L'intégration de FEAX permet des flux de travail d'éléments finis différentiables qui restent compatibles avec les solveurs JAX (par exemple, Diffrax, Optimistix) ou des bibliothèques externes (par exemple, SciPy).

Importance

Les auteurs positionnent jNO comme une étape cruciale vers la réduction de la fragmentation dans l'écosystème SciML. Son importance réside dans :

• Réduction des barrières : Il rend les flux de travail avancés d'opérateurs neuronaux et de modèles de fondation plus accessibles aux chercheurs en fournissant une pile native JAX unifiée.
• Permettre des flux de travail hybrides : Il facilite des flux de travail auparavant difficiles, tels que les méthodes d'ensemble, l'apprentissage par transfert progressif et les contraintes multi-modèles, en standardisant les contrôles d'entraînement et l'optimisation au niveau des paramètres.
• Fondement pour la recherche future : En consolidant les modèles de fondation dans un cadre unique, il permet une comparaison systématique et des flux de travail compositionnels qui étaient auparavant entravés par des dépôts disparates.

L'article conclut que jNO offre une pertinence de recherche crédible en rendant ces flux de travail avancés plus accessibles et efficaces au sein de l'écosystème JAX, servant de pile logicielle de recherche réutilisable pour la communauté.