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

arXiv: https://arxiv.org/abs/2402.13412 · HTML: https://mlsanity.com/paper/2402.13412/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Placed in MoE theory and gating functions, one of 12 directions in Mixture-of-experts.
- Builds on 35 of its 48 references that are in the map.
- First posted on arXiv 20 Feb 2024.

## Placed in

- Mixture-of-experts / [MoE theory and gating functions](https://mlsanity.com/moe/moe-theory-and-gating-functions/)

## Builds on (most cited first, 30 of 35)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Physics-informed neural networks with hard constraints for inverse design](https://mlsanity.com/paper/2102.04626/) (2021)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [DC3: A learning method for optimization with hard constraints](https://mlsanity.com/paper/2104.12225/) (2021)
- [End-to-End Constrained Optimization Learning: A Survey](https://mlsanity.com/paper/2103.16378/) (2021)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Exact imposition of boundary conditions with distance functions in physics-informed deep neural networks](https://mlsanity.com/paper/2104.08426/) (2021)
- [Learning Differential Equations that are Easy to Solve](https://mlsanity.com/paper/2007.04504/) (2020)
- [Scalable Differentiable Physics for Learning and Control](https://mlsanity.com/paper/2007.02168/) (2020)
- [Neural Conservation Laws: A Divergence-Free Perspective](https://mlsanity.com/paper/2210.01741/) (2022)
- [Theseus: A Library for Differentiable Nonlinear Optimization](https://mlsanity.com/paper/2207.09442/) (2022)
- [Learning Neural Constitutive Laws From Motion Observations for Generalizable PDE Dynamics](https://mlsanity.com/paper/2304.14369/) (2023)
- [FluidLab: A Differentiable Environment for Benchmarking Complex Fluid Manipulation](https://mlsanity.com/paper/2303.02346/) (2023)
- [Learned Turbulence Modelling with Differentiable Fluid Solvers: Physics-based Loss-functions and Optimisation Horizons](https://mlsanity.com/paper/2202.06988/) (2022)
- [Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation](https://mlsanity.com/paper/2106.03273/) (2021)
- [Learning differentiable solvers for systems with hard constraints](https://mlsanity.com/paper/2207.08675/) (2022)
- [Sequential Quadratic Optimization for Nonlinear Equality Constrained Stochastic Optimization](https://mlsanity.com/paper/2007.10525/) (2020)
- [Boundary Graph Neural Networks for 3D Simulations](https://mlsanity.com/paper/2106.11299/) (2021)
- [Fully Stochastic Trust-Region Sequential Quadratic Programming for Equality-Constrained Optimization Problems](https://mlsanity.com/paper/2211.15943/) (2022)

## Built on by (most cited first, 25 of 25)

- [AB-UPT: Scaling Neural CFD Surrogates for High-Fidelity Automotive Aerodynamics Simulations via Anchored-Branched Universal Physics Transformers](https://mlsanity.com/paper/2502.09692/) (2025)
- [Generalizing Weather Forecast to Fine-grained Temporal Scales via Physics-AI Hybrid Modeling](https://mlsanity.com/paper/2405.13796/) (2024)
- [Model Decides How to Tokenize: Adaptive DNA Sequence Tokenization with MxDNA](https://mlsanity.com/paper/2412.13716/) (2024)
- [Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model](https://mlsanity.com/paper/2406.19905/) (2024)
- [TRACE: Learning 3D Gaussian Physical Dynamics from Multi-view Videos](https://mlsanity.com/paper/2508.09811/) (2025)
- [End-to-End Probabilistic Framework for Learning with Hard Constraints](https://mlsanity.com/paper/2506.07003/) (2025)
- [Calibrated Physics-Informed Uncertainty Quantification](https://mlsanity.com/paper/2502.04406/) (2025)
- [Enabling Automatic Differentiation with Mollified Graph Neural Operators](https://mlsanity.com/paper/2504.08277/) (2025)
- [Towards a Foundation Model for Partial Differential Equations Across Physics Domains](https://mlsanity.com/paper/2511.21861/) (2025)
- [Robust Biharmonic Skinning Using Geometric Fields](https://mlsanity.com/paper/2406.00238/) (2024)
- [Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates](https://mlsanity.com/paper/2506.05513/) (2025)
- [EddyFormer: Accelerated Neural Simulations of Three-Dimensional Turbulence at Scale](https://mlsanity.com/paper/2510.24173/) (2025)
- [Learning PDE Solvers with Physics and Data: A Unifying View of Physics-Informed Neural Networks and Neural Operators](https://mlsanity.com/paper/2601.14517/) (2026)
- [Projected Neural Differential Equations for Learning Constrained Dynamics](https://mlsanity.com/paper/2410.23667/) (2024)
- [Solving Differential Equations with Constrained Learning](https://mlsanity.com/paper/2410.22796/) (2024)
- [M2M: Learning controllable Multi of experts and multi-scale operators are the Partial Differential Equations need](https://mlsanity.com/paper/2410.11617/) (2024)
- [GITO: Graph-Informed Transformer Operator for Learning Complex Partial Differential Equations](https://mlsanity.com/paper/2506.13906/) (2025)
- [Towards Generalizable PDE Dynamics Forecasting via Physics-Guided Invariant Learning](https://mlsanity.com/paper/2509.24332/) (2025)
- [Differentiable Autoencoding Neural Operator for Interpretable and Integrable Latent Space Modeling](https://mlsanity.com/paper/2510.00233/) (2025)
- [Physics-Informed Heterogeneous Graph Neural Networks for DC Blocker Placement](https://mlsanity.com/paper/2405.10389/) (2024)
- [LE-PDE++: Mamba for accelerating PDEs Simulations](https://mlsanity.com/paper/2411.01897/) (2024)
- [Spatio-temporal, multi-field deep learning of shock propagation in meso-structured media](https://mlsanity.com/paper/2509.16139/) (2025)
- [PHASE: Physics-Integrated, Heterogeneity-Aware Surrogates for Scientific Simulations](https://mlsanity.com/paper/2509.23453/) (2025)
- [Structure tensor Reynolds-averaged Navier-Stokes turbulence models with equivariant neural networks](https://mlsanity.com/paper/2511.09769/) (2025)
- [NUCLEUS-MoE: Unified Model of Pool Boiling for Liquid Cooling](https://mlsanity.com/paper/2605.27722/) (2026)

## Abstract (arXiv)

Imposing known physical constraints, such as conservation laws, during neural network training introduces an inductive bias that can improve accuracy, reliability, convergence, and data efficiency for modeling physical dynamics. While such constraints can be softly imposed via loss function penalties, recent advancements in differentiable physics and optimization improve performance by incorporating PDE-constrained optimization as individual layers in neural networks. This enables a stricter adherence to physical constraints. However, imposing hard constraints significantly increases computational and memory costs, especially for complex dynamical systems. This is because it requires solving an optimization problem over a large number of points in a mesh, representing spatial and temporal discretizations, which greatly increases the complexity of the constraint. To address this challenge, we develop a scalable approach to enforce hard physical constraints using Mixture-of-Experts (MoE), which can be used with any neural network architecture. Our approach imposes the constraint over smaller decomposed domains, each of which is solved by an "expert" through differentiable optimization. During training, each expert independently performs a localized backpropagation step by leveraging the implicit function theorem; the independence of each expert allows for parallelization across multiple GPUs. Compared to standard differentiable optimization, our scalable approach achieves greater accuracy in the neural PDE solver setting for predicting the dynamics of challenging non-linear systems. We also improve training stability and require significantly less computation time during both training and inference stages.