1 of 1

Scaling physics-informed hard constraints with mixture-of-experts

Scientific Achievement

We scale PDE-constrained optimization using mixture-of-experts to enforce differential equation constraints into neural networks (NNs) with a high degree of accuracy. Our new scaled approach is significantly faster and can solve much more challenging problems.

Significance and Impact

Partial differential equations (PDEs) are crucial for describing the complex phenomena of climate dynamics, and numerous other energy-related areas. NNs provide a way to approximate solutions to such systems much faster than numerical methods, but current approaches only enforce physical constraints approximately and are not as accurate – we address this key problem through our method and scale the method to handle larger, much more complex systems.

Technical Approach

We have a differentiable layer in a NN that enforces hard PDE constraints. The layer is differentiable using implicit differentiation, thereby allowing us to train our model with gradient-based optimization methods and train the NN end-to-end. The approach is computationally and memory expensive for complex systems that require fine discretizations, and we scale our approach with mixture-of-experts.

Schematic showing our method with a scaled differentiable layer, which can be added on top of any NN architecture. We scale the differentiable layer using mixture-of-experts, enabling us to scale to larger and finer mesh discretizations to solve spatiotemporal problems.

PI(s)/Facility Lead(s): Lenny Oliker (LBL)

Collaborating Institutions: UC Berkeley, LBNL

ASCR Program: SciDAC RAPIDS2

ASCR PM: Kalyan Perumalla (SciDAC RAPIDS2)

Publication: N. Chalapathi, Y. Du, A. S. Krishnapriyan. International Conference on Learning Representations (2024)

LOCAL LAB POC: Aditi Krishnapriyan

TALKING POINTS:

PDEs are crucial for describing chemical, physical, and biological processes. However, it can be very slow to solve such systems.
Methods based on neural networks (NNs) have shown promise in recent years for physics-based problems
Current NN approaches try to enforce PDEs as “soft constraints” in the loss function, which have numerous optimization issues.
We introduce a new scaled version of a differentiable PDE-constrained layer that enforces the relevant physics much more precisely, and can be trained within a NN framework.
Our results show that by scaling the differentiable PDE-constrained layer with mixture-of-experts and incorporating hard constraints directly into the NN architecture, we can achieve much lower test error when compared to training on an unconstrained objective (previous approaches). This extends for complex, non-linear PDEs.

METADATA:

Name of the associated awarded project: SciDac-5 Institutes RAPIDS2

PI name(s): Lenny Oliker (LBL)

Name of the program manager: Kalyan Permualla (SciDac RAPIDS2)

CITATIONS:

N. Chalapathi, Y. Du, A. S. Krishnapriyan, International Conference on Learning Representations (2024).