Transformer神经算子的零样本泛化到更大域

Zero-shot generalization of transformer neural operators to larger domains

精选理由

零样本推理更大域的新方法

AI 摘要

Transformer神经算子在复杂几何PDE求解中表现出色,但现有方法假设固定域大小限制泛化。本文提出可分解注意力偏置与旋转位置编码,实现空间局部性和平移等变性,使模型在训练域2倍、4倍更大的域上零样本推理。在2个PDE基准(Navier-Stokes、Darcy)和1个3D工业大气流动应用中,该方法显著提升零样本泛化性能。代码和数据集已公开在GitHub。

原文 · arXiv cs.LG

Zero-shot generalization of transformer neural operators to larger domains

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.