作者 | 論文團(tuán)隊(duì)

編輯 | ScienceAI

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

這正是 Geometry Aware Operator Transformer（GAOT）所嘗試實(shí)現(xiàn)的目標(biāo)。該研究由 ETH Zurich 與 CMU 合作完成，并被 NeurIPS 2025 正式接收。

GAOT 提出了一種面向任意幾何域的高效神經(jīng)算子框架，不僅顯著提升了 PDE 學(xué)習(xí)的精度與魯棒性，更在計(jì)算可擴(kuò)展性（scalability）上取得了實(shí)質(zhì)性突破。

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

1. 背景 | 當(dāng) PDE 遇上復(fù)雜幾何

偏微分方程（PDE）是描述自然規(guī)律的基本語言，從流體力學(xué)到材料模擬，它支撐著幾乎所有連續(xù)物理系統(tǒng)的建模與設(shè)計(jì)。傳統(tǒng)數(shù)值方法（如有限元、有限差分）雖然精確，但在復(fù)雜幾何或高分辨率網(wǎng)格上計(jì)算代價(jià)極高。神經(jīng)算子（Neural Operator）希望通過數(shù)據(jù)驅(qū)動(dòng)的方式學(xué)習(xí) PDE 的解算映射，從而實(shí)現(xiàn)近乎實(shí)時(shí)的預(yù)測。

然而，現(xiàn)有神經(jīng)算子仍面臨兩大瓶頸：

• 多數(shù)模型只能在規(guī)則網(wǎng)格上工作，難以推廣至任意幾何；
• 在大規(guī)模輸入下，計(jì)算效率和顯存占用成為關(guān)鍵限制。

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

GAOT 的目標(biāo)是讓神經(jīng)算子同時(shí)具備幾何感知能力與計(jì)算效率。

它在經(jīng)典的 編碼–處理–解碼（encode–process–decode） 框架上引入了兩項(xiàng)核心創(chuàng)新：

1. 多尺度注意力圖神經(jīng)算子（MAGNO）

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

2. 幾何嵌入（Geometry Embedding）

為每個(gè)點(diǎn)提取局部統(tǒng)計(jì)特征（鄰域密度、平均距離、主方向等），將幾何信息顯式注入模型，提高對(duì)復(fù)雜幾何域的敏感性。

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

圖 1: GAOT 框架結(jié)構(gòu)示意。結(jié)合幾何嵌入與多尺度注意力，實(shí)現(xiàn)任意幾何域上的 PDE 求解。

3. 工程洞察 | 從計(jì)算瓶頸出發(fā)的可擴(kuò)展性優(yōu)化

GAOT 的另一核心亮點(diǎn)，是基于對(duì)計(jì)算與顯存瓶頸的系統(tǒng)分析所提出的一系列架構(gòu)優(yōu)化。

團(tuán)隊(duì)發(fā)現(xiàn)：

• 顯存瓶頸主要來自 Encoder 與 Decoder，它們?cè)谙∈鑸D上的鄰域聚合需要大量邊級(jí)計(jì)算；
• 計(jì)算瓶頸主要集中在 Processor（Transformer 部分），其注意力計(jì)算隨 token 數(shù)量平方增長，是決定訓(xùn)練吞吐率的關(guān)鍵環(huán)節(jié)。

為此，GAOT 設(shè)計(jì)了針對(duì)性優(yōu)化策略：

1. 圖緩存與異步加載

預(yù)存點(diǎn)云鄰接結(jié)構(gòu)，避免重復(fù)構(gòu)圖，并支持異步的圖構(gòu)建與模型計(jì)算。

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

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

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

在極大規(guī)模數(shù)據(jù)上采用邊采樣策略，在幾乎不損失精度的情況下顯著節(jié)省內(nèi)存，并提升模型泛化能力。

這些優(yōu)化使 GAOT 在輸入可擴(kuò)展性（input scalability） 和模型可擴(kuò)展性（model scalability）上都較主流算子模型表現(xiàn)更優(yōu)。

圖 2：訓(xùn)練吞吐率與模型規(guī)模可擴(kuò)展性。隨輸入網(wǎng)格與模型規(guī)模增大時(shí)，GAOT 仍保持最高訓(xùn)練吞吐率，展現(xiàn)出優(yōu)異的輸入與模型可擴(kuò)展性。

同時(shí)，它成為首個(gè)在 9 百萬節(jié)點(diǎn)（9M-point） 的 DrivAerML 數(shù)據(jù)集上實(shí)現(xiàn)全分辨率（full-resolution）訓(xùn)練的神經(jīng)算子模型 —— 這一規(guī)模此前被認(rèn)為難以在單卡 GPU 上完成。

圖 3: 流場預(yù)測結(jié)果對(duì)比（DrivAerML 與 NASA-CRM）

4. 實(shí)驗(yàn) | 同時(shí)實(shí)現(xiàn)精度、效率與可擴(kuò)展性

GAOT 在 28 個(gè) PDE 基準(zhǔn)任務(wù)上，與共計(jì) 14 個(gè)代表性基線模型進(jìn)行了全面評(píng)測，覆蓋時(shí)間相關(guān)與時(shí)間無關(guān)問題、規(guī)則網(wǎng)格與隨機(jī)點(diǎn)云、二維到三維工業(yè)級(jí)仿真等多種場景。結(jié)果顯示：

表 1: 時(shí)間相關(guān)與時(shí)間無關(guān)任務(wù)基準(zhǔn)結(jié)果

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

表 2: DrivAerNet++ 數(shù)據(jù)集誤差分析

值得注意的是，GAOT 在一個(gè)綜合評(píng)測中 —— 涵蓋八個(gè)關(guān)鍵維度：

精度（Accuracy）、魯棒性（Robustness）、訓(xùn)練吞吐率（Throughput）、推理延遲（Latency）、輸入可擴(kuò)展性（Input Scalability）、模型可擴(kuò)展性（Model Scalability）、以及在時(shí)間相關(guān) / 無關(guān)任務(wù)上的穩(wěn)定性（Time-Dependent / Time-Independent Performance）—— 均達(dá)到了最優(yōu)或并列最優(yōu)表現(xiàn)。

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

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

都展現(xiàn)出顯著的綜合優(yōu)勢，首次在同一框架下兼顧了精度、效率與可擴(kuò)展性。

圖 4: 八維度綜合性能評(píng)估

5. 總結(jié) | 面向大規(guī)?？茖W(xué)計(jì)算的高效神經(jīng)算子

GAOT 不僅是一個(gè)新的神經(jīng)算子架構(gòu)，更是一種對(duì)「效率與可擴(kuò)展性」的系統(tǒng)性重思。它通過對(duì)計(jì)算結(jié)構(gòu)的深入理解，實(shí)現(xiàn)了從算法設(shè)計(jì)到工程實(shí)現(xiàn)的協(xié)同優(yōu)化。這一成果表明，AI 不僅可以「近似」數(shù)值求解器，還能在理解幾何與物理規(guī)律的基礎(chǔ)上，以高效、可擴(kuò)展的方式重現(xiàn)科學(xué)計(jì)算的核心過程。未來，團(tuán)隊(duì)計(jì)劃將 GAOT 擴(kuò)展到上億網(wǎng)格的全精度訓(xùn)練，并作為 backbone 開發(fā)多模態(tài)微分方程大模型，探索其在更廣泛科學(xué)問題中的應(yīng)用潛力。