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

jNO é uma biblioteca unificada e nativa do JAX que otimiza o treinamento de operadores neurais e modelos fundamentais ao integrar abordagens baseadas em dados e informadas pela física em um único sistema de rastreamento simbólico, permitindo transições contínuas entre regressão de operadores, avaliação de resíduos conscientes de malha e otimização com restrições de EDP sem reestruturação de código.

O essencial

Imagine que você está tentando ensinar um computador a entender as leis da física, como o calor flui através de uma placa de metal ou como a água gira em torno de uma pedra. No passado, fazer isso com inteligência artificial era como tentar construir uma casa onde o arquiteto, o encanador, o eletricista e o carpinteiro falavam idiomas diferentes e usavam plantas baixas distintas. Você tinha que escrever um conjunto de código para a forma do cômodo (geometria), outro para as equações matemáticas (física) e um terceiro para o processo real de aprendizado. Se você quisesse trocar de um tipo de matemática para outro, frequentemente precisava demolir a casa inteira e começar do zero.

jNO (jax Neural Operators) é uma nova ferramenta que atua como um "tradutor universal" e um "construtor mestre" em um só. É uma biblioteca de software projetada para tornar o treinamento desses modelos de IA inteligentes em física muito mais fácil, rápido e flexível, especificamente para a linguagem de programação JAX (uma ferramenta popular para computação científica de alta velocidade).

Veja como funciona, usando algumas analogias simples:

1. A Magia do "Único Script" (Sistema de Rastreamento)

Pense no jNO como um único script mágico que controla todo o canteiro de obras.

• Antes: Você tinha que escrever um script para a planta baixa, um script separado para a matemática e outro para as regras de aprendizado. Se você quisesse mudar a matemática, tinha que reescrever o script da planta baixa também.
• Com o jNO: Você escreve tudo em uma única linguagem. Você define a forma do cômodo, as equações da física e os objetivos de aprendizado tudo de uma vez. O software "rastreia" (ou registra) suas instruções como um diretor de cinema filmando uma cena. Posteriormente, ele compila esse filme em um programa super eficiente e de alta velocidade. Isso significa que você pode trocar entre diferentes tipos de problemas matemáticos ou adicionar novas regras de física sem reescrever seu código.

2. Os Modelos de Fundação "Lego"

Atualmente, existem muitos "modelos de fundação" (cérebros de IA pré-treinados) diferentes para física, mas eles são como conjuntos de Lego de fabricantes diferentes que não se encaixam. Uma marca usa blocos vermelhos, outra usa azuis, e eles não podem ser empilhados.

• Papel do jNO: Ele atua como um adaptador universal. Ele pega esses diferentes modelos de IA (como Poseidon, Walrus e Morph) e os traduz para que todos se encaixem no mesmo ecossistema JAX. Agora, um pesquisador pode pegar um "cérebro" pré-treinado, ajustá-lo ligeiramente e combiná-lo com suas próprias regras de física personalizadas, tudo sem precisar trocar de ferramentas de software.

3. A "Malha Inteligente" (Manipulação de Formas)

Ao simular física, os computadores precisam quebrar formas (como um tubo curvo ou um edifício complexo) em pequenas peças de grade chamadas "malha".

• A Inovação: O jNO possui um sistema de "malha inteligente" embutido. É como ter um robô que pode desenhar instantaneamente uma grade sobre qualquer forma que você descreva, seja um quadrado simples ou um objeto 3D complexo com buracos. Ele mantém o registro de qual parte da grade é o "interior", qual é a "parede" e qual é a "fronteira", para que a IA saiba exatamente onde aplicar as regras de física.

4. O Botão de "Ajuste Fino"

Às vezes, você quer pegar uma IA pré-treinada e ensinar a ela uma nova tarefa específica.

• O Painel de Controle: O jNO oferece um painel de controle muito detalhado. Você pode dizer à IA: "Congele estas partes do seu cérebro para que elas não mudem", ou "Aprenda apenas dessas conexões específicas", ou "Use uma velocidade de aprendizado específica". Você pode fazer isso para partes individuais do modelo sem precisar reconstruir tudo. É como ser capaz de ajustar o volume apenas dos tambores em uma música sem mudar o violão ou os vocais.

5. O Motor de "Duplo Modo" (FEM e PINNs)

O artigo destaca que o jNO pode lidar com duas maneiras diferentes de resolver problemas de física:

• Ponto a Ponto: Verificando a física em pontos específicos (como verificar a temperatura em locais específicos em um mapa).
• Forma Inteira (Elemento Finito): Observando a física como um fluxo contínuo sobre uma forma inteira (como calcular o estresse total em uma ponte).
• O Benefício: O jNO permite que você alterne entre esses dois modos usando o mesmo código. É como ter um carro que pode dirigir tanto em estradas de terra quanto em rodovias sem que você precise trocar o motor ou o volante.

Por Que Isso Importa?

O objetivo principal do jNO é interromper a "fragmentação" do software científico. Em vez de pesquisadores equilibrarem cinco ferramentas diferentes para treinar um único modelo de IA, o jNO traz tudo para um só lugar.

• Velocidade: Como utiliza recursos especiais de compilação do JAX, ele roda mais rápido em chips de computador modernos.
• Simplicidade: Você não precisa ser um arquiteto de software para alternar entre diferentes tipos de problemas de física.
• Reutilização: Uma vez que você escreve um programa em jNO, pode salvá-lo, compartilhá-lo e executá-lo novamente mais tarde, mesmo em computadores diferentes, com a confiança de que funcionará da mesma maneira.

Em resumo, o jNO está tentando fazer com que o mundo complexo da "aprendizagem de máquina científica" pareça tão simples e unificado quanto escrever uma única história coerente, em vez de costurar um colcha de retalhos de fragmentos de código diferentes.

Resumo técnico

Resumo Técnico: jNO – Uma Biblioteca JAX para Treinamento de Operadores Neurais e Modelos de Fundação

Declaração do Problema

O cenário atual do aprendizado de máquina científico (SciML), particularmente no que diz respeito a operadores neurais e Redes Neurais Informadas pela Física (PINNs), sofre de fragmentação significativa de software. Modelos de fundação de última geração (por exemplo, Poseidon, Walrus, PDEformer2) e ferramentas de aprendizado de operadores são predominantemente mantidos em repositórios separados dentro do ecossistema PyTorch. Essa separação cria barreiras para comparação sistemática, aprendizado por transferência e ajuste fino reproduzível entre diferentes famílias de modelos. Além disso, fluxos de trabalho práticos são frequentemente desconexos, exigindo que os usuários gerenciem abstrações separadas para geometria, construção de perda, definição de modelo e lógica de treinamento. Essa fragmentação é exacerbada ao tentar combinar aprendizado de operador supervisionado com resíduos informados pela física ou ao transferir backbones pré-treinados para novos cenários de Equações Diferenciais Parciais (PDE).

Metodologia

O jNO (jax Neural Operators) aborda esses desafios fornecendo uma biblioteca nativa JAX unificada que trata todo o fluxo de trabalho de treinamento como um único programa rastreado. A metodologia central baseia-se em uma Linguagem de Domínio Específico (DSL) simbólica centrada em rastreamento e execução adiada.

1. Design Focado em Rastreamento

Ao contrário de modelos de execução ansiosa (eager), o jNO constrói um grafo de computação adiada onde domínios, chamadas de modelo, resíduos, perdas supervisionadas e diagnósticos são escritos em uma única linguagem simbólica.

• Nós Simbólicos: O sistema utiliza nós Placeholder para representar variáveis, tensores, modelos e operadores diferenciais. Operações aritméticas e diferenciais sobrecarregam os primitivos padrão do JAX para construir árvores de expressão em vez de executar imediatamente.
• Pipeline Unificado: Os usuários definem domínios e variáveis, compõem modelos com expressões diferenciais e resolvem tudo por meio de uma única interface. Isso permite transições transparentes entre regressão de operador, avaliação de resíduos conscientes de malha e treinamento com restrições de PDE sem reestruturar o código.
• Otimização: Antes da execução, o sistema realiza eliminação de subexpressões comuns (CSE) na árvore rastreada, canonicando subárvores estruturalmente idênticas em um único programa XLA. Isso unifica objetivos, definições de operadores e restrições físicas.

2. Fluxo de Trabalho de Domínio e Malha

A classe domain atua como uma camada de organização de dados consciente de malha, fazendo a ponte entre definições geométricas e expressões rastreadas.

• Integração: Suporta geração de malha via PyGmsh/Gmsh e carregamento via meshio.
• Estrutura: Os dados são organizados em um pool de malha (para conjuntos de pontos rotulados) e um dicionário de contexto de tempo de execução. Suporta treinamento em lote onde cada elemento do lote pode representar uma parametrização de PDE ou campo de força diferente.
• Adaptatividade: O sistema inclui ganchos para gerenciamento adaptativo de pontos, permitindo que o treinamento atualize conjuntos de pontos sem alterar equações simbólicas.

3. Integração do Método de Elementos Finitos (FEM)

O jNO estende seu modelo de programação rastreada para montagem de forma fraca via o backend FEAX.

• DSL Unificada: Em vez de uma frente de FEM separada, o jNO usa a mesma camada simbólica para formas fracas baseadas em malha, PINNs variacionais e aprendizado de operadores.
• Montagem: Formas fracas simbólicas são reduzidas para representações neutras em relação ao backend. Dependendo do alvo, o sistema pode despachar para vpinn (PINN variacional), fem_system (estacionário linear), fem_residual (estacionário não linear) ou fem_time (transiente).
• Diferenciabilidade: A integração mantém operadores de elementos finitos montados dentro do ecossistema JAX, permitindo fluxos de trabalho diferenciáveis de ponta a ponta (montagem, avaliação de resíduos, construção de Jacobiano) compatíveis com transformações JAX.

4. Interface de Modelo e Modelos de Fundação

O jNO trata arquiteturas de operadores neurais como operações ordinárias dentro do programa rastreado.

• Modelos de Fundação: Suporta traduções nativas JAX de modelos de fundação de PDE (Poseidon, Walrus, Morph, etc.) mantidos no repositório externo foundax. Esses modelos podem ser encapsulados via jno.nn.wrap() e compostos com expressões rastreadas arbitrárias.
• Controle Granular: Os usuários podem configurar o modelo, otimizador e comportamentos de taxa de aprendizado no nível de parâmetro. Recursos incluem mascaramento seletivo de parâmetros, congelamento/descongelamento e suporte a LoRA (Adaptação de Baixo Rango) para modelos Equinox e Flax.

5. Execução e Treinamento

O módulo jno.core compila as restrições rastreadas em um único programa de treinamento compilado JIT.

• Desempenho: Aproveita o XLA para eliminação agressiva de subexpressões, otimização de memória e paralelismo multi-dispositivo (via uma malha de dispositivos de 2 eixos).
• Ajuste de Hiperparâmetros: A interface ArchSpace suporta busca em grade e otimização sem gradiente (via Nevergrad) para parâmetros de arquitetura e treinamento, integrados diretamente no loop de treinamento.
• Persistência: Experimentos podem ser serializados usando jno.save (com suporte a assinaturas RSA para integridade) para preservar estados do solucionador, objetos de domínio e wrappers de inferência compilados.

Contribuições Principais

1. Linguagem Rastreada Unificada: O jNO é o primeiro framework consolidado a expressar operadores neurais, resíduos informados pela física e ajuste fino de modelos de fundação dentro de uma única interface simbólica no JAX.
2. Consolidação do Ecossistema: Atua como um backend unificado para traduzir e integrar diversas famílias de modelos de fundação de PDE (por exemplo, Poseidon, Walrus) para o JAX, permitindo avaliação lado a lado e composição sem manter bases de código separadas.
3. Unificação de FEM e Operadores: Ao estender a DSL rastreada para montagem de forma fraca, o jNO permite que formas fracas baseadas em malha, PINNs variacionais e modelos de operadores sejam compostos sem alterar o modelo de programação.
4. Controle no Nível de Parâmetro: A biblioteca fornece controle granular sobre a dinâmica de treinamento, incluindo mascaramento seletivo e adaptação LoRA, diretamente dentro do fluxo simbólico.
5. Reprodutibilidade e Portabilidade: A biblioteca é empacotada para PyPI, inclui testes automatizados e suporta compartilhamento de artefatos assinados para fluxos de trabalho verificáveis.

Resultados e Alegações

O artigo não apresenta novos resultados de benchmark em conjuntos de dados de PDE específicos, mas foca nas contribuições arquitetônicas e de fluxo de trabalho:

• Eficiência do Fluxo de Trabalho: O artigo alega que o design rastreado elimina a necessidade de reestruturação de código ao alternar entre aprendizado de operador, avaliação de resíduos e treinamento com restrições de PDE.
• Desempenho: Ao compilar para um único grafo XLA, o jNO alega reduzir a sobrecarga de compilação e melhorar a estabilidade numérica em comparação com alternativas centradas no PyTorch, mantendo a simplicidade de prototipagem.
• Interoperabilidade: A integração do FEAX permite fluxos de trabalho de elementos finitos diferenciáveis que permanecem compatíveis com solucionadores JAX (por exemplo, Diffrax, Optimistix) ou bibliotecas externas (por exemplo, SciPy).

Significância

Os autores posicionam o jNO como um passo crítico para reduzir a fragmentação no ecossistema SciML. Sua significância reside em:

• Redução de Barreiras: Torna fluxos de trabalho avançados de operadores neurais e modelos de fundação mais acessíveis aos pesquisadores, fornecendo uma pilha nativa JAX unificada.
• Habilitação de Fluxos de Trabalho Híbridos: Facilita fluxos de trabalho anteriormente difíceis, como métodos de ensemble, aprendizado por transferência progressivo e restrições multi-modelo, padronizando controles de treinamento e otimização no nível de parâmetro.
• Fundação para Pesquisa Futura: Ao consolidar modelos de fundação em um único framework, permite comparação sistemática e fluxos de trabalho composicionais que anteriormente eram impedidos por repositórios díspares.

O artigo conclui que o jNO fornece uma significância de pesquisa credível ao tornar esses fluxos de trabalho avançados mais acessíveis e eficientes dentro do ecossistema JAX, servindo como uma pilha de software de pesquisa reutilizável para a comunidade.