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

jNO es una biblioteca unificada y nativa de JAX que agiliza el entrenamiento de operadores neuronales y modelos fundamentales al integrar enfoques basados en datos e informados por la física en un único sistema de trazado simbólico, permitiendo transiciones fluidas entre la regresión de operadores, la evaluación de residuos conscientes de la malla y la optimización con restricciones de EDP sin reestructurar el código.

Lo esencial

Imagina que estás intentando enseñar a una computadora a entender las leyes de la física, como cómo fluye el calor a través de una placa metálica o cómo el agua gira alrededor de una roca. En el pasado, hacer esto con inteligencia artificial era como intentar construir una casa donde el arquitecto, el fontanero, el electricista y el carpintero hablaban idiomas diferentes y usaban planos distintos. Tenías que escribir un conjunto de código para la forma de la habitación (geometría), otro para las ecuaciones matemáticas (física) y un tercero para el proceso de aprendizaje real. Si querías cambiar de un tipo de matemáticas a otro, a menudo tenías que derribar toda la casa y empezar de nuevo.

jNO (jax Neural Operators) es una nueva herramienta que actúa como un "traductor universal" y un "maestro constructor" combinados en uno. Es una biblioteca de software diseñada para hacer que el entrenamiento de estos modelos de IA inteligentes en física sea mucho más fácil, rápido y flexible, específicamente para el lenguaje de programación JAX (una herramienta popular para la computación científica de alta velocidad).

Así es como funciona, usando algunas analogías simples:

1. La Magia del "Un Solo Script" (Sistema de Rastreo)

Piensa en jNO como un único script mágico que controla toda la obra de construcción.

• Antes: Tenías que escribir un script para el plano, un script separado para las matemáticas y otro para las reglas de aprendizaje. Si querías cambiar las matemáticas, también tenías que reescribir el script del plano.
• Con jNO: Escribes todo en un solo lenguaje. Defines la forma de la habitación, las ecuaciones de física y los objetivos de aprendizaje todo de una vez. El software "rastrea" (o registra) tus instrucciones como un director de cine filmando una escena. Más tarde, compila esta película en un programa de alta velocidad y súper eficiente. Esto significa que puedes cambiar entre diferentes tipos de problemas matemáticos o agregar nuevas reglas de física sin reescribir tu código.

2. Los Modelos Base de "Lego"

Actualmente, existen muchos "modelos base" (cerebros de IA preentrenados) diferentes para la física, pero son como conjuntos de Lego de diferentes fabricantes que no encajan entre sí. Una marca usa ladrillos rojos, otra usa azules, y no se pueden apilar.

• El Rol de jNO: Actúa como un adaptador universal. Toma estos diferentes modelos de IA (como Poseidon, Walrus y Morph) y los traduce para que todos encajen en el mismo ecosistema JAX. Ahora, un investigador puede tomar un "cerebro" preentrenado, ajustarlo ligeramente y combinarlo con sus propias reglas de física personalizadas, todo sin necesidad de cambiar de herramientas de software.

3. La "Malla Inteligente" (Manejo de Formas)

Al simular física, las computadoras necesitan dividir formas (como una tubería curva o un edificio complejo) en pequeñas piezas de cuadrícula llamadas "malla".

• La Innovación: jNO tiene un sistema de "malla inteligente" integrado. Es como tener un robot que puede dibujar instantáneamente una cuadrícula sobre cualquier forma que describas, ya sea un cuadrado simple o un objeto 3D complejo con agujeros. Mantiene un registro de qué parte de la cuadrícula es el "interior", cuál es la "pared" y cuál es el "límite", para que la IA sepa exactamente dónde aplicar las reglas de física.

4. El Dial de "Ajuste Fino"

A veces quieres tomar una IA preentrenada y enseñarle una tarea nueva específica.

• El Panel de Control: jNO te ofrece un panel de control muy detallado. Puedes decirle a la IA: "Congela estas partes de tu cerebro para que no cambien", o "Aprende solo de estas conexiones específicas", o "Usa una velocidad de aprendizaje específica". Puedes hacer esto para partes individuales del modelo sin tener que reconstruir todo. Es como poder ajustar el volumen solo de la batería en una canción sin cambiar la guitarra ni los vocales.

5. El Motor de "Doble Modo" (FEM y PINNs)

El artículo destaca que jNO puede manejar dos formas diferentes de resolver problemas de física:

• Punto por Punto: Verificar la física en puntos específicos (como verificar la temperatura en lugares específicos de un mapa).
• Forma Completa (Elemento Finito): Observar la física como un flujo continuo sobre toda una forma (como calcular el estrés total en un puente).
• El Beneficio: jNO te permite cambiar entre estos dos modos usando el mismo código. Es como tener un coche que puede conducir tanto por caminos de tierra como por autopistas sin que necesites cambiar el motor ni el volante.

¿Por Qué Importa Esto?

El objetivo principal de jNO es detener la "fragmentación" del software científico. En lugar de que los investigadores manejen cinco herramientas diferentes para entrenar un modelo de IA, jNO lo reúne todo en un solo lugar.

• Velocidad: Porque utiliza las características especiales de compilación de JAX, se ejecuta más rápido en los chips de computadora modernos.
• Simplicidad: No necesitas ser un arquitecto de software para cambiar entre diferentes tipos de problemas de física.
• Reutilización: Una vez que escribes un programa en jNO, puedes guardarlo, compartirlo y ejecutarlo de nuevo más tarde, incluso en computadoras diferentes, con la confianza de que funcionará de la misma manera.

En resumen, jNO está tratando de hacer que el mundo complejo del "aprendizaje automático científico" se sienta tan simple y unificado como escribir una sola historia coherente, en lugar de unir un patchwork de diferentes fragmentos de código.

Resumen técnico

Resumen Técnico: jNO – Una Biblioteca JAX para Entrenamiento de Operadores Neuronales y Modelos Fundacionales

Declaración del Problema

El panorama actual del aprendizaje automático científico (SciML), particularmente en lo que respecta a operadores neuronales y Redes Neuronales Informadas por la Física (PINNs), sufre de una fragmentación de software significativa. Los modelos fundacionales más avanzados (por ejemplo, Poseidon, Walrus, PDEformer2) y las herramientas de aprendizaje de operadores se mantienen predominantemente en repositorios separados dentro del ecosistema de PyTorch. Esta separación crea barreras para la comparación sistemática, el aprendizaje por transferencia y el ajuste fino reproducible entre diferentes familias de modelos. Además, los flujos de trabajo prácticos suelen estar desconectados, lo que obliga a los usuarios a gestionar abstracciones separadas para la geometría, la construcción de la función de pérdida, la definición del modelo y la lógica de entrenamiento. Esta fragmentación se ve exacerbada al intentar combinar el aprendizaje de operadores supervisado con residuos informados por la física o al transferir espaldas preentrenadas a nuevos entornos de Ecuaciones Diferenciales Parciales (EDP).

Metodología

jNO (jax Neural Operators) aborda estos desafíos proporcionando una biblioteca unificada y nativa de JAX que trata todo el flujo de trabajo de entrenamiento como un único programa rastreado. La metodología central se basa en un Lenguaje Específico de Dominio (DSL) simbólico centrado en el rastreo y la ejecución diferida.

1. Diseño Basado en Rastreo Primero

A diferencia de los modelos de ejecución ansiosa, jNO construye un gráfico de cómputo diferido donde los dominios, las llamadas a modelos, los residuos, las pérdidas supervisadas y los diagnósticos se escriben en un único lenguaje simbólico.

• Nodos Simbólicos: El sistema utiliza nodos Placeholder para representar variables, tensores, modelos y operadores diferenciales. Las operaciones aritméticas y diferenciales sobrecargan los primitivos estándar de JAX para construir árboles de expresión en lugar de ejecutarlos inmediatamente.
• Pipeline Unificado: Los usuarios definen dominios y variables, componen modelos con expresiones diferenciales y resuelven todo a través de una única interfaz. Esto permite transiciones fluidas entre la regresión de operadores, la evaluación de residuos conscientes de la malla y el entrenamiento restringido por EDP sin reestructurar el código.
• Optimización: Antes de la ejecución, el sistema realiza la eliminación de subexpresiones comunes (CSE) en el árbol rastreado, canonicizando subárboles estructuralmente idénticos a un único programa XLA. Esto unifica objetivos, definiciones de operadores y restricciones físicas.

2. Flujo de Trabajo de Dominio y Malla

La clase domain actúa como una capa de organización de datos consciente de la malla, puenteando las definiciones geométricas y las expresiones rastreadas.

• Integración: Soporta la generación de mallas a través de PyGmsh/Gmsh y la carga mediante meshio.
• Estructura: Los datos se organizan en un grupo de mallas (para conjuntos de puntos etiquetados) y un diccionario de contexto de tiempo de ejecución. Soporta el entrenamiento por lotes donde cada elemento del lote puede representar una parametrización de EDP o un campo de forzamiento diferente.
• Adaptabilidad: El sistema incluye ganchos para la gestión adaptativa de puntos, permitiendo que el entrenamiento actualice los conjuntos de puntos sin alterar las ecuaciones simbólicas.

3. Integración del Método de Elementos Finitos (MEF)

jNO extiende su modelo de programación rastreada al ensamblaje de formas débiles a través del backend FEAX.

• DSL Unificado: En lugar de un front-end de MEF separado, jNO utiliza la misma capa simbólica para formas débiles basadas en malla, PINNs variacionales y aprendizaje de operadores.
• Ensamblaje: Las formas débiles simbólicas se reducen a representaciones neutrales del backend. Dependiendo del objetivo, el sistema puede despachar a vpinn (PINN variacional), fem_system (estacionario lineal), fem_residual (estacionario no lineal) o fem_time (transitorio).
• Diferenciabilidad: La integración mantiene los operadores de elementos finitos ensamblados dentro del ecosistema de JAX, habilitando flujos de trabajo diferenciables de extremo a extremo (ensamblaje, evaluación de residuos, construcción de Jacobianos) compatibles con las transformaciones de JAX.

4. Interfaz de Modelo y Modelos Fundacionales

jNO trata las arquitecturas de operadores neuronales como operaciones ordinarias dentro del programa rastreado.

• Modelos Fundacionales: Soporta traducciones nativas de JAX de modelos fundacionales de EDP (Poseidon, Walrus, Morph, etc.) mantenidos en el repositorio externo foundax. Estos modelos pueden envolverse mediante jno.nn.wrap() y componerse con expresiones rastreadas arbitrarias.
• Control Fino: Los usuarios pueden configurar el modelo, el optimizador y los comportamientos de la tasa de aprendizaje a nivel de parámetro. Las características incluyen enmascaramiento selectivo de parámetros, congelación/descongelación y soporte de LoRA (Adaptación de Bajo Rango) tanto para modelos Equinox como Flax.

5. Ejecución y Entrenamiento

El módulo jno.core compila las restricciones rastreadas en un único programa de entrenamiento compilado JIT.

• Rendimiento: Aprovecha XLA para la eliminación agresiva de subexpresiones, la optimización de memoria y el paralelismo multi-dispositivo (a través de una malla de dispositivos de 2 ejes).
• Ajuste de Hiperparámetros: La interfaz ArchSpace soporta búsqueda en cuadrícula y optimización sin gradiente (mediante Nevergrad) para parámetros de arquitectura y entrenamiento, integrada directamente en el bucle de entrenamiento.
• Persistencia: Los experimentos pueden serializarse utilizando jno.save (soportando firmas RSA para integridad) para preservar estados del solucionador, objetos de dominio y envoltorios de inferencia compilados.

Contribuciones Clave

1. Lenguaje Rastreado Unificado: jNO es el primer marco consolidado para expresar operadores neuronales, residuos informados por la física y ajuste fino de modelos fundacionales dentro de una única interfaz simbólica en JAX.
2. Consolidación del Ecosistema: Actúa como un backend unificado para traducir e integrar diversas familias de modelos fundacionales de EDP (por ejemplo, Poseidon, Walrus) en JAX, habilitando la evaluación y composición lado a lado sin mantener bases de código separadas.
3. Unificación de MEF y Operadores: Al extender el DSL rastreado al ensamblaje de formas débiles, jNO permite componer formas débiles basadas en malla, PINNs variacionales y modelos de operadores sin cambiar el modelo de programación.
4. Control a Nivel de Parámetro: La biblioteca proporciona un control fino sobre la dinámica de entrenamiento, incluido el enmascaramiento selectivo y la adaptación LoRA, directamente dentro del flujo simbólico.
5. Reproducibilidad y Portabilidad: La biblioteca está empaquetada para PyPI, incluye pruebas automatizadas y soporta el intercambio de artefactos firmados para flujos de trabajo verificables.

Resultados y Afirmaciones

El artículo no presenta nuevos resultados de referencia en conjuntos de datos de EDP específicos, sino que se centra en las contribuciones arquitectónicas y de flujo de trabajo:

• Eficiencia del Flujo de Trabajo: El artículo afirma que el diseño rastreado elimina la necesidad de reestructurar el código al cambiar entre el aprendizaje de operadores, la evaluación de residuos y el entrenamiento restringido por EDP.
• Rendimiento: Al compilar a un único gráfico XLA, jNO afirma reducir la sobrecarga de compilación y mejorar la estabilidad numérica en comparación con las alternativas centradas en PyTorch, manteniendo al mismo tiempo la simplicidad de la prototipación.
• Interoperabilidad: La integración de FEAX permite flujos de trabajo de elementos finitos diferenciables que permanecen compatibles con solucionadores de JAX (por ejemplo, Diffrax, Optimistix) o bibliotecas externas (por ejemplo, SciPy).

Significancia

Los autores posicionan a jNO como un paso crítico hacia la reducción de la fragmentación en el ecosistema SciML. Su significancia radica en:

• Reducción de Barreras: Hace que los flujos de trabajo avanzados de operadores neuronales y modelos fundacionales sean más accesibles para los investigadores al proporcionar una pila unificada y nativa de JAX.
• Habilitación de Flujos de Trabajo Híbridos: Facilita flujos de trabajo previamente difíciles, como métodos de conjunto, aprendizaje por transferencia progresivo y restricciones multi-modelo, al estandarizar los controles de entrenamiento y la optimización a nivel de parámetro.
• Fundamento para la Investigación Futura: Al consolidar los modelos fundacionales en un único marco, habilita la comparación sistemática y los flujos de trabajo composicionales que anteriormente se veían impedidos por repositorios dispares.

El artículo concluye que jNO proporciona una significancia de investigación creíble al hacer que estos flujos de trabajo avanzados sean más accesibles y eficientes dentro del ecosistema de JAX, sirviendo como una pila de software de investigación reutilizable para la comunidad.