[论文解读] Choose a Transformer: Fourier or Galerkin

Shuhao Cao|arXiv (Cornell University)|May 31, 2021

Model Reduction and Neural Networks参考文献 97被引用 58

一句话总结

本文引入无 softmax 的傅立叶型和 Galerkin 型自注意力变体，用于 PDE 运算符学习，证明它们的近似能力，并在 PDE 基准上展示相较于基于 softmax 的变换器的效率与准确性提升。

ABSTRACT

In this paper, we apply the self-attention from the state-of-the-art Transformer in Attention Is All You Need for the first time to a data-driven operator learning problem related to partial differential equations. An effort is put together to explain the heuristics of, and to improve the efficacy of the attention mechanism. By employing the operator approximation theory in Hilbert spaces, it is demonstrated for the first time that the softmax normalization in the scaled dot-product attention is sufficient but not necessary. Without softmax, the approximation capacity of a linearized Transformer variant can be proved to be comparable to a Petrov-Galerkin projection layer-wise, and the estimate is independent with respect to the sequence length. A new layer normalization scheme mimicking the Petrov-Galerkin projection is proposed to allow a scaling to propagate through attention layers, which helps the model achieve remarkable accuracy in operator learning tasks with unnormalized data. Finally, we present three operator learning experiments, including the viscid Burgers' equation, an interface Darcy flow, and an inverse interface coefficient identification problem. The newly proposed simple attention-based operator learner, Galerkin Transformer, shows significant improvements in both training cost and evaluation accuracy over its softmax-normalized counterparts.

研究动机与目标

• 鼓励数据驱动的 PDE 运算符学习，以及需要网格不变、分辨率不变的运算符。
• 引入一个无 softmax 的自注意力变体，使其与希尔伯特空间和 Petrov–Galerkin 原理对齐。
• 开发 Galerkin Transformer 以及相应的层归一化，以在未归一化数据上稳定训练。
• 在 PDE 基准上展示相较于基于 softmax 的注意力，训练效率和精度的提升。

提出的方法

• 定义没有 softmax 且带有网格长度加权归一化的简单自注意力变体（Galerkin 型和 Fourier 型）。
• 将注意力解释为希尔伯特空间中对积分变换或 Petrov–Galerkin 投影的矩阵自由近似。
• 显示 Galerkin 型注意力对应一个可学习的 Petrov–Galerkin 投影，具有 n 独立近似的性质；给出 Céa 型误差界（定理 4.3）。
• 提出一种层归一化方案，模仿 Petrov–Galerkin 投影以在训练中稳定并在各层传播缩放。
• 通过前馈网络和位置编码动态丰富潜在基，适应算子响应。
• 将 FT 和 GT 与基于 softmax 的 ST 和 LT 变体在 PDE 任务上进行比较，突出记忆和速度方面的优势。

实验结果

研究问题

• RQ1 无 softmax 的注意力是否能在 PDE 运算符学习中达到与 Petrov–Galerkin 投影相媲美的近似能力？
• RQ2 Galerkin 型注意力是否能在未归一化数据上实现稳定、可扩展的训练，并在效率上超越传统的 softmax 注意力？
• RQ3 Fourier 型和 Galerkin 型注意力在基准 PDE 运算符学习任务（粘性 Burgers 方程、Darcy 流、反问题系数识别）上的表现如何？
• RQ4 层归一化和动态基更新在提高算子学习者的准确性和稳定性方面起到什么作用？

主要发现

• 无 softmax 的 Galerkin 型注意力可以明确表示 Petrov–Galerkin 投影及其近似能力与序列长度无关（定理 4.3）。
• 一种新的 Galerkin 投影型层归一化实现通过注意力层的稳定缩放，并改善物理性质跟踪（如能量衰减）。
• Galerkin Transformer 和 Fourier Transformer 降低了训练内存，相较于基于 softmax 的对手，在算子学习任务中可加速训练。
• 在粘性 Burgers 方程、界面 Darcy 流以及反问题界面系数识别的实验中显示 GT/FT 提高了准确性和效率。
• 相比 softmax，Fourier 型注意力的内存下降 40-50%，在某些情形下 Galerkin 型注意力还可带来显著的速度提升。
• 作者提供用于复现实验结果的开源 PyTorch 代码。