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, WA, Seattle
We are looking for a Principal Applied Scientist to drive the research and development of real-time multimodal conversational AI. You will operate across two focus areas: advancing foundation models for speech and audio, and building the post-training systems (reward modeling, reinforcement learning) that shape natural, human-like conversational behavior. You will be the expert in your area while contributing across the full model lifecycle — from pre-training and architecture design through post-training alignment and real-time deployment. You will work at the frontier of what's possible in conversational AI, with the compute, data, and runway to pursue problems that few teams in the world have the resources to tackle. As a Principal Scientist, you will set the technical direction for your research area, influence the broader roadmap, and work closely with inference engineers to ensure your models are designed for real-time production deployment from inception. Key job responsibilities Foundation Model Scaling - Build and train large-scale multimodal foundation models for real-time speech and audio generation, from architecture design through production-scale training - Advance the scaling and efficiency of conversational modes, including the relationship between data, model size, and real time performance. - Design model architectures informed by hardware constraints and inference requirements, working with inference engineers to ensure models are servable from inception - Develop training methodologies for multimodal models that jointly process and generate speech, language, and audio in real-time streaming contexts Post-Training & Reinforcement Learning - Design and build reward models and reward functions for speech systems — capturing naturalness, fluency, conversational quality, and real-time responsiveness - Develop and apply reinforcement learning methods to shape conversational behavior — teaching models natural timing, responsiveness, and fluid interaction - Build the post-training pipeline from SFT through RL alignment, optimized for real-time multimodal outputs rather than text-only generation - Design evaluation frameworks that capture the quality dimensions unique to real-time conversation Real-Time Perception & Generation - Advance the team's capabilities in real-time perception - Work at the intersection of model architecture and production constraints to ensure multimodal capabilities function within hard real-time latency budgets
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Member of Technical Staff with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Member of Technical Staff with the AGI team, you will lead the development of algorithms and modeling techniques, to advance the state of the art with LLMs. You will lead the foundational model development in an applied research role, including model training, dataset design, and pre- and post-training optimization. Your work will directly impact our customers in the form of products and services that make use of GenAI technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in LLMs. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
US, WA, Redmond
We are searching for a talented candidate with experience in orbital mechanics, orbit determination, launch vehicle trajectories, and launch vehicle mission planning. In this position, a successful candidate would serve as a Research Scientist in support of Amazon Leo’s constellation with particular focus on Launch Vehicle support. Strong analysis skills are required to develop engineering studies of complex large-scale dynamical systems. This position requires demonstrated expertise in computational analysis automation and tool development. Export Control Requirement: Due to applicable export control laws and regulations, candidates must be a U.S. citizen or national, U.S. permanent resident (i.e., current Green Card holder), or lawfully admitted into the U.S. as a refugee or granted asylum. Key job responsibilities Working with the Leo GNC team, you will: • Perform spacecraft maneuver or navigation analysis in support of multi-disciplinary trades within the Amazon Leo team. • Contribute to prototype software development of flight algorithms. • Test and assess navigation software for integration into flight systems. • Assess and trouble-shoot the performance of Leo on-board GNSS hardware and software systems. • Work closely with GNC engineers to manage on-orbit performance and develop flight dynamics operations processes. • Manage engineering trades as needed for various launch vehicle mission designs. • Support Leo’s Launch Vehicle Mission Management team with technical expertise in Launch Vehicle trajectory requirements specification • Evaluate Launch Vehicle performance and compliance with mission requirements • Develop tools to support Mission Management planning for over 80 launches! • Work collaboratively with launch vehicle system technical teams About the team The Flight Dynamics team is responsible for the guidance, navigation, control and safety of the spaceflight of the Amazon Leo constellation. This team provides solutions to spaceflight challenges in constellation design, orbit selection, launch vehicle insertion requirements, navigation, trajectory design, space situational awareness, and space traffic coordination.
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Member of Technical Staff with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Member of Technical Staff with the AGI team, you will lead the development of algorithms and modeling techniques, to advance the state of the art with LLMs. You will lead the foundational model development in an applied research role, including model training, dataset design, and pre- and post-training optimization. Your work will directly impact our customers in the form of products and services that make use of GenAI technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in LLMs. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Member of Technical Staff with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Member of Technical Staff with the AGI team, you will lead the development of algorithms and modeling techniques, to advance the state of the art with LLMs. You will lead the foundational model development in an applied research role, including model training, dataset design, and pre- and post-training optimization. Your work will directly impact our customers in the form of products and services that make use of GenAI technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in LLMs. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Member of Technical Staff with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Member of Technical Staff with the AGI team, you will lead the development of algorithms and modeling techniques, to advance the state of the art with LLMs. You will lead the foundational model development in an applied research role, including model training, dataset design, and pre- and post-training optimization. Your work will directly impact our customers in the form of products and services that make use of GenAI technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in LLMs. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Member of Technical Staff with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Member of Technical Staff with the AGI team, you will lead the development of algorithms and modeling techniques, to advance the state of the art with LLMs. You will lead the foundational model development in an applied research role, including model training, dataset design, and pre- and post-training optimization. Your work will directly impact our customers in the form of products and services that make use of GenAI technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in LLMs. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
IN, KA, Bengaluru
Amazon is looking for a passionate, talented, and inventive Data Scientist with machine learning background to help build industry-leading Speech and Language technology. Our mission is to provide a delightful experience to Amazon’s customers by pushing the envelope in Automatic Speech Recognition (ASR), Natural Language Understanding (NLU), Machine Learning (ML). Key job responsibilities Key job responsibilities Amazon is looking for a passionate, talented, and inventive Data Scientist with machine learning background to help build industry-leading Speech and Language technology. Our mission is to provide a delightful experience to Amazon’s customers by pushing the envelope in Automatic Speech Recognition (ASR), Natural Language Understanding (NLU), Machine Learning (ML) and Computer Vision (CV). As part of our AI team in Amazon AWS, you will work alongside internationally recognized experts to develop data experiments, novel algorithms and data techniques to advance the state-of-the-art in human language technology. Your work will directly impact millions of our customers in the form of products and services that make use of speech and language technology. You will gain hands on experience with Amazon’s heterogeneous speech, text, and structured data sources, and large-scale computing resources to accelerate advances in spoken language understanding.
US, WA, Seattle
We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to apply their causal inference / structural econometrics skillsets to solve real world problems. The intern will work in the area of Store Economics and Science (SEAS) and develop models to SEAS. Our PhD Economist Internship Program offers hands-on experience in applied economics, supported by mentorship, structured feedback, and professional development. Interns work on real business and research problems, building skills that prepare them for full-time economist roles at Amazon and beyond. You will learn how to build data sets and perform applied econometric analysis collaborating with economists, scientists, and product managers. These skills will translate well into writing applied chapters in your dissertation and provide you with work experience that may help you with placement. About the team The Stores Economics and Science Team (SEAS) is a Stores-wide interdisciplinary team at Amazon with a "peak jumping" mission focused on disruptive innovation. The team applies science, economics, and engineering expertise to tackle the business's most critical problems, working to move from local to global optima across Amazon Stores operations. SEAS builds partnerships with organizations throughout Amazon Stores to pursue this mission, exploring frontier science while learning from the experience and perspective of others. Their approach involves testing solutions first at a small scale, then aligning more broadly to build scalable solutions that can be implemented across the organization. The team works backwards from customers using their unique scientific expertise to add value, takes on long-run and high-risk projects that business teams typically wouldn't pursue, helps teams with kickstart problems by building practical prototypes, raises the scientific bar at Amazon, and builds and shares software that makes Amazon more productive.
US, WA, Seattle
Prime Video is a first-stop entertainment destination offering customers a vast collection of premium programming in one app available across thousands of devices. Prime members can customize their viewing experience and find their favorite movies, series, documentaries, and live sports – including Amazon MGM Studios-produced series and movies; licensed fan favorites; and programming from Prime Video add-on subscriptions such as Apple TV+, Max, Crunchyroll and MGM+. All customers, regardless of whether they have a Prime membership or not, can rent or buy titles via the Prime Video Store, and can enjoy even more content for free with ads. Are you interested in shaping the future of entertainment? Prime Video's technology teams are creating best-in-class digital video experience. As a Prime Video technologist, you’ll have end-to-end ownership of the product, user experience, design, and technology required to deliver state-of-the-art experiences for our customers. You’ll get to work on projects that are fast-paced, challenging, and varied. You’ll also be able to experiment with new possibilities, take risks, and collaborate with remarkable people. We’ll look for you to bring your diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. With global opportunities for talented technologists, you can decide where a career Prime Video Tech takes you! We are looking for a self-motivated, passionate and resourceful Applied Scientist to bring diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. You will spend your time as a hands-on machine learning practitioner and a research leader. You will play a key role on the team, building and guiding machine learning models from the ground up. At the end of the day, you will have the reward of seeing your contributions benefit millions of Amazon.com customers worldwide. Key job responsibilities Develop AI solutions for various Prime Video Personalization systems using Deep learning, GenAI, Reinforcement Learning, and optimization methods; Work closely with engineers and product managers to design, implement and launch AI solutions end-to-end; Design and conduct offline and online (A/B) experiments to evaluate proposed solutions based on in-depth data analyses; Effectively communicate technical and non-technical ideas with teammates and stakeholders; Stay up-to-date with advancements and the latest modeling techniques in the field; Publish your research findings in top conferences and journals. About the team Prime Video Personalization and Discovery team owns science solution to power personalized experience on various devices, from sourcing, relevance, ranking, to name a few. We work closely with the engineering and product teams to launch our solutions in production.