Stabilizing Test-Time Adaptation of High-Dimensional Simulation Surrogates via D-Optimal Statistics

Anna Zimmel; Paul Setinek; Gianluca Galletti; Johannes Brandstetter; Werner Zellinger

cs.LGFeb 17, 2026

Stabilizing Test-Time Adaptation of High-Dimensional Simulation Surrogates via D-Optimal Statistics

Authors: Anna Zimmel, Paul Setinek, Gianluca Galletti, Johannes Brandstetter, Werner Zellinger

Organizations: ELLIS Unit, LIT AI Lab, Institute for Machine Learning, JKU Linz, Austria · Emmi AI GmbH, Linz, Austria

Abstract

Machine learning surrogates are increasingly used in engineering to accelerate costly simulations, yet distribution shifts between training and deployment often cause severe performance degradation (e.g., unseen geometries or configurations). Test-Time Adaptation (TTA) can mitigate such shifts, but existing methods are largely developed for lower-dimensional classification with structured outputs and visually aligned input-output relationships, making them unstable for the high-dimensional, unstructured and regression problems common in simulation. We address this challenge by proposing a TTA framework based on storing maximally informative (D-optimal) statistics, which jointly enables stable adaptation and principled parameter selection at test time. When applied to pretrained simulation surrogates, our method yields up to 7% out-of-distribution improvements at negligible computational cost. 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 on the SIMSHIFT and EngiBench benchmarks.

Figures & tables

Figure 1: Our method applied to hot rolling task from ( Setinek et al., 2025 ) . (a) Pre-training on the source domain with fixed input parameters, thickness ( $\tau$ ), post-rolling reduction ( $r$ ), and temperature coefficients ( $\lambda_{a}$ , $\lambda_{b}$ ). The representation learner $\phi$ and the predictor $g$ are optimized and maximally informative (D-optimal) statistics are computed. (b) Test-time adaptation of $\phi$ without source data using D-optimal statistics for realizing three TTA pillars: adaptation (KL-based feature alignment), source knowledge preservation (statistics-based regularization) and parameter tuning (importance weighted validation).

Table 1: Comparison of current baselines with TTA methods for all simulation datasets. Results are averaged across 20 TTA runs, over a pretrained model with standard deviation reported. Reported RMSE is normalized over all fields.

Figure 2: Comparison of Equivalent Plastic Strain (PEEQ) predictions on a hot rolling sample. Displaying the Ground Truth (GT), the unadapted Source model, and the SATTS results in the top row, with the absolute residuals,

|\text{GT}-\text{Source}|

and

|\text{GT}-\text{SATTS}|

in the bottom row.

Figure 3: Relative performance improvements of SATTS and the Oracle (lower bound for model selection) compared to the Source model, measured by RMSE.

Dataset	Rolling	Forming	Motor	Heatsink
PAD	1.765	0.929	0.314	1.767

Table 2: PAD values for the simulation datasets.

Table 3: Comparison of current baselines with TTA methods for design optimization datasets. Results are averaged across 20 TTA runs, over one model with standard deviation reported.

Figure 4: Comparison of 2D beam topology results based on ground truth, source prediction, and SATTS. The heatmaps illustrate material density $\rho\in[0,1]$ . SATTS shows stronger alignment with the original design, resulting in more robust design outcomes.

$\hat{\mathcal{R}}_{\mathrm{src}}(f_{\theta})$	IWV	Rolling	Motor	Forming
		$0.566_{\pm 0.020}$	$0.336_{\pm 0.000}$	$0.215_{\pm 0.005}$
✓		$0.550_{\pm 0.020}$	$0.204_{\pm 0.010}$	$0.195_{\pm 0.005}$
✓	✓	$\mathbf{0.545_{\pm 0.019}}$	$\mathbf{0.109_{\pm 0.000}}$	$\mathbf{0.157_{\pm 0.001}}$
Source		$0.561_{\pm 0.001}$	$0.109_{\pm 0.001}$	$0.161_{\pm 0.001}$

Table 4: RMSE scores of SATTS with and without importance weighting and model selection. The baseline and IWV results were evaluated with

lr=0.01

. Best scores are bolded.

Selection Method	Rolling	Motor	Forming
Source Best	$0.550$	$0.203$	$0.157$
IWV	$\mathbf{0.545}$	$\mathbf{0.109}$	$\mathbf{0.157}$

Table 5: RMSE comparison of two model selection algorithms: IWV and source-best on the SIMSHIFT dataset. Best scores are bolded.

TTA Method	Runtime	Increase
SSA	$0.472_{\pm 0.053}$
SATTS	$0.889_{\pm 0.085}$	$(\uparrow\,1.88\times)$

Table 6: Runtime comparison between SSA and SATTS on the Rolling dataset, highlighting the additional overhead of the proposed method. Mean

\pm

std across 10 runs.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value	HF Class Argument Name
Block channels (low $\to$ high)	$[32,64,128,256]$	block_out_channels
Layers per block	$2$	layers_per_block
Transformer layers / block	$1$	transformer_layers_per_block
Cross-attention dim	$64$	cross_attention_dim
Only cross-attention	True	only_cross_attention=True
Normalization groups	$16$	norm_num_groups

Appendix

Table 7: Hyperparameters for our conditional diffusion U-Net. This setup leads to a model size of 17.5M parameters.

Figure 5: t-SNE visualization of the conditioner’s latent space on the structural beam bending dataset. While overhang_constraint and forcedist are either constant or exhibit almost a uniform distribution, volfrac and rmin exhibit a clear structure.

Figure 6: t-SNE visualization of the conditioner’s latent space on the structural beam bending dataset. Both conditions ( volume and lentgh ) exhibit a clear structure in the latent space.

Dataset	Parameter	Description	Source range (no. samples)	Target range (no. samples)
Beams2D	rmin	Minimum feature length of beam members.	[1.5, 3.25) (3087)	[3.25, 4] (353)
HeatConduction2D	volume	Volume limits on the material distributions.	[0.3, 0.465) (231)	[0.465, 0.6] (39)

Appendix

Table 8: Defined distribution shifts (source and target domains) for each dataset.

Test-Time Adaptation

Test-time adaptation updates or recalibrates pretrained models at inference to handle distribution shifts, often without target labels.

See all 322 papers

Multi-Hypothesis Test-Time Adaptation to Mitigate Underspecification
Jun 30, 2026Afshar Shamsi, Xiao-Yu Guo, Hamid Alinejad-Rokny +3Distribution Shift Robustness Test-Time Adaptation
Concordia University, Canada · Australian Institute for Machine Learning, Australia · University of New South Wales, Australia +2
Adaptive surrogate modeling for high-dimensional spatio-temporal output
Aug 18, 2026Berkcan Kapusuzoglu, Shunsaku Matsumoto, Yoshitomo Miyagi +2Surrogate Modeling Adaptive Sampling
Department of Civil and Environmental Engineering, Vanderbilt University, Nashville, TN 37235, USA · Strength Research Department, Research and Innovation Center, Mitsubishi Heavy Industries, Ltd., Nagasaki, 851-0392, Japan · Strength Research Department, Research and Innovation Center, Mitsubishi Heavy Industries, Ltd., Takasago, 676-8686, Japan
Revalidation Beats Stateful Routing for Scientific Surrogates Under Distribution Shift
Aug 31, 2026Harshil LodhiyaSurrogate Modeling Distribution Shift Robustness
SlicedHealth

Jun 30, 2026J/K move · Enter open · S save

Multi-Hypothesis Test-Time Adaptation to Mitigate Underspecification

Afshar Shamsi, Xiao-Yu Guo, Hamid Alinejad-Rokny, Arash Mohammadi+2

Concordia University, Canada · Australian Institute for Machine Learning, Australia · University of New South Wales, Australia+2

Distribution Shift Robustness Test-Time Adaptation

Open paper AskPDF arXiv

Test-Time Adaptation (TTA) seeks to improve model robustness under distribution shifts by adapting parameters using unlabeled target data. However, in the absence of supervision, entropy-based adaptation is fundamentally underconstrained: multiple distinct parameter updates can achieve similarly low entropy while inducing drastically different decision boundaries. This phenomenon, known as underspecification, renders standard TTA brittle and prone to collapse into spurious modes. In this work, we reinterpret TTA through a posterior-inspired lens induced by entropy minimization, where low-entropy solutions define a pseudo-likelihood over parameters. Instead of committing to a single point estimate, we introduce a particle-based diversification framework that explores multiple plausible adaptation trajectories simultaneously. Our method can be viewed as a structured exploration of multiple plausible adaptation solutions, implemented through multi-level diversification at the output, parameter, optimizer, and input levels. Crucially, the framework acts as a plug-and-play wrapper compatible with existing TTA methods. Extensive experiments on challenging benchmarks demonstrate consistent gains in stability and robustness, achieving improvements of 3-4% under mixed shifts, 2-3% with batch size one, and 1-2.5% under label shifts, outperforming state-of-the-art baselines. Our results suggest that treating TTA as a multi-hypothesis inference problem, rather than a single-point optimization task, is key to mitigating underspecification and enabling reliable real-world deployment.

Abstract

Test-Time Adaptation

Explore similar work

Multi-Hypothesis Test-Time Adaptation to Mitigate Underspecification

Adaptive surrogate modeling for high-dimensional spatio-temporal output

Revalidation Beats Stateful Routing for Scientific Surrogates Under Distribution Shift