Learned PDE Solvers: Von Neumann Stability Analysis for Cell-Average Neural Networks

by Akshat Valse, Mengyuan Yang, & Jue Yan

Department of Mathematics, Iowa State University

We study the stability of cell-average neural network solvers for PDEs, which advance a solution one step at a time with a learned flux. The same network is evaluated at both faces of every cell, so these solvers conserve mass exactly for any weights, but nothing in the construction guarantees stability: networks differing only in their random seed, with training losses you cannot tell apart, can behave completely differently over long rollouts. We are adapting von Neumann stability analysis to them, toward stability conditions on the trained weights.

Write-up

Manuscript in Preparation

Fifteen tall rounded columns on a purple background, their heights rising and falling in one smooth, symmetric mound. Each glows coral low down and warms to pale apricot at its tip.

Introduction

A neural network can be trained to advance a PDE solution one step at a time, and it can be built so that it never creates or destroys mass, whatever its weights are. What that construction does not give you is stability, and telling a stable network from an unstable one is hard for two reasons:

The training objective: Training minimizes a one-step prediction error against a short reference trajectory, and nothing in that objective says what happens when the scheme is iterated for thousands of steps on states the network never saw. Networks trained identically up to the random seed, with training losses you cannot tell apart, can behave completely differently under long rollouts: some track the reference solution, others develop oscillations that grow until the output is worthless (Figure 1).
The usual checks: Since the losses are indistinguishable, the loss cannot tell you which network you have. A short validation rollout is not much better, because the instability may take far longer to show itself than anyone is willing to simulate.

So you are left choosing among minima with no principled basis for the choice. It is the autoregressive-surrogate failure mode, in the simplest setting where the questions of classical numerical analysis can be put to it. The question we ask is:

Can you tell, from a trained network alone and before committing to an expensive simulation, whether its rollout will hold together?

Neural PDE solvers are usually assessed by running them: train, roll out, measure the error. That is expensive, and it only certifies the trajectories you actually ran. Anything you can check on the weights themselves is cheap by comparison and does not depend on having picked the right test case. We are adapting von Neumann stability analysis, the classical Fourier-mode test for finite-difference schemes, to these trained solvers, toward stability conditions on the trained weights.

This page sets out the problem and what we are doing about it. The results are held back until the paper, in preparation with Prof. Jue Yan and Mengyuan Yang, is submitted.

Figure 1. The problem, as a schematic rather than data. Networks trained identically up to the random seed reach training losses you cannot tell apart, yet can behave completely differently under long rollouts. The question at the bottom paraphrases the one we ask.

Cell-Average Neural Network Solvers

A conservative finite-volume scheme advances cell averages by differencing a numerical flux across cell interfaces. In a cell-average-based neural network solver that flux is learned rather than designed: a small network is trained to reproduce the time-averaged flux at an interface. One step of the solver reads

\bar v_j^{\,n+1} = \bar v_j^{\,n} - \tfrac{\Delta t}{\Delta x}\big[\mathcal{N}(V_{j+1/2};\Theta) - \mathcal{N}(V_{j-1/2};\Theta)\big], \tag{1}

where $\bar v_j^{\,n}$ is the average of the solution over cell $j$ at time step $n$ , $\Delta t$ and $\Delta x$ are the time step and the cell width, $\mathcal N(\cdot\,;\Theta)$ is the flux network with weights $\Theta$ , and $V_{j\pm1/2}$ is its input at the face $x_{j\pm1/2}$ . The public write-up does not say which cell averages make up $V_{j+1/2}$ , so Figures 2 and 3 leave the network's input unspecified, and the toy network in Figure 4 uses a placeholder input.

Figure 2 puts the learned flux in the slot that a designed flux would fill. The designed flux on the right is one textbook example, first-order upwind for linear advection, shown only for comparison: it does not stand for the project's equation or reference scheme, neither of which is public.

Learned flux

Designed flux (textbook, for comparison)

Figure 2. Schematics; bar heights are illustrative. Left: in a cell-average neural network solver, one network with the same weights $\Theta$ is evaluated at every face. Right, for comparison only: a designed flux, where the face value is a formula, here textbook first-order upwind for $u_t + a\,u_x = 0$ with $a > 0$ . Either way, each face value is subtracted from the cell on its left and added to the cell on its right: what leaves cell $j$ enters cell $j+1$ .

Conservation, Whatever the Weights

Sum (1) over the cells $j = 1, \dots, M$ of the grid:

\sum_{j=1}^{M} \bar v_j^{\,n+1} = \sum_{j=1}^{M} \bar v_j^{\,n} - \tfrac{\Delta t}{\Delta x}\big[\mathcal N(V_{M+1/2};\Theta) - \mathcal N(V_{1/2};\Theta)\big]. \tag{2}

Every interior face value appears twice, once with each sign, because what leaves cell $j$ through $x_{j+1/2}$ is what enters cell $j+1$ (Figure 3). On a periodic grid $V_{M+1/2} = V_{1/2}$ , so the bracket vanishes for every $\Theta$ : the update telescopes, and mass is conserved exactly, whatever the weights turn out to be. We assume a periodic grid here only for the illustration; the write-up does not say which boundary conditions the project uses.

Figure 3. Why conservation comes for free, drawn for a periodic row of five cells with illustrative bar heights; $\mathcal N_{j+1/2}$ is short for $\mathcal N(V_{j+1/2};\Theta)$ . The network is evaluated once per face, and that one value enters the updates of the two cells that share the face, with opposite signs. Under each face the two terms cancel, and on a periodic grid the two end terms belong to the same face. Nothing here depends on training.

The argument never uses the values of the weights, so it holds just as well for a network with random, untrained weights. Figure 4 checks this live. In floating point the identity can only hold up to rounding, so the check computes the drift from an exact sum of the cell averages, rounded once at the end, and compares it with a rigorous bound on the rounding committed by the updates so far.

Face flux

Network seed 20261110

λ = Δt/Δx 0.40

Cells M 64

Rollout length 500 steps

step 0 of 500

Σ v̄ⁿ (mass/Δx) –

drift 0 (exact 0)

|drift| / (M·max|v̄ⁿ|) 0

rounding bound 0

|drift| ≤ bound: yes

max |v̄ⁿ| –

Figure 4. A live check of identity (2) on a periodic grid, starting from a square pulse and a smooth bump on a constant background, with $\lambda = \Delta t/\Delta x$ . The face flux is either a designed one, the textbook upwind flux $a\,\bar v_j$ at $x_{j+1/2}$ with $a = 1$ , or a tanh network of width 8 whose weights are random and untrained, drawn from the seed shown. In exact arithmetic the drift is zero for either flux and every seed; whatever the profile does, the mass stays put up to rounding. The network is an illustration of "whatever the weights", not a trained solver: its architecture, activation and two-cell input are placeholders chosen for this illustration, not a description of the project's networks, and the rollout says nothing about which trained networks are stable.

Conservation is guaranteed by construction. Stability is not.

Question	Answer	Why
Is mass conserved?	Yes, exactly, for any weights	The update telescopes, identity (2)
Are long rollouts stable?	Not guaranteed	Stability is not built into the construction
Does the training loss tell stable from unstable networks?	No	The losses are indistinguishable
Is a short validation rollout a better test?	Not much better	The instability may take far longer to show itself than anyone is willing to simulate

Adapting von Neumann Analysis

A trained solver is still a numerical scheme. Its weights are frozen once training ends, so what you have in front of you is a fixed, deterministic update rule that happens to have been fitted rather than derived, and numerical analysis has a long history of asking exactly this question about update rules of that kind. Whether that history carries over to a learned flux, and what it says when it does, is what the project is about.

We are adapting von Neumann stability analysis to cell-average neural network solvers of the form (1): deriving the linearized amplification factor of the learned time-stepping scheme and characterizing how perturbations propagate through it, toward stability conditions on the trained weights analogous to those of classical finite-difference methods. I reimplemented the group's solver in Python to enable systematic stability experiments across schemes and step sizes.

The Textbook Case, for Reference

In its textbook form, von Neumann analysis applies to a linear scheme with constant coefficients on a periodic grid. One step carries a single Fourier mode $\bar v_j^{\,n} = \hat v^{\,n} e^{\mathrm{i} j\theta}$ , with $\mathrm{i} = \sqrt{-1}$ and $\theta \in [-\pi, \pi]$ , to a multiple of itself, $\hat v^{\,n+1} = G(\theta)\,\hat v^{\,n}$ , and the scheme is stable in this sense when $|G(\theta)| \le 1$ for every $\theta$ (more generally, $|G(\theta)| \le 1 + C\Delta t$ for a constant $C$ ). For the designed flux of Figure 2, first-order upwind with $\nu = a\,\Delta t/\Delta x$ ,

G(\theta) = 1 - \nu\big(1 - e^{-\mathrm{i}\theta}\big), \qquad |G(\theta)|^2 = 1 - 2\nu(1-\nu)(1-\cos\theta). \tag{3}

so $|G(\theta)| \le 1$ for every $\theta$ exactly when $0 \le \nu \le 1$ . This is the textbook case, shown for comparison only. It is not the project's scheme or result, and this page does not say how the analysis carries over to a learned flux.

Conclusion

We are asking whether you can tell, from a trained cell-average neural network solver alone and before an expensive simulation, whether its rollout will hold together. Conservation comes with the construction, for any weights; stability does not, and neither the training loss nor a short validation rollout settles it. We are adapting von Neumann stability analysis to the question.

Bears on: learned operators · neural surrogates · rollout stability.

Status

The code and the results are in a private repository, pending submission; the problem statement on this page is also on GitHub.

Item	Status
Project	Ongoing, May 2026 – present
Paper	Stability Certificates for Cell-Average-Based Neural Network Solvers: von Neumann Analysis of Trained Flux Networks (working title). In preparation, with Prof. Jue Yan and Mengyuan Yang (Iowa State, Mathematics)
Advisor	Prof. Jue Yan, Department of Mathematics, Iowa State University
Funding	NSF D4 NRT
What we are doing	Adapting von Neumann stability analysis to trained cell-average solvers, toward stability conditions on the weights
Results	Held back pending submission
Code	Private repository, pending submission

Problem write-up

Rare Events All research Early-Feedback Calibration