ETHz提出GAOT，讓神經算子在任意形狀上實現高效可擴展的PDE求解

作者 | 論文團隊

編輯 | ScienceAI

想象一下，如果 AI 能真正理解復雜幾何的形狀，并快速預測其中的物理場分布 —— 無論是汽車外殼的氣流走向，還是飛機機翼上的壓力變化，都能在幾秒內被準確模擬，而不再依賴昂貴的數值仿真。

這正是 Geometry Aware Operator Transformer（GAOT）所嘗試實現的目標。該研究由 ETH Zurich 與 CMU 合作完成，并被 NeurIPS 2025 正式接收。

GAOT 提出了一種面向任意幾何域的高效神經算子框架，不僅顯著提升了 PDE 學習的精度與魯棒性，更在計算可擴展性（scalability）上取得了實質性突破。

• 論文標題：Geometry Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domains
• 論文連接：https://arxiv.org/abs/2505.18781
• 項目主頁：https://camlab-ethz.github.io/GAOT

1. 背景 | 當 PDE 遇上復雜幾何

偏微分方程（PDE）是描述自然規律的基本語言，從流體力學到材料模擬，它支撐著幾乎所有連續物理系統的建模與設計。傳統數值方法（如有限元、有限差分）雖然精確，但在復雜幾何或高分辨率網格上計算代價極高。神經算子（Neural Operator）希望通過數據驅動的方式學習 PDE 的解算映射，從而實現近乎實時的預測。

然而，現有神經算子仍面臨兩大瓶頸：

• 多數模型只能在規則網格上工作，難以推廣至任意幾何；
• 在大規模輸入下，計算效率和顯存占用成為關鍵限制。

2. 方法 | Geometry-Aware Operator Transformer (GAOT)

GAOT 的目標是讓神經算子同時具備幾何感知能力與計算效率。

它在經典的 編碼–處理–解碼（encode–process–decode） 框架上引入了兩項核心創新：

1. 多尺度注意力圖神經算子（MAGNO）

在編碼與解碼階段，通過多尺度鄰域聚合與注意力機制整合不同尺度的幾何特征，實現跨尺度的物理建模。

2. 幾何嵌入（Geometry Embedding）

為每個點提取局部統計特征（鄰域密度、平均距離、主方向等），將幾何信息顯式注入模型，提高對復雜幾何域的敏感性。

同時，GAOT 在潛空間中采用了 Vision Transformer 式的 patch 機制，兼顧全局建模能力與計算效率。

圖 1: GAOT 框架結構示意。結合幾何嵌入與多尺度注意力，實現任意幾何域上的 PDE 求解。

3. 工程洞察 | 從計算瓶頸出發的可擴展性優化

GAOT 的另一核心亮點，是基于對計算與顯存瓶頸的系統分析所提出的一系列架構優化。

團隊發現：

• 顯存瓶頸主要來自 Encoder 與 Decoder，它們在稀疏圖上的鄰域聚合需要大量邊級計算；
• 計算瓶頸主要集中在 Processor（Transformer 部分），其注意力計算隨 token 數量平方增長，是決定訓練吞吐率的關鍵環節。

為此，GAOT 設計了針對性優化策略：

1. 圖緩存與異步加載

預存點云鄰接結構，避免重復構圖，并支持異步的圖構建與模型計算。

2.Encoder/Decoder 順序處理 + Processor 批處理

對編碼器和解碼器采用順序化處理，能有效緩解大 batch size 下的顯存峰值；同時對 Processor 采用計算密集型操作， 進行批量并行計算，在顯存受限時仍保持高吞吐率。

3. 可選邊采樣（Edge Mask）策略

在極大規模數據上采用邊采樣策略，在幾乎不損失精度的情況下顯著節省內存，并提升模型泛化能力。

這些優化使 GAOT 在輸入可擴展性（input scalability） 和模型可擴展性（model scalability）上都較主流算子模型表現更優。

圖 2：訓練吞吐率與模型規模可擴展性。隨輸入網格與模型規模增大時，GAOT 仍保持最高訓練吞吐率，展現出優異的輸入與模型可擴展性。

同時，它成為首個在 9 百萬節點（9M-point） 的 DrivAerML 數據集上實現全分辨率（full-resolution）訓練的神經算子模型 —— 這一規模此前被認為難以在單卡 GPU 上完成。

圖 3: 流場預測結果對比（DrivAerML 與 NASA-CRM）

4. 實驗 | 同時實現精度、效率與可擴展性

GAOT 在 28 個 PDE 基準任務上，與共計 14 個代表性基線模型進行了全面評測，覆蓋時間相關與時間無關問題、規則網格與隨機點云、二維到三維工業級仿真等多種場景。結果顯示：

表 1: 時間相關與時間無關任務基準結果

• 在幾乎所有任務中取得最佳或次優性能；
• 訓練吞吐量較代表性模型（GINO、RIGNO、Transolver）提升約 50%；
• 推理延遲降低 15–30%；
• 在三維工業 CFD 數據集（DrivAerNet++、DrivAerML、NASA-CRM）上均取得最新最優結果（SOTA）；
• 在 DrivAerNet++ 數據集上，壁面壓力和剪應力預測精度較次優模型分別提升約 15% 和 30%。

表 2: DrivAerNet++ 數據集誤差分析

值得注意的是，GAOT 在一個綜合評測中 —— 涵蓋八個關鍵維度：

精度（Accuracy）、魯棒性（Robustness）、訓練吞吐率（Throughput）、推理延遲（Latency）、輸入可擴展性（Input Scalability）、模型可擴展性（Model Scalability）、以及在時間相關 / 無關任務上的穩定性（Time-Dependent / Time-Independent Performance）—— 均達到了最優或并列最優表現。

換言之，GAOT 在三類主流模型范式中 ——

• Transformer-based（如 Transolver、GNOT）
• Graph-based（如 RIGNO）
• FNO-based（如 GINO、GeoFNO）

都展現出顯著的綜合優勢，首次在同一框架下兼顧了精度、效率與可擴展性。

圖 4: 八維度綜合性能評估

5. 總結 | 面向大規模科學計算的高效神經算子

GAOT 不僅是一個新的神經算子架構，更是一種對「效率與可擴展性」的系統性重思。它通過對計算結構的深入理解，實現了從算法設計到工程實現的協同優化。這一成果表明，AI 不僅可以「近似」數值求解器，還能在理解幾何與物理規律的基礎上，以高效、可擴展的方式重現科學計算的核心過程。未來，團隊計劃將 GAOT 擴展到上億網格的全精度訓練，并作為 backbone 開發多模態微分方程大模型，探索其在更廣泛科學問題中的應用潛力。