Physics-based machine learning for mantle convection

Mantle convection simulations are central to understanding the thermal evolution of rocky planets. However, their high computational cost limits the feasibility of extensive parameter studies needed to constrain models with observations. While scaling laws provide low-cost alternatives, they are limited in the physical processes they can capture and typically predict only reduced quantities rather than the full spatio-temporal fields.

Machine learning (ML) promises to accelerate mantle convection simulations, yet purely data-driven approaches can fail to match the accuracy and stability of numerical solvers, even when trained on thousands of simulations. To address this, we propose a physics-based ML framework that combines neural networks with numerical time integration. The ML model predicts creeping-flow velocities as a function of temperature, thereby bypassing the numerical solution of the Stokes equations, which poses the primary computational bottleneck in mantle convection simulations. Mass conservation is enforced as a hard constraint through a stream-function formulation. The predicted velocity field is then used by a finite-volume solver to advect and diffuse temperature forward in time.

The model is trained on temperature–velocity snapshots from 94 two-dimensional simulations of statistically steady mantle convection, where three parameters are varied: internal heating as well as pressure- and temperature-dependence of viscosity. Compared to the direct numerical solver, our model is 89 times faster. For some parameter combinations, the model outperforms an under-relaxed iterative numerical solver in speed and accuracy, further underscoring the potential of ML in geodynamics. Ablation studies demonstrate the importance of mass conservation, learned boundary padding, and loss scaling in achieving stable and accurate predictions over long time-integration scales.

Despite being trained exclusively on snapshots from statistically steady simulations, the model successfully performs thermal evolution, demonstrating generalization in this unseen setting. Performance degrades, however, when additional compressibility effects are introduced at inference or when initial conditions deviate substantially from the training data, highlighting directions for future improvements.