Attention-Weighted Physics-Informed Neural Networks for the Unsteady Navier–Stokes Equations

---
title: Attention-Weighted Physics-Informed Neural Networks for the Unsteady Navier–Stokes Equations
author: Ming Wang^1^, Hua Li^1,2^, Siyuan Chen^2^
date: September 2026
abstract: Physics-informed neural networks (PINNs) often fail to converge on unsteady flows at high Reynolds numbers because the gradients of their loss terms are badly out of balance. We propose Attn-PINN, which uses an attention mechanism to weight the residuals adaptively, and prove an upper bound on its approximation error. For flow past a cylinder ($\mathrm{Re}=3\,900$), Attn-PINN reaches a relative $L^2$ error of $1.7\times10^{-3}$ — **an order of magnitude** below the baseline — while training 41% faster.
bibliography: refs.bib
number-sections: true
header-left: Attn-PINN for the Unsteady Navier–Stokes Equations
---

# Introduction

PINNs embed the residual of a partial differential equation in the loss of a neural network [@raissi2019]. They need no mesh and solve forward and inverse problems alike, and have become a core method of scientific machine learning [@karniadakis2021]. For the strongly nonlinear Navier–Stokes equations, however, the gradients of the individual loss terms can differ by more than $10^{3}$, trapping training in a “gradient pathology” [@wang2021]. Neural operators [@li2021] ease the problem but need large amounts of labelled data. Borrowing the attention mechanism of the Transformer [@vaswani2017], we let the network **learn** the weight of every collocation point during training.

# Problem statement {#sec:problem}

Consider incompressible viscous flow in a domain $\Omega\subset\mathbb{R}^2$. The velocity $\mathbf{u}=(u,v)^{\top}$ and pressure $p$ satisfy

$$
\left\{
\begin{aligned}
\frac{\partial \mathbf{u}}{\partial t} + (\mathbf{u}\cdot\nabla)\mathbf{u}
  &= -\frac{1}{\rho}\nabla p + \nu\,\nabla^{2}\mathbf{u} + \mathbf{f},
  && (\mathbf{x},t)\in\Omega\times(0,T],\\
\nabla\cdot\mathbf{u} &= 0,
  && (\mathbf{x},t)\in\Omega\times(0,T],\\
\mathbf{u}(\mathbf{x},0) &= \mathbf{u}_0(\mathbf{x}),
  && \mathbf{x}\in\Omega,
\end{aligned}
\right.
$$ {#eq:ns}

where $\nu$ is the kinematic viscosity. With networks $\hat{\mathbf{u}}_{\theta}$ and $\hat{p}_{\theta}$ approximating the solution, the momentum residual is
$\mathbf{r}_{\theta} = \partial_t\hat{\mathbf{u}}_{\theta} + (\hat{\mathbf{u}}_{\theta}\cdot\nabla)\hat{\mathbf{u}}_{\theta} + \rho^{-1}\nabla\hat{p}_{\theta} - \nu\nabla^{2}\hat{\mathbf{u}}_{\theta} - \mathbf{f}$.

# Method {#sec:method}

## Attention-weighted residuals

For $N_r$ collocation points $\{\mathbf{x}_i\}$, the weights come from scaled dot-product attention over residual features:

$$
\boldsymbol{\lambda} = \operatorname{softmax}\!\left(\frac{Q K^{\top}}{\sqrt{d_k}}\right)V,
\qquad
Q = \begin{pmatrix}
  q_{11} & q_{12} & \cdots & q_{1d}\\
  q_{21} & q_{22} & \cdots & q_{2d}\\
  \vdots & \vdots & \ddots & \vdots\\
  q_{N_r 1} & q_{N_r 2} & \cdots & q_{N_r d}
\end{pmatrix}
$$ {#eq:attn}

The total loss adds the weighted residual, the initial and boundary conditions and the divergence constraint:

$$
\mathcal{L}(\theta)
= \underbrace{\frac{1}{N_r}\sum_{i=1}^{N_r} \lambda_i \bigl\lVert \mathbf{r}_{\theta}(\mathbf{x}_i,t_i) \bigr\rVert_2^{2}}_{\text{PDE residual}}
+ \underbrace{\frac{\beta}{N_b}\sum_{j=1}^{N_b} \bigl\lVert \hat{\mathbf{u}}_{\theta}(\mathbf{x}_j,t_j)-\mathbf{g}_j \bigr\rVert_2^{2}}_{\text{boundary}}
+ \gamma \int_{\Omega} \bigl(\nabla\cdot\hat{\mathbf{u}}_{\theta}\bigr)^{2}\,\mathrm{d}\mathbf{x}
$$ {#eq:loss}

The learning rate follows a piecewise cosine schedule:

$$
\eta(k) =
\begin{cases}
\eta_0 \dfrac{k}{k_w}, & 0 \le k < k_w,\\[6pt]
\eta_{\min} + \dfrac{\eta_0-\eta_{\min}}{2}\left[1+\cos\!\left(\pi\dfrac{k-k_w}{K-k_w}\right)\right], & k_w \le k \le K.
\end{cases}
$$ {#eq:lr}

::: {.important title="Theorem 1 (error bound)"}
Let the exact solution satisfy $\mathbf{u}\in H^{2}(\Omega)$ and let the network have width $m$. If the training loss satisfies $\mathcal{L}(\theta^{*})\le\varepsilon$, there is a constant $C$ independent of $m$ such that

$$
\bigl\lVert \mathbf{u}-\hat{\mathbf{u}}_{\theta^{*}} \bigr\rVert_{L^{2}(\Omega)}
\;\le\; C\left(\sqrt{\varepsilon} + \frac{\ln m}{\sqrt{m}}\right)
\exp\!\left(\int_{0}^{T}\lVert\nabla\mathbf{u}(\cdot,s)\rVert_{L^{\infty}}\,\mathrm{d}s\right).
$$ {#eq:bound}
:::

*Proof sketch.* An energy estimate for the error $\mathbf{e}=\mathbf{u}-\hat{\mathbf{u}}$ and Grönwall’s inequality give

$$
\begin{aligned}
\frac{1}{2}\frac{\mathrm{d}}{\mathrm{d}t}\lVert\mathbf{e}\rVert^{2} + \nu\lVert\nabla\mathbf{e}\rVert^{2}
&\le \lVert\nabla\mathbf{u}\rVert_{L^\infty}\lVert\mathbf{e}\rVert^{2} + \lVert\mathbf{r}\rVert\,\lVert\mathbf{e}\rVert \\
&\le \Bigl(\lVert\nabla\mathbf{u}\rVert_{L^\infty} + \tfrac{1}{2}\Bigr)\lVert\mathbf{e}\rVert^{2} + \tfrac{1}{2}\lVert\mathbf{r}\rVert^{2}.
\end{aligned}
$$

Integrating and substituting [@eq:loss] yields [@eq:bound]. $\blacksquare$

# Experiments {#sec:exp}

We test on flow past a cylinder ($\mathrm{Re}=U D/\nu=3\,900$) with a network of 8 layers of 128 neurons[^1].

<table>
  <caption>Error and training cost of each method {#tbl:cmp}</caption>
  <tr><th rowspan="2">Method</th><th colspan="3" align="center">Relative $L^2$ error ($\times10^{-3}$)</th><th colspan="2" align="center">Training cost</th></tr>
  <tr><th>$u$</th><th>$v$</th><th>$p$</th><th>Time (h)</th><th>Memory (GB)</th></tr>
  <tr><td>PINN [@raissi2019]</td><td align="right">18.6</td><td align="right">24.1</td><td align="right">41.7</td><td align="right">9.8</td><td align="right">6.2</td></tr>
  <tr><td>GradNorm-PINN [@wang2021]</td><td align="right">6.4</td><td align="right">8.9</td><td align="right">15.2</td><td align="right">11.3</td><td align="right">6.9</td></tr>
  <tr><td>FNO [@li2021]</td><td align="right">3.1</td><td align="right">4.0</td><td align="right">7.7</td><td align="right">7.6</td><td align="right">14.5</td></tr>
  <tr><td><b>Attn-PINN (ours)</b></td><td align="right"><b>1.7</b></td><td align="right"><b>2.2</b></td><td align="right"><b>4.9</b></td><td align="right"><b>5.8</b></td><td align="right">7.4</td></tr>
</table>

::: {#fig:result .subfigures}
![Vorticity predicted by Attn-PINN](images/field.png){width=49%} ![Loss during training](images/curve.png){width=49%}

Flow past a cylinder: Attn-PINN reproduces the Kármán vortex street
:::

As [@tbl:cmp] shows, Attn-PINN has the lowest error on all three fields and trains 41% faster than the baseline PINN ([@fig:result]).

# Conclusion

Attn-PINN balances the residual weights with an attention mechanism, comes with an error bound ([@eq:bound]), and clearly outperforms existing methods on a high-Reynolds-number benchmark.

[^1]: Activation $\tanh$, Adam optimiser, $\eta_0=10^{-3}$, $k_w=2\,000$, $K=20\,000$.

# References