SCALING PHYSICS-INFORMED HARD CONSTRAINTS WITH MIXTURE-OF-EXPERTS

SCALING PHYSICS-INFORMED HARD CONSTRAINTS WITH MIXTURE-OF-EXPERTS

2024 | Nithin Chalapathi, Yiheng Du, Aditi S. Krishnapriyan
This paper introduces a scalable approach to enforce hard physical constraints in neural networks using Mixture-of-Experts (MoE), called Physics-Informed Hard Constraint Mixture-of-Experts (PI-HC-MoE). The method enables strict adherence to physical laws by decomposing the domain into smaller subdomains, each handled by an "expert" through differentiable optimization. This approach improves accuracy, training stability, and computational efficiency compared to traditional soft constraint methods and standard differentiable optimization. The PI-HC-MoE framework is applied to two challenging non-linear systems: 1D diffusion-sorption and 2D turbulent Navier-Stokes equations. The method outperforms both physics-informed soft constraint (PI-SC) and physics-informed hard constraint (PI-HC) approaches in terms of accuracy and efficiency. For the diffusion-sorption problem, PI-HC-MoE achieves significantly lower error rates and faster convergence. For the Navier-Stokes problem, it also shows better performance, with lower relative $ L_2 $ error and faster inference times. The key advantage of PI-HC-MoE is its scalability. It maintains near-constant execution time as the number of sampled points increases, while traditional methods become significantly slower. This is achieved by decomposing the domain into smaller subdomains, allowing each expert to solve a localized optimization problem independently. This parallelization across multiple GPUs improves training stability and reduces computational costs. The method is implemented using JAX and is available for reproducibility. It demonstrates sub-linear scaling in terms of computational efficiency, making it suitable for large-scale physical modeling tasks. The framework is particularly effective in scenarios where physical constraints are critical, such as in fluid dynamics, heat transfer, and chemical reaction modeling. The results show that PI-HC-MoE provides a more accurate and efficient solution to physical modeling problems compared to existing methods.This paper introduces a scalable approach to enforce hard physical constraints in neural networks using Mixture-of-Experts (MoE), called Physics-Informed Hard Constraint Mixture-of-Experts (PI-HC-MoE). The method enables strict adherence to physical laws by decomposing the domain into smaller subdomains, each handled by an "expert" through differentiable optimization. This approach improves accuracy, training stability, and computational efficiency compared to traditional soft constraint methods and standard differentiable optimization. The PI-HC-MoE framework is applied to two challenging non-linear systems: 1D diffusion-sorption and 2D turbulent Navier-Stokes equations. The method outperforms both physics-informed soft constraint (PI-SC) and physics-informed hard constraint (PI-HC) approaches in terms of accuracy and efficiency. For the diffusion-sorption problem, PI-HC-MoE achieves significantly lower error rates and faster convergence. For the Navier-Stokes problem, it also shows better performance, with lower relative $ L_2 $ error and faster inference times. The key advantage of PI-HC-MoE is its scalability. It maintains near-constant execution time as the number of sampled points increases, while traditional methods become significantly slower. This is achieved by decomposing the domain into smaller subdomains, allowing each expert to solve a localized optimization problem independently. This parallelization across multiple GPUs improves training stability and reduces computational costs. The method is implemented using JAX and is available for reproducibility. It demonstrates sub-linear scaling in terms of computational efficiency, making it suitable for large-scale physical modeling tasks. The framework is particularly effective in scenarios where physical constraints are critical, such as in fluid dynamics, heat transfer, and chemical reaction modeling. The results show that PI-HC-MoE provides a more accurate and efficient solution to physical modeling problems compared to existing methods.
Reach us at info@study.space
[slides and audio] Scaling physics-informed hard constraints with mixture-of-experts