Towards Test-Time Adaptation for Neural Surrogates

Neural surrogates that map input configurations (e.g., initial conditions and meshes) to simulation outputs are increasingly used in practical applications such as engineering design optimization. However, pre-trained models often experience significant performance drop on unseen problem configurations, such as different geometries, structural dimensions, and physical parameters. TTA mitigates distribution shifts by leveraging target configurations, online and at test-time. It avoids the need for costly re-training and doesn't require access to the original dataset, which is typically unavailable in practice. In this work, we propose Representation Alignment for Simulations (SimRA), a novel method to improve performance at deployment, specific for multi-dimensional regression on simulation data. SimRA extends prior work on univariate regression [Adachi et al., ICLR 2025] with a novel feature weighting mechanism, ensuring stability in high-dimensional simulation settings. To our knowledge, this is the first study of TTA for neural surrogates. Empirical evaluations on diverse engineering tasks demonstrate strong performance and highlight the potential of TTA in the field.