Toward a Physics-Aware TokaMind: Learning Governing Laws of Magnetohydrodynamic Equilibrium

Although foundation models for scientific domains promise scalable prediction across diverse modalities, they often lack explicit constraints to the physical laws governing their prediction targets. In fusion plasma physics, this deficit is particularly significant since surrogates for tokamak dynamics must predict magnetic equilibrium and stability features that are essential for safe operation, but standard data-driven learning approaches do not guarantee that predictions respect the fundamental magnetohydrodynamic equations that dominate plasma dynamics in magnetic confinement, especially considering that the amount of experimental data in fusion is limited compared to other domains. In this paper, we present a general methodology to augment the multi-modal foundation model TokaMind with physics-informed losses derived from the Grad-Shafranov equation, the cornerstone constraint on tokamak equilibrium geometry. We introduce two loss formulations based on strong-form and weak-form (variational) residuals, and develop new evaluation metrics that target physical consistency and operator-space accuracy alongside standard data fidelity. Fine-tuning experiments based on the TokaMark benchmark reveal that physics-informed regularization substantially reduces mutual consistency errors between predicted flux and current density, improves flux accuracy in the operator space through which the physics acts, and softens small-scale structural disruptions in reconstructions. All these improvements come at a modest cost in signal-space normalized RMSE, indicating a fundamental trade-off between data fidelity and physical consistency. Our work bridges foundation models and physical reasoning, providing both a reusable template for physics-informed training in scientific machine learning, and strong evidence that embedding governing laws into TokaMind's learning objective enhances the reliability and interpretability of resulting trained models for high-stakes plasma physics applications. The implementation of the Grad-Shafranov residuals, losses, and diagnostics is open-sourced and will be available upon acceptance.