Stable Test-Time Adaptation for High-Dimensional Simulation Regression

Machine learning surrogates are increasingly used in engineering to accelerate costly simulations of partial differential equations (PDEs), yet distribution shifts between training and deployment often cause severe performance degradation, e.g., under unseen geometries or configurations. Test-Time Adaptation (TTA) can mitigate such shifts, but existing methods are largely developed for classification or image-based tasks on regular grids. When applied to high-dimensional regression on unstructured meshes, a standard setting in engineering simulation, estimating distributional statistics becomes severely ill-conditioned, leading to unstable adaptation. In this work, we propose a stable test-time adaptation method for high-dimensional simulation regressors. Our method uses E-optimal source subset selection to stabilize KL-based alignment between source and target representation distributions. The retained maximally informative source samples further enable stable source-risk regularization and parameter selection at test time. When applied to pretrained simulation surrogates, our method yields up to 12% improvement in out-of-distribution RMSE with modest computational overhead. To the best of our knowledge, this is the first systematic demonstration of effective TTA for high-dimensional simulation regression and generative design optimization, validated across diverse state-of-the-art benchmark datasets.