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

jNO は、データ駆動型アプローチと物理情報型アプローチを単一の記号的トレーシングシステムに統合することで、ニューラル演算子および基盤モデルの学習を効率化する統合された JAX ネイティブライブラリであり、コードの再構成なしに演算子回帰、メッシュ対応残差評価、PDE 制約最適化間のシームレスな移行を可能にします。

原著者: Leon Armbruster, Rathan Ramesh, Georg Kruse, Christopher Straub

公開日 2026-05-12
📖 1 分で読めます☕ さくっと読める

原著者: Leon Armbruster, Rathan Ramesh, Georg Kruse, Christopher Straub

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。 免責事項の全文を読む

あなたが、金属板を伝わる熱の流れや、岩の周りを渦巻く水の流れといった、物理法則をコンピュータに理解させようとしていると想像してみてください。かつて、人工知能(AI)でこれを実現しようとするのは、建築家、配管工、電気工事士、大工がそれぞれ異なる言語を話し、異なる設計図を使っている家を建てようとするようなものでした。部屋の形状(幾何学)のためのコード一式、数学的方程式(物理学)のための別のコード、そして実際の学習プロセスのための第三のコードを記述する必要がありました。ある種類の数学から別の種類へ切り替えたい場合、往々にして家全体を解体し、最初からやり直す必要がありました。

jNO(jax Neural Operators)は、「万能翻訳機」と「熟練の建築家」を一つにまとめたような新しいツールです。これは、特に高速科学計算用の人気ツールである JAX プログラミング言語向けに設計されたソフトウェアライブラリであり、物理法則を理解する AI モデルの学習を、より容易に、高速に、柔軟に行えるようにします。

以下に、いくつかの簡単な比喩を用いてその仕組みを説明します。

1. 「ワンスクリプト」の魔法(トレーシングシステム)

jNO を、建設現場全体を制御する単一の魔法の脚本だと考えてください。

  • 以前: 設計図用のスクリプト、数式用の別スクリプト、学習ルール用の別のスクリプトを記述する必要がありました。数式を変更したければ、設計図用のスクリプトも書き直す必要がありました。
  • jNO を使えば: すべてを一つの言語で記述できます。部屋の形状、物理方程式、学習目標を一度に定義します。ソフトウェアは、映画監督がシーンを撮影するように、あなたの指示を「トレーシング(記録)」します。その後、この記録を超効率的で高速なプログラムにコンパイルします。これにより、コードを書き直すことなく、異なる種類の数学的問題に切り替えたり、新しい物理法則を追加したりできるようになります。

2. 「レゴ」の基盤モデル

現在、物理学向けのさまざまな「基盤モデル」(事前学習済み AI 脳)が存在しますが、これらは互いに適合しない異なるメーカーのレゴセットのようです。あるブランドは赤いブロックを使い、別のブランドは青いブロックを使い、それらを積み重ねることができません。

  • jNO の役割: これは万能アダプターとして機能します。Poseidon、Walrus、Morph といった異なる AI モデルを取り込み、それらをすべて同じ JAX 生態系に適合するように変換します。これにより、研究者は事前学習済みの「脳」を少し調整し、独自の物理法則と組み合わせることが可能になり、ソフトウェアツールを切り替える必要がなくなります。

3. 「スマートメッシュ」(形状の処理)

物理学のシミュレーションを行う際、コンピュータは曲がったパイプや複雑な建物などの形状を、「メッシュ」と呼ばれる小さな格子の断片に分割する必要があります。

  • 革新: jNO には組み込みの「スマートメッシュ」システムがあります。これは、単純な正方形であれ、穴のある複雑な 3 次元物体であれ、あなたが記述する任意の形状の上に瞬時に格子を描くロボットのようなものです。格子のどの部分が「内部」で、どの部分が「壁」で、どの部分が「境界」かを追跡するため、AI は物理法則を適用すべき場所を正確に知ることができます。

4. 「微調整」のダイヤル

時には、事前学習済みの AI を特定の新しいタスクに教えたい場合があります。

  • コントロールパネル: jNO は非常に詳細なコントロールパネルを提供します。「これらの脳の一部を凍結して変更させない」「特定の接続からのみ学習させる」「特定の学習速度を使用する」といった指示を AI に与えることができます。全体を再構築することなく、モデルの個々の部分に対してこれを行うことができます。まるで、ギターやボーカルを変えずに、曲の中のドラムの音量だけを調整できるようなものです。

5. 「ダブルモード」エンジン(FEM と PINNs)

この論文は、jNO が物理問題の解決法として 2 つの異なるアプローチを扱えることを強調しています。

  • 点ごとのアプローチ: 地図上の特定の地点の温度をチェックするように、特定の点における物理を検証します。
  • 全体形状アプローチ(有限要素法): 橋全体の応力を計算するように、形状全体にわたる連続的な流れとして物理を捉えます。
  • メリット: jNO を使えば、同じコードでこれら 2 つのモード間を切り替えることができます。エンジンやハンドルを変えることなく、土の道も高速道路も走行できる車を持っているようなものです。

なぜこれが重要なのか

jNO の主な目的は、科学ソフトウェアの「分断」を終わらせることです。研究者が 1 つの AI モデルを学習させるために 5 つの異なるツールをやり取りする代わりに、jNO はすべてを一つの場所に集約します。

  • 速度: 最新のコンピュータチップ上で、JAX の特殊なコンパイル機能を利用することで高速に実行されます。
  • 簡素さ: 異なる種類の物理問題間を切り替えるために、ソフトウェア建築家である必要はありません。
  • 再利用性: jNO でプログラムを書けば、それを保存し、共有し、後で別のコンピュータ上でも実行できます。同じように動作することに確信を持って行えます。

要約すれば、jNO は、複雑な「科学機械学習」の世界を、異なるコード断片を継ぎ接ぎしたパッチワークキルトのように扱うのではなく、単一の整合性のある物語を書くように、シンプルで統合されたものにするよう努めています。

自分の分野の論文に埋もれていませんか?

研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。

Digest を試す →