🔢 mathematics

Preconditioning and Numerical Stability in Neural Network Training for Parametric PDEs

This paper investigates the impact of preconditioning via well-conditioned frame representations on training neural networks for parameter-dependent PDEs, demonstrating significant performance improvements and proposing a novel stable matrix representation that enables accurate computations in single- and half-precision floating point formats.

Original authors: Markus Bachmayr, Wolfgang Dahmen, Chenguang Duan, Mathias Oster

Published 2026-02-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Markus Bachmayr, Wolfgang Dahmen, Chenguang Duan, Mathias Oster

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart but slightly clumsy robot (a Neural Network) to solve a complex puzzle. This puzzle isn't just one picture; it's a whole library of pictures that change slightly depending on a set of knobs you turn (these are the parameters). In the real world, these puzzles are often equations describing how heat flows, how fluids move, or how structures bend under stress.

The paper by Bachmayr, Dahmen, Duan, and Oster is about making this robot learn faster, more accurately, and without getting "confused" by the math, even when the robot is working with very limited mental energy (low-precision computer numbers).

Here is the breakdown of their discovery using simple analogies:

1. The Problem: The "Squishy" Puzzle

When the robot tries to learn, it minimizes a "loss" (a score that tells it how wrong it is). To do this, it has to navigate a landscape of hills and valleys.

The Issue: In many of these physics problems, the landscape is incredibly "squishy" and distorted. Some hills are incredibly steep, and some valleys are incredibly flat. This is called being ill-conditioned.
The Result: If the robot tries to walk down these hills, it might get stuck, bounce back and forth, or take a million tiny steps to get anywhere. This makes training slow and inaccurate.

2. The First Fix: "Preconditioning" (Flattening the Terrain)

The authors suggest using a tool called Preconditioning.

The Analogy: Imagine the robot is trying to walk down a muddy, uneven mountain path. It's hard to get traction. Preconditioning is like laying down a smooth, flat board over the mud. Now, the robot can walk straight down the hill without slipping.
What they did: They used a specific mathematical structure called a Frame Representation (think of it as a special set of building blocks) to reshape the problem. This made the "hills" much more uniform.
The Result: The robot (using an optimizer called Adam) learned much faster and got much closer to the correct answer. The error dropped by thousands of times compared to not using this tool.

3. The Second Problem: The "Fuzzy Glasses" (Numerical Instability)

Here is where the paper gets really clever.

The Issue: Even with the smooth board (preconditioning), the robot still had a problem. The mathematical tools used to build that smooth board were themselves "fuzzy." When the robot tried to use them, it lost tiny bits of information.
The Analogy: Imagine the robot is wearing glasses that are slightly scratched. Even if the path is flat, the robot can't see the details clearly. If the robot is forced to work with low-precision numbers (like 16-bit or 32-bit math, which are like using a ruler with only big markings instead of tiny ones), these scratches cause the robot to make big mistakes. It loses "significant digits," meaning the answer becomes garbage.
The Paper's Claim: Standard ways of doing this math break down when you try to save memory by using low-precision numbers.

4. The Ultimate Fix: "Stable Representations" (Clearing the Glasses)

The authors didn't just stop at flattening the hill; they fixed the robot's glasses.

The Solution: They proposed a new way to write down the math equations. Instead of multiplying several "fuzzy" numbers together, they broke the math down into a specific format where every single step is crystal clear, even if the numbers are small and simple.
The Analogy: They replaced the scratched glasses with a high-definition lens that works perfectly even when the robot is using a cheap, low-resolution ruler.
The Result: This allowed them to train the neural network using half-precision (16-bit) numbers—which is much faster and uses less computer memory—without losing any accuracy. The robot got the same perfect answer as if it were using super-precise (64-bit) math.

5. The "ResNet" Architecture (The Robot's Brain)

The paper also tested different ways to build the robot's brain (the neural network architecture).

They used Residual Networks (ResNets), which are like a robot that checks its own work at every step and corrects small errors before moving on.
They tried three different layouts for this brain. They found that a standard, "all-in-one" brain worked just as well as more complicated, split-up brains. The most important factor wasn't the brain's layout, but the smooth path (preconditioning) and the clear glasses (stable representation).

Summary of the "Win"

Without their method: The robot struggles, takes a long time, and gets stuck in low-precision mode.
With their method: The robot glides down the hill, learns quickly, and can run on cheap, low-power hardware (using 16-bit math) while still producing a perfect, high-precision answer.

In a nutshell: They figured out how to rearrange the math so that a neural network can solve complex physics puzzles quickly and accurately, even when the computer is working with limited precision. They did this by smoothing out the mathematical terrain and ensuring the math doesn't "blur" when simplified.

Technical Summary: Preconditioning and Numerical Stability in Neural Network Training for Parametric PDEs

Problem Formulation
The paper addresses the challenge of efficiently approximating solutions to parameter-dependent partial differential equations (PDEs) using neural networks. Specifically, it focuses on "hybrid" approximations where the dependence on spatial variables $x$ is handled by classical discretization methods (e.g., finite elements), while the dependence on parametric variables $y$ is modeled by neural networks. The solution is approximated as $u(x, y) \approx \sum u_i(y; \theta)\phi_i(x)$ , where the coefficients $u_i$ are outputs of a neural network.

The core difficulty lies in the optimization of the network parameters $\theta$ . Whether minimizing the empirical risk (approximating the solution directly) or minimizing a residual loss (variational correctness), the objective functions involve quadratic forms derived from differential operators. Standard finite element bases often lead to ill-conditioned Gramian matrices, which severely degrade the performance of gradient-based optimization schemes. Furthermore, the authors identify a critical, often overlooked issue: even when preconditioning improves the condition number of the problem, standard implementations of preconditioners can still cause a catastrophic loss of numerical precision, particularly when using low-precision arithmetic (single or half-precision floating-point numbers).

Methodology
The authors propose a two-pronged approach to address ill-conditioning and numerical instability:

Well-Conditioned Frame Representations: Instead of standard finite element bases, the authors utilize multilevel frame representations (specifically, hierarchical bases derived from additive Schwarz methods). By employing a frame $\Phi$ that satisfies the Riesz basis property (or is an overcomplete frame with bounded frame constants), the resulting discretization matrices become well-conditioned. The paper utilizes First-Order System Least Squares (FOSLS) formulations, which allow for direct error quantification and variational correctness.
Stable Matrix Representations: The paper demonstrates that simply preconditioning the system matrix (e.g., forming $H^\top A_y H$ ) is insufficient for numerical stability in low precision. The product of ill-conditioned factors (the preconditioner $H$ and the stiffness matrix $A_y$ ) can still lead to precision loss. To resolve this, the authors propose a "stable representation" of the preconditioned operator. They factorize the operator into $D^\top C_y D$ , where:
- $D$ maps frame coefficients to a space of piecewise polynomials (via differential operators).
- $C_y$ is a diagonal or block-diagonal matrix representing the inner product in the image space.
- Crucially, both $D$ and $C_y$ are constructed to be uniformly well-conditioned independently of the discretization level. This decomposition avoids the multiplication of ill-conditioned matrices, preserving precision even in half-precision (float16) arithmetic.

The authors implement these concepts using a matrix-free numerical realization, evaluating matrix-vector products via tensor contractions on local elements rather than assembling global sparse matrices. They test three neural network architectures: a Full Low-Rank ResNet, a Separate Low-Rank ResNet (one network per frame level), and a Separate Frame Representation (using prolongation operators between levels).

Key Results
Numerical experiments on a 1D parametric elliptic equation yield the following findings:

Impact of Preconditioning: Preconditioning significantly improves the convergence of optimization schemes. For the Adam optimizer, preconditioning reduces the loss by approximately four orders of magnitude compared to unpreconditioned training.
Optimizer Sensitivity: The Adam optimizer benefits most from preconditioning, achieving high accuracy in both single and double precision. In contrast, Stochastic Gradient Descent (SGD) shows limited improvement from preconditioning alone. Second-order methods like L-BFGS and Natural Gradient Descent (NGD) perform well in double precision but are highly sensitive to numerical precision; their performance degrades significantly in half-precision without the stable representation.
Numerical Stability: The proposed "stable representation" ( $D^\top C_y D$ ) enables training with half-precision (float16) arithmetic without loss of accuracy. In half-precision, the Adam optimizer with stable preconditioning achieves a loss value close to the machine precision limit ( $10^{-4}$ ), matching the performance of double-precision training. Standard preconditioning ( $H^\top A_y H$ ) fails to achieve this in half-precision due to precision loss during the matrix-vector product.
Architecture Comparison: The three tested network architectures (Full ResNet, Separate ResNet, Separate Frame Representation) exhibit comparable performance across different precisions, provided the number of trainable parameters is similar. The simpler Full ResNet architecture performed competitively with more complex separate representations.

Significance and Claims
The paper claims that while preconditioning is a known technique for accelerating convergence, its standard application is insufficient for ensuring numerical stability in neural network training, particularly when leveraging low-precision hardware. The primary contribution is the demonstration that a specific factorization of the preconditioned operator—avoiding the explicit product of ill-conditioned matrices—is necessary to maintain full precision in low-precision arithmetic.

The authors assert that this approach allows for the use of half-precision floating-point numbers in training neural networks for parametric PDEs without sacrificing solution accuracy. This is significant for computational efficiency, as half-precision arithmetic is faster and requires less memory, enabling the training of larger models or more complex problems on modern hardware. The work bridges the gap between operator preconditioning theory (specifically Full Operator Preconditioning concepts) and practical neural network training, showing that stable representations can be achieved without reformulating the continuous operator equation, but rather through a non-standard factorization of the discrete matrix representation.