Physics-constrained machine learning for scientific computing

Amazon researchers draw inspiration from finite-volume methods and adapt neural operators to enforce conservation laws and boundary conditions in deep-learning models of physical systems.

Commercial applications of deep learning have been making headlines for years — never more so than this spring. More surprisingly, deep-learning methods have also shown promise for scientific computing, where they can be used to predict solutions to partial differential equations (PDEs). These equations are often prohibitively expensive to solve numerically; using data-driven methods has the potential to transform both scientific and engineering applications of scientific computing, including aerodynamics, ocean and climate, and reservoir modeling.

A fundamental challenge is that the predictions of deep-learning models trained on physical data typically ignore fundamental physical principles. Such models might, for instance, violate system conservation laws: the solution to a heat transfer problem may fail to conserve energy, or the solution to a fluid flow problem may fail to conserve mass. Similarly, a model’s solution may violate boundary conditions — say, allowing heat flow through an insulator at the boundary of a physical system. This can happen even when the model’s training data includes no such violations: at inference time, the model may simply extrapolate from patterns in the training data in an illicit way.

In a pair of recent papers accepted at the International Conference on Machine Learning (ICML) and the International Conference on Learning Representations (ICLR), we investigate the problems of adding known physics constraints to the predictive outputs of machine learning (ML) models when computing the solutions to PDEs.

Related content
Danielle Maddix Robinson's mathematics background helps inform robust models that can predict everything from retail demand to epidemiology.

The ICML paper, “Learning physical models that can respect conservation laws”, which we will present in July, focuses on satisfying conservation laws with black-box models. We show that, for certain types of challenging PDE problems with propagating discontinuities, known as shocks, our approach to constraining model outputs works better than its predecessors: it more sharply and accurately captures the physical solution and its uncertainty and yields better performance on downstream tasks.

In this paper, we collaborated with Derek Hansen, a PhD student in the Department of Statistics at the University of Michigan, who was an intern at AWS AI Labs at the time, and Michael Mahoney, an Amazon Scholar in Amazon’s Supply Chain Optimization Technologies organization and a professor of statistics at the University of California, Berkeley.

In a complementary paper we presented at this year’s ICLR, “Guiding continuous operator learning through physics-based boundary constraints”, we, together with Nadim Saad, an AWS AI Labs intern at the time and a PhD student at the Institute for Computational and Mathematical Engineering (ICME) at Stanford University, focus on enforcing physics through boundary conditions. The modeling approach we describe in this paper is a so-called constrained neural operator, and it exhibits up to a 20-fold performance improvement over previous operator models.

So that scientists working with models of physical systems can benefit from our work, we’ve released the code for the models described in both papers (conservation laws | boundary constraints) on GitHub. We also presented on both works in March 2023 at AAAI's symposium on Computational Approaches to Scientific Discovery.

Danielle Maddix Robinson on physics-constrained machine learning for scientific computing
A talk presented in April 2023 at the Machine Learning and Dynamical Systems Seminar at the Alan Turing Institute.

Conservation laws

Recent work in scientific machine learning (SciML) has focused on incorporating physical constraints into the learning process as part of the loss function. In other words, the physical information is treated as a soft constraint or regularization.

Related content
Hybrid model that combines machine learning with differential equations outperforms models that use either strategy by itself.

A main issue with these approaches is that they do not guarantee that the physical property of conservation is satisfied. To address this issue, in “Learning physical models that can respect conservation laws”, we propose ProbConserv, a framework for incorporating constraints into a generic SciML architecture. Instead of expressing conservation laws in the differential forms of PDEs, which are commonly used in SciML as extra terms in the loss function, ProbConserv converts them into their integral form. This allows us to use ideas from finite-volume methods to enforce conservation.

In finite-volume methods, a spatial domain — say, the region through which heat is propagating — is discretized into a finite set of smaller volumes called control volumes. The method maintains the balance of mass, energy, and momentum throughout this domain by applying the integral form of the conservation law locally across each control volume. Local conservation requires that the out-flux from one volume equals the in-flux to an adjacent volume. By enforcing the conservation law across each control volume, the finite-volume method guarantees global conservation across the whole domain, where the rate of change of the system’s total mass is given by the change in fluxes along the domain boundaries.

Flux Volume Edit-01_230525135151.jpg
The integral form of a conservation law states that the rate of change of the total mass of the system over a domain (Ω) is equal to the difference between the in-flux and out-flux along the domain boundaries (∂Ω).

More specifically, the first step in the ProbConserv method is to use a probabilistic machine learning model — such as a Gaussian process, attentive neural process (ANP), or ensembles of neural-network models — to estimate the mean and variance of the outputs of the physical model. We then use the integral form of the conservation law to perform a Bayesian update to the mean and covariance of the distribution of the solution profile such that it satisfies the conservation constraint exactly in the limit.

Related content
Learning the complete quantile function, which maps probabilities to variable values, rather than building separate models for each quantile level, enables better optimization of resource trade-offs.

In the paper, we provide a detailed analysis of ProbConserv’s application to the generalized porous-medium equation (GPME), a widely used parameterized family of PDEs. The GPME has been used in applications ranging from underground flow transport to nonlinear heat transfer to water desalination and beyond. By varying the PDE parameters, we can describe PDE problems with different levels of complexity, ranging from “easy” problems, such as parabolic PDEs that model smooth diffusion processes, to “hard” nonlinear hyperbolic-like PDEs with shocks, such as the Stefan problem, which has been used to model two-phase flow between water and ice, crystal growth, and more complex porous media such as foams.

For easy GPME variants, ProbConserv compares well to state-of-the-art competitors, and for harder GPME variants, it outperforms other ML-based approaches that do not guarantee volume conservation. ProbConserv seamlessly enforces physical conservation constraints, maintains probabilistic uncertainty quantification (UQ), and deals well with the problem of estimating shock propagation, which is difficult given ML models’ bias toward smooth and continuous behavior. It also effectively handles heteroskedasticity, or fluctuation in variables’ standard deviations. In all cases, it achieves superior predictive performance on downstream tasks, such as predicting shock location, which is a challenging problem even for advanced numerical solvers.

Examples

Conservation of mass.png
Conservation of mass can be violated by a black-box deep-learning model (here, the ANP), even when the PDE is applied as a soft constraint (here, SoftC-ANP) on the loss function, à la physics-informed neural networks (PINNs). This figure shows the variation of total mass over time for the smooth constant coefficient diffusion equation (an “easy” GPME example). The true mass remains zero, since there is zero net flux from the domain boundaries, and thus mass cannot be created or destroyed in the domain interior.
Uncertainty quantification.png
Density solution profiles with uncertainty quantification. In the “hard” version of the GPME problem, also known as the Stefan problem, the solution profile may contain a moving sharp interface in space, known as a shock. The shock here separates the region with fluid from the degenerate one with zero fluid density. The uncertainty is largest in the shock region and becomes smaller in the areas away from it. The main idea behind ProbConserv’s UQ method is to use the uncertainty in the unconstrained black box to modify the mean and covariance at the locations where the variance is largest, to satisfy the conservation constraint. The constant-variance assumption in the HardC-ANP baseline does not result in improvement on this hard task, while ProbConserv results in a better estimate of the solution at the shock and a threefold improvement in the mean squared error (MSE).
Shock position.png
Downstream task. Histogram of the posterior of the shock position computed by ProbConserv and the other baselines. While the baseline models skew the distribution of the shock position, ProbConserv computes a distribution that is well-centered around the true shock position. This illustrates that enforcing physical constraints such as conservation is necessary in order to provide reliable and accurate estimations of the shock position.

Boundary conditions

Boundary conditions (BCs) are physics-enforced constraints that solutions of PDEs must satisfy at specific spatial locations. These constraints carry important physical meaning and guarantee the existence and the uniqueness of PDE solutions. Current deep-learning-based approaches that aim to solve PDEs rely heavily on training data to help models learn BCs implicitly. There is no guarantee, though, that these models will satisfy the BCs during evaluation. In our ICLR 2023 paper, “Guiding continuous operator learning through physics-based boundary constraints”, we propose an efficient, hard-constrained, neural-operator-based approach to enforcing BCs.

Related content
Amazon quantum computing scientist recognized for ‘outstanding contributions to physics’.

Where most SciML methods (for example, PINNs) parameterize the solution to PDEs with a neural network, neural operators aim to learn the mapping from PDE coefficients or initial conditions to solutions. At the core of every neural operator is a kernel function, formulated as an integral operator, that describes the evolution of a physical system over time. For our study, we chose the Fourier neural operator (FNO) as an example of a kernel-based neural operator.

We propose a model we call the boundary-enforcing operator network (BOON). Given a neural operator representing a PDE solution, a training dataset, and prescribed BCs, BOON applies structural corrections to the neural operator to ensure that the predicted solution satisfies the system BCs.

BOON architecture full.png
BOON architectures. Kernel correction architectures for commonly used Dirichlet, Neumann, and periodic boundary conditions that carry different physical meanings.

We provide our refinement procedure and demonstrate that BOON’s solutions satisfy physics-based BCs, such as Dirichlet, Neumann, and periodic. We also report extensive numerical experiments on a wide range of problems including the heat and wave equations and Burgers's equation, along with the challenging 2-D incompressible Navier-Stokes equations, which are used in climate and ocean modeling. We show that enforcing these physical constraints results in zero boundary error and improves the accuracy of solutions on the interior of the domain. BOON’s correction method exhibits a 2-fold to 20-fold improvement over a given neural-operator model in relative L2 error.

Examples

Insulator at boundary.png
Nonzero flux at an insulator on the boundary. The solution to the unconstrained Fourier-neural-operator (FNO) model for the heat equation has a nonzero flux at the left insulating boundary, which means that it allows heat to flow through an insulator. This is in direct contradiction to the physics-enforced boundary constraint. BOON, which satisfies this so-called Neumann boundary condition, ensures that the gradient is zero at the insulator. Similarly, at the right boundary, we see that the FNO solution has a negative gradient at a positive heat source and that the BOON solution corrects this nonphysical result. Guaranteeing no violation of the underlying physics is critical to the practical adoption of these deep-learning models by practitioners in the field.
Stokes's second problem.png
Stokes’s second problem. This figure shows the velocity profile and corresponding absolute errors over time obtained by BOON (top). BOON improves the accuracy at the boundary, which, importantly, also improves accuracy on the interior of the domain compared to the unconstrained Fourier-neural-operator (FNO) model (bottom), where the errors at the boundary propagate inward over time.
Initial condition.png
2-D Navier-Stokes lid-driven cavity flow initial condition. The initial vorticity field (perpendicular to the screen), which is defined as the curl of the velocity field. At the initial time step, t = 0, the only nonzero component of the horizontal velocity is given by the top constant Dirichlet boundary condition, which drives the viscous incompressible flow at the later time steps. The other boundaries have the common no-slip Dirichlet boundary condition, which fixes the velocity to be zero at those locations.

Navier-Stokes lid-driven flow
2-D Navier-Stokes lid-driven cavity flow vorticity field. The vorticity field (perpendicular to the screen) within a square cavity filled with an incompressible fluid, which is induced by a fixed nonzero horizontal velocity prescribed by the Dirichlet boundary condition at the top boundary line for a 25-step (T=25) prediction until final time t = 2.
2-D Navier-Stokes lid-driven cavity flow relative error.
The L2 relative-error plots show significantly higher relative error over time for the data-driven Fourier neural operator (FNO) compared to that of our constrained BOON model on the Navier-Stokes lid-driven cavity flow problem for both a random test sample and the average over the test samples.

Acknowledgements: This work would have not been possible without the help of our coauthor Michael W. Mahoney, an Amazon Scholar; coauthors and PhD student interns Derek Hansen and Nadim Saad; and mentors Yuyang Wang and Margot Gerritsen.

Research areas

Related content

US, NY, New York
We are seeking a Robotics/AI Motor Control Scientist to develop cutting-edge machine learning algorithms for motor control systems in robots. In this role, you will focus on creating and optimizing intelligent motor control strategies to enable robots to perform complex, whole-body tasks. Your contributions will be essential in advancing robotics by enabling fluid, reliable, and safe interactions between robots and their environments. Key job responsibilities - Develop controllers that leverage reinforcement learning, imitation learning, or other advanced AI techniques to achieve natural, robust, and adaptive motor behaviors - Collaborate with multi-disciplinary teams to integrate motor control systems with robotic hardware, ensuring alignment with real-world constraints such as actuator dynamics and energy efficiency - Use simulation and real-world testing to refine and validate control algorithms - Stay updated on advancements in robotics, AI, and control systems to apply advanced techniques to robotic motion challenges - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers - Bridge research initiatives with practical engineering implementation About the team Fauna Robotics, an Amazon company, is building capable, safe, and genuinely delightful robots for everyday life. Our goal is simple: make robots people actually want to live and interact with in everyday human spaces. We believe that future won’t arrive until building for robotics becomes far more accessible. Today, too much effort is spent reinventing the fundamentals. We’re changing that by developing tightly integrated hardware and software systems that make it faster, safer, and more intuitive to create real-world robotic products. Our work spans the full stack: mechanical design, control systems, dynamic modeling, and intelligent software. The focus is not just functionality, but experience. We’re building robots that feel responsive, expressive, and genuinely useful. At Fauna, you’ll work at the frontier of this space, helping define how robots move, manipulate, and interact with people in natural environments. It’s an opportunity to solve hard problems across hardware and software with a team focused on making robotics accessible and joyful to build. If you care about making robotics real for everyone and building systems that are as delightful as they are capable, we’re interested in hearing from you. an opportunity to solve hard problems across hardware and software with a team focused on making robotics accessible and joyful to build. If you care about making robotics real for everyone and building systems that are as delightful as they are capable, we’re interested in hearing from you.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists in the Forecasting, Macroeconomics & Finance field document, interpret and forecast Amazon business dynamics. This track is well suited for economists adept at combining times-series statistical methods with strong economic analysis and intuition. This track could be a good fit for candidates with research experience in: macroeconometrics and/or empirical macroeconomics; international macroeconomics; time-series econometrics; forecasting; financial econometrics and/or empirical finance; and the use of micro and panel data to improve and validate traditional aggregate models. Economists at Amazon are expected to work directly with our senior management and scientists from other fields on key business problems faced across Amazon, including retail, cloud computing, third party merchants, search, Kindle, streaming video, and operations. The Forecasting, Macroeconomics & Finance field utilizes methods at the frontier of economics to develop formal models to understand the past and the present, predict the future, and identify relevant risks and opportunities. For example, we analyze the internal and external drivers of growth and profitability and how these drivers interact with the customer experience in the short, medium and long-term. We build econometric models of dynamic systems, using our world class data tools, formalizing problems using rigorous science to solve business issues and further delight customers.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company.
US, WA, Seattle
Economists in the Forecasting, Macroeconomics & Finance field document, interpret and forecast Amazon business dynamics. This track is well suited for economists adept at combining times-series statistical methods with strong economic analysis and intuition. This track could be a good fit for candidates with research experience in: macroeconometrics and/or empirical macroeconomics; international macroeconomics; time-series econometrics; forecasting; financial econometrics and/or empirical finance; and the use of micro and panel data to improve and validate traditional aggregate models. Economists at Amazon are expected to work directly with our senior management and scientists from other fields on key business problems faced across Amazon, including retail, cloud computing, third party merchants, search, Kindle, streaming video, and operations. The Forecasting, Macroeconomics & Finance field utilizes methods at the frontier of economics to develop formal models to understand the past and the present, predict the future, and identify relevant risks and opportunities. For example, we analyze the internal and external drivers of growth and profitability and how these drivers interact with the customer experience in the short, medium and long-term. We build econometric models of dynamic systems, using our world class data tools, formalizing problems using rigorous science to solve business issues and further delight customers.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists at Amazon partner closely with senior management, business stakeholders, scientist and engineers, and economist leadership to solve key business problems ranging from Amazon Web Services, Kindle, Prime, inventory planning, international retail, third party merchants, search, pricing, labor and employment planning, effective benefits (health, retirement, etc.) and beyond. Amazon Economists build econometric models using our world class data systems and apply approaches from a variety of skillsets – applied macro/time series, applied micro, econometric theory, empirical IO, empirical health, labor, public economics and related fields are all highly valued skillsets at Amazon. You will work in a fast moving environment to solve business problems as a member of either a cross-functional team embedded within a business unit or a central science and economics organization. You will be expected to develop techniques that apply econometrics to large data sets, address quantitative problems, and contribute to the design of automated systems around the company.
US, WA, Seattle
Amazon.com strives to be Earth's most customer-centric company where customers can shop in our stores to find and discover anything they want to buy. We hire the world's brightest minds, offering them a fast paced, technologically sophisticated and friendly work environment. Economists in the Forecasting, Macroeconomics & Finance field document, interpret and forecast Amazon business dynamics. This track is well suited for economists adept at combining times-series statistical methods with strong economic analysis and intuition. This track could be a good fit for candidates with research experience in: macroeconometrics and/or empirical macroeconomics; international macroeconomics; time-series econometrics; forecasting; financial econometrics and/or empirical finance; and the use of micro and panel data to improve and validate traditional aggregate models. Economists at Amazon are expected to work directly with our senior management and scientists from other fields on key business problems faced across Amazon, including retail, cloud computing, third party merchants, search, Kindle, streaming video, and operations. The Forecasting, Macroeconomics & Finance field utilizes methods at the frontier of economics to develop formal models to understand the past and the present, predict the future, and identify relevant risks and opportunities. For example, we analyze the internal and external drivers of growth and profitability and how these drivers interact with the customer experience in the short, medium and long-term. We build econometric models of dynamic systems, using our world class data tools, formalizing problems using rigorous science to solve business issues and further delight customers.