让Transformer神经算子在未训练的大域上零样本泛化,提升跨域预测能力。
Zero-shot generalization of transformer neural operators to larger domains

- 通过可分解偏置设计注意力机制,实现可控空间局部性。
- 在两个PDE基准和3D大气流模拟中显著提升大域零样本性能。
- 无需修改架构即可适配大域,适合科学计算与工业仿真场景。
基于Transformer的神经算子在复杂几何上逼近偏微分方程解算子方面表现出色。然而,现有方法隐含假设域大小固定,限制了推理时的泛化能力。本文研究域扩展问题,即在训练时未见的大空间域上进行零样本推理。我们指出,该设定需具备空间局部性和平移等变性。提出通过注意力逻辑中引入可分解偏置来实现局部性,既能精细控制空间支持范围,又可完全分解为查询-键内积,兼容优化注意力核。结合旋转位置编码,实现表达力强且支持可控的空间嵌入,无需改变Transformer结构。实验证明,该方法在两个PDE基准和一个3D工业大气流应用中均显著提升对更大域的零样本泛化能力。代码与数据集见https://github.com/cerea-daml/domain-extension。
原文摘要 · Abstract (English)
Transformer-based neural operators have shown remarkable performance for approximating solution operators of partial differential equations on complex geometries. However, existing approaches implicitly assume a fixed domain size, which limits their ability to generalize at inference. In this work, we investigate domain extension, namely zero-shot inference on spatial domains that are significantly larger than those encountered during training. We argue that this setting fundamentally requires spatial locality and translation equivariance. We propose to implement this locality via a decomposable bias in the attention logits computation, enabling finely controllable locality while remaining fully decomposable into query-key inner products and directly compatible with optimized attention kernels. Combined with rotary positional embeddings, it enables expressive embeddings with controllable spatial support without altering the transformer architecture. We empirically show that our approach substantially improves zero-shot generalization to larger domains across two PDE benchmarks and a 3D industrial atmospheric flow application. Our code and datasets are available at https://github.com/cerea-daml/domain-extension.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。