Research map / Architectures and efficiency / Attention and transformers
Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes NIPS, a Fourier-domain linear-attention neural operator that jointly acts as a forward PDE solver and inverse kernel/parameter discoverer, reducing NAO's quadratic cost on Darcy flow, Mechanical MNIST, and synthetic tissue tasks.
Problem
Learning PDE operators and discovering hidden governing parameters (inverse PDE problems) from limited function pairs is ill-posed, and attention-based operators like NAO scale quadratically in the number of spatial tokens.
Approach
Reformulate NAO's attention as linear attention and replace the explicit square spatial projection with a learnable Fourier-space kernel (channel-independent convolution), so spatial interactions are computed via FFT with O(N d log N) cost and the learned data-dependent kernel gives both forward solutions and the hidden permeability field.
Key results
- On 21x21 Darcy flow with 2 layers, NIPS achieves 2.28% test error vs 3.09% for NAO and 52.09% for AFNO, with 98,396 parameters.
- On Mechanical MNIST with 4 layers, NIPS reaches 1.11% test error vs 5.27% for NAO.
- NIPS scales to 14,641 tokens on a single 40GB A100 whereas quadratic NAO fails at 3.7k tokens and linear NAO at 6.5k tokens.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace NAO with NIPS Fourier kernel learning | Darcy flow · Darcy flow, 441 tokens, 2 layers | test error: 3.09% -> 2.28% | Helped |
| Apply full LayerNorm in every layer (ablation) | NAO · Darcy flow, 441 tokens, 2 layers | test error: 3.09% -> 3.45% | Hurt |
| Add a learnable square projection (ablation) | NAO · Darcy flow, 441 tokens, 2 layers | test error: 3.09% -> 4.00% | Hurt |
| Use NIPS rather than NAO at higher sequence length | Darcy flow · Darcy flow, 1681 tokens, 4 layers | test error: 1.49% -> 1.14% | Helped |
| Increase the number of random permutations (ablation) | NIPS · Darcy flow, 441 tokens, 2 layers, d_k=40 | test error: 42.04% -> 2.28% | Helped |
| Increase the query-key projection dimension d_k (ablation) | NIPS · Darcy flow, 441 tokens, 2 layers, n_rand=50 | test error: 14.67% -> 2.45% | Helped |
| Replace NAO with NIPS | Mechanical MNIST · Mechanical MNIST, 1682 tokens, 2 layers | test error: 5.17% -> 2.11% | Helped |
| Replace NAO with NIPS | Synthetic tissue · Synthetic tissue, 1682 tokens, 4 layers | test error: 5.62% -> 4.95% | Helped |
| Add a two-layer MLP with skip connection to the final V vector (ablation) | NIPS · Synthetic tissue, 1682 tokens, 4 layers | test error: 4.95% -> 4.52% | Helped |
Surprising findings
- Full LayerNorm in every layer (NAO-f) and a learnable square projection (NAO-Wp) both degraded accuracy relative to NAO despite more parameters.
- AFNO baseline fails at zero-shot generalization to unseen PDE parameters, with test errors above 50%.
Limitations
- Performance may degrade when target kernels deviate from training distribution, e.g., non-symmetric kernels from advection-dominated problems.
- Learned kernel is smooth, requiring thresholding to recover two-phase microstructure; boundary regions are unidentifiable under Dirichlet conditions.
- Results rely on in-house synthetic data and a limited set of baselines.
Experimental setup
Datasets: Darcy flow synthetic data (in-house finite difference) Mechanical MNIST
Benchmarks: Darcy flow Mechanical MNIST synthetic tissue learning
Metrics: relative L2 test error kernel error microstructure error per-epoch runtime peak memory trainable parameters
Baselines: NAO NAO-f NAO-Wp AFNO NIPS-mlp
Builds on: Nonlocal Attention Operators Fourier Neural Operator Galerkin-type linear attention Performer
Closest papers by meaning
- Nonlocal Attention Operator: Materializing Hidden Knowledge Towards Interpretable Physics Discovery (NeurIPS 2024 · similarity 0.95)
- Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs (NeurIPS 2024 · similarity 0.86)
- Improved Operator Learning by Orthogonal Attention (ICML 2024 · similarity 0.86)
- Data-Efficient Operator Learning via Unsupervised Pretraining and In-Context Learning (NeurIPS 2024 · similarity 0.86)
- Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators (NeurIPS 2024 · similarity 0.86)
- From Neurons to Neutrons: A Case Study in Interpretability (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.