Transformer-Powered Surrogates Close the ICF Simulation-Experiment Gap with Extremely Limited Data

Recent advances in machine learning, specifically transformer architecture, have led to significant advancements in commercial domains. These powerful models have demonstrated superior capability to learn complex relationships and often generalize better to new data and problems. This paper presents a novel transformer-powered approach for enhancing prediction accuracy in multi-modal output scenarios, where sparse experimental data is supplemented with simulation data. The proposed approach integrates transformer-based architecture with a novel graph-based hyper-parameter optimization technique. The resulting system not only effectively reduces simulation bias, but also achieves superior prediction accuracy compared to the prior method. We demonstrate the efficacy of our approach on inertial confinement fusion experiments, where only 10 shots of real-world data are available, as well as synthetic versions of these experiments.


Introduction
Simulation-driven science relies on the premise that sophisticated computational simulations can enable researchers to explore complex phenomena which can be challenging to explore experimentally due to prohibitive costs, time constraints, or both.Over the recent years, we have witnessed a major interest surge [1,2,3,4,5] in leveraging such large-scale simulation data along with machine learning (ML) methodologies to drive our understanding of complex physical systems.Despite its flexibility, this approach comes with an implicit understanding that simulations are often lower fidelity representations of the true physical phenomena and can hence contain critical gaps when translating insights to real experiments [6].In other words, ML models trained purely on simulation data can inherit its biases, and limitations, and can eventually lead to severe miscalibration with respect to the experiments.
A viable approach to mitigate this gap is to systematically adapt simulation-trained models using a handful of experimental observations, enabling the models to adjust their biases to match experimental measurements more closely through transfer learning, a method where a model developed for one task is repurposed for another [7].When successful, this strategy can be remarkably effective at predicting experiment outcomes (or even intermediate states) accurately, while requiring only a small fraction of the experimental observations that typically would be needed to train sophisticated ML models (e.g., deep neural networks) if experimental data alone were used [5].However, two critical challenges need to be addressed when building practical, transfer learning protocols: (i) the heightened risk for overfitting in cases of extremely few-shot data (∼10-20 experiments), since surrogates can typically contain a large number of parameters of the order of hundreds of thousands or even millions; and (ii) the lack of clear guidance for hyper-parameter selection (e.g., learning rate, number of epochs for optimization).Imprecise choice of hyper-parameters during model fine-tuning can lead to several undesirable effects (e.g., excessive feature distortion or an increased risk of simplicity bias [8]), therefore resulting in poor generalization.The conventional practice of using a held-out validation dataset for hyper-parameter selection is no longer applicable to our setting of performing transfer learning with very limited data.
In this work, we address these issues using inertial confinement fusion (ICF) [9] as a test bed, where the simulation-experiment gap is well documented [3,4] and the number of available experimental observations are very few (10) due to their high cost (∼$1M/experiment).First, recognizing the need for a more generalizable base model to enhance transfer learning performance, we introduce a novel framework specifically for training in transformer-based architectures [10].Transformers have demonstrated their adaptability and effectiveness across many domains, including language [11,12,13,14,15], vision [16,17,18,19], audio [20,21,22], chemistry [23,24,25,26], and biology [27,28].Building upon the versatility of transformers, we introduce a novel framework specifically designed for masked training in transformerbased architectures by utilizing masked auto-encoders [29], where masking involves selectively hiding parts of the data to enhance model learning without labels.Targeted at enhancing the adaptability and expressiveness of ML models trained on simulations, this framework accommodates a variety of designer-specified masking strategies, such as forward modeling, inverse modeling, or combinations thereof.While the framework is flexible enough to support any masked modeling approach, we specifically focus on two strategies: forward modeling for predictive learning and random masking of an entire data sample.These strategies enable the model to jointly learn the complex dependencies between simulation inputs and outputs-akin to a standard surrogate model-as well as the correlations across disparate output modalities resembling modern representation learners.Second, we introduce a novel hyper-parameter selection approach for model fine-tuning.Our approach models different hyper-parameter choices as nodes of a graph, their corresponding validation errors as the function at each node, and adopts a graph filtering strategy for reliable hyper-parameter recommendation.To demonstrate that our proposed techniques are statistically meaningful, we also show improvements using a larger, synthetic ICF dataset, where the simulation-experiment gap is artificially built by splitting the datasets along known physics parameters [5].

Experiment Description
Reference Our primary scalar prediction results showing significant improvements from our methods versus Kustowski et al. [5] Table 2 Our primary image predictions from our model versus Kustowski et al. [5] , with consistent improvements for the synthetic data scenario.Figure 3 An analysis between pretrained learned embeddings and fine-tuned embeddings showing consistent simulation bias for simple hyperparameter selection.Table 4 Experiment using significantly more synthetic data (50 training points) to show graph smoothing matches minimum validation error.Main Findings We evaluate our methods on a real-world benchmark from the literature [4], which comprises ICF simulations and real experiments curated at the National Ignition Facility (NIF); and a more recent Hydra simulation-based synthetic benchmark [5] that emulates the large distribution shifts typically observed in the real world.We find that our transformer-based surrogate, combined with our robust hyperparameter selection strategy, is significantly more effective at bridging the simulation-experiment gap, offering a relative gain of ∼ 40% in terms of predictive error over the state-of-the-art neural network surrogates.More specifically, we find that our richer class of transformer-based surrogates enables us to employ a simpler transfer learning protocol (a simple linear-bias correction as opposed to extensive neural network weight finetuning), therefore, making it ideal for applications operating in very small experimental data regimes.Next, we find that the graph-based hyper-parameter selection strategy yields much more robust and generalizable models that outperform traditional validation techniques significantly.We present an overview of our method in figure 1, and we summarize our experiments in table 1.

Experimental Setup and Results
In our effort to bridge the gap between simulation and experimental data, we employ our proposed framework integrating masked training in transformer-based architectures and a graph-based hyper-parameter selection strategy that is particularly effective when the number of experimental observations is very small.We begin by assessing the framework's performance on the inertial confinement fusion (ICF) [30,9] datasets, which present substantial challenges due to limited availability and high costs of experiments.To demonstrate the effectiveness of our proposed approaches, we build upon the work of Kustowski et al. [5] by using the benchmarks presented in their study.
Datasets: Specifically, we use two datasets in our experiments: The first, referred to as R, stems from real inertial confinement fusion (ICF) experiments conducted during a "Bigfoot" campaign in 2018 [31] at the National Ignition Facility (NIF) in Livermore, California.This multi-modal dataset comprises 10 ICF shots and is accompanied by a large set of simulations produced using a 1D physics simulator [5], denoted as S. The dataset consists of nine scalar inputs corresponding to the design space of the simulator and experiments, ten output scalar values, and an output X-ray image.Most of the inputs relate to the laser energy's conversion into X-rays and its impact on capsule compression, including energy, power, and geometric asymmetry; the other inputs concern hydrodynamic scaling, fuel preheat, and capsule material properties.The 10 scalar outputs capture key phenomena such as the precise moments of peak neutron and X-ray emissions, referred to as "bang times", alongside essential thermodynamic variables like temperature and velocity.Additionally, the dataset includes detailed profiles of X-ray emissions and neutron yields, the latter being a critical indicator of the experimental yield.The overarching aim is to enhance our predictive capabilities, thereby enabling us to maximize the experimental energy yield.The second dataset, denoted as Y, was generated from a multi-modal surrogate [32], previously trained on all the aforementioned simulations.The domain shift here is synthetically induced by obtaining predictions from the surrogate across a disjoint set of input parameters [4].This allows us to test our hypothesis on a much larger set of data (1000 samples in total) to obtain more statistically significant results.Even with the synthetic set, we always assume access only to a very few number of samples for fine-tuning, but here we Scalar ID.Name Evaluation metrics: To assess the efficacy of our proposed methods, we use the Mean Square Error (MSE) as the primary evaluation metric for both scalar and imagebased predictions.Each experimental setup was executed 10 times, with a leave-one-out cross-validation across the 10 available data samples in the real dataset.In each crossvalidation fold, one sample is used for testing, one for validation, and the remaining 8 for fine-tuning.For consistency, we use the same setup in the synthetic dataset (8 train, 1 validation) during fine-tuning and model selection but increase the test set to all the remaining available samples (991).This is repeated 10 times, with the train and val data chosen at random without replacement.

Results:
The aggregated results are presented in Table 2 (top).Our approach is compared against baseline methods on both experimental and synthetic datasets.Across the board, our method demonstrates a substantial reduction in the MSE values for both scalar and image predictions.Specifically, on the experimental dataset, we observed an average reduction of nearly 50% in the MSE, declining from 0.87 to 0.492.For the synthetic dataset, the error rate decreased from 6.580 to 5.155, nearly a 20% improvement.
For a more comprehensive evaluation, we also conducted additional experiments with seven training data points, as shown in Table 2 (bottom), aligning with the experimental setup described in Kustowski et al. [5].In this setting, we trained models using all possible combinations of seven data points, leading to a total of 120 individual experiments.The performance degraded slightly when utilizing fewer training samples, as expected, but our proposed method still significantly outperformed the baseline, exhibiting remarkable gains in predictive accuracy for both scalar and image outputs.
Comparative Statistical Evaluation of Hyper-parameter Selection Strategies For additional experimental evaluation, we use our "leave-3-out" experiments for further statistical analysis shown in Table 2 (top).It is evident that our proposed method consistently outperforms the baseline algorithm.However, to offer a quantitative comparison, we focus on contrasting our Minimum Smoothed Error Graph (hereinafter denoted as GSE min ) with the Traditional Minimum Validation Error (V E min ).A detailed table for the leave-3-out experiment results (and results for leave-one-out with V E min ) can be found in the supplement.
To ascertain the statistical significance of the performance differences between GSE min and V E min , we conducted a series of paired-sample t-tests.For the Mean Squared Errors (MSE) averaged over scalars, the test yields µ 1 = 1.027, µ 2 = 0.631, t = 2.3134, and p = 0.0108, confirming the superiority of GSE min at a 95% confidence level.Similarly, for the average pixel-wise MSE, we find µ 1 = 0.208, µ 2 = 0.189, t = 2.0124, and p = 0.0227, which again corroborates the enhanced performance of GSE min .While the small sample size is relatively small, we emphasize the thoroughness of our approach in partitioning the dataset into all possible configurations, thereby enhancing the reliability of our statistical inferences.

Diagnostic X-ray Images
We commence our discussion with an analysis of the model's efficacy on the reconstructed images, as depicted in Figure 2, in comparison to the baseline method.Our model exhibits a superior ability to approximate the underlying distribution of the training set.In particular, we draw attention to the synthetic image results, which demonstrate a marked reduction in simulation bias in our approach.
Although our generated images display minor artifacts attributable to the use of transformer-based patching techniques, they successfully approximate the overarching geometric structures.It is crucial to note the primary focus of our study lies not in image reconstruction but in the accurate prediction of scalar values.Our dataset is multimodal, comprising diagnostic images and scalar values; however, the latter serve as the principal targets of interest.The notable improvement in the prediction of these scalar attributes for the experimental dataset underlines the practical significance of our approach.Our models' predictions on the held-out test X-ray images after fine-tuning on the real training data, compared to the baseline.Pixels here represent energy outputs of the experimental implosion.White pixels are high energy, purple are lower energy, and black are no energy.Zoom in to better see the results.The MSE for our method is lower than the baseline for the test predictions.While the image quality is not perfect, we find that our new model has modest improvements over the baseline both qualitatively and quantitatively.
Robust Hyper-parameter Optimization via Graph Smoothing Figure 3 elucidates the efficacy of our graph smoothing hyper-parameter optimization method, elaborated in Section 4.6.The primary utility of this method lies in its ability to remap instances characterized by a disparity between validation and test errors into a refined validation error space.By applying this smoothing operation, we uncover regions within the hyperparameter landscape that robustly yield low test errors.
The figure plots validation against test errors for multiple hyper-parameter configurations, thereby empirically demonstrating the algorithm's robustness.Notably, configurations that initially exhibit high test errors, despite low validation errors, are effectively smoothed out.This results in a more reliable selection of well-performing hyper-parameters, as evidenced by the sparsity of such points in the modified validation space.
While our primary results are consistent improvements over the baseline, we take a  Detailed results comparing the masking strategies L masked and L pred , as well as using the smoothed graph validation error rate GSE min versus the non-graph minimum validation error rate V E min .We find the interesting result: masking is useful for the Y dataset, but not for R. Furthermore, using the graph is always an improvement.
deeper look at how our pretraining losses affect the results of our models.We compare two pretraining strategies.The first is forward surrogate modeling prediction loss (L pred ): only predict simulation outputs given simulation inputs.The second is forward loss in addition to masked auto-encoding loss (L masked ): the model randomly sees partial inputs and partial outputs and, then, predicts what it does not see.Furthermore, we take a detailed look at how GSE min and V E min perform when separately analyzing the losses.
The graphs in figure 4 provide insights into the relationship between hyperparameters and model performance, showing an interesting behavior for the use of the masked auto-encoding loss.We see that on the synthetic dataset Y, the use of masking is a large improvement over the pure prediction loss.We also find the opposite to be true for R.This is most likely due to the former dataset's size: It is a much smaller distribution shift between the pretraining dataset and the fine-tuning dataset, such that the learned correlations from L masked can easily be accounted for, whereas the changes in R are so dramatic that deeper correlations learned from L masked result in overfitting.We also highlight that for all our experiments shown in Figure 4, using the smoothed graph validation errors, GSE min consistently results in enhanced performances for all experiments over simply using the minimum validation error V E min .

Effects of Increased Training Data
In an effort to understand the model's performance in data-rich scenarios, we conduct an ablation study utilizing 50 data points for finetuning, as presented in Table 3.As the definition of "few-shot" learning can be ambiguous in the literature, we consider the scenario with 50 points to not be few.Nevertheless, our findings indicate that both V E min and GSE min yield comparable performance, significantly surpassing the baseline.This suggests two critical insights: First, our transformer-based model consistently outperforms the non-transformer baseline.Second, in scenarios with cleaner, less noisy validation data, the graph smoothing operation poses no detriment to model performance.Another effect of additional data is a change in the optimal hyper-parameters.We compare the hyper-parameter configurations selected by all runs between the data-scarce experiments and this relatively data-rich experiments.Our analysis revealed a degree of consistency in hyper-parameters across different data scales, such as identical learning rates and a high number of training epochs.However, variations were observed in the Scalar ID 4. Using CKA to compare the similarity between pretrained feature embeddings and finetuned feature embeddings from one left-out test point.We show how our proposed method for GSE min results in less simulation bias (lower similarity score compared to the pretrained embedded features) for all scalar embeddings when compared against using V E min .
selection of fine-tuning layers and other hyper-parameters, showing the importance of validation metrics within a dataset.
Extreme Case: One-Shot Learning To explore the limitations of our method, we conducted an experiment with only one data point for training and another for validation.As anticipated, the results are markedly sub-optimal; however, the performance of V E min and GSE min is indistinguishable in this extreme setting.This results serves to corroborate that GSE min essentially reduces to V E min when the data becomes extremely sparse.

Analysis of feature Embeddings using CKA
Centered Kernel Alignment (CKA) is a technique used to measure the similarity between two sets of features [33].It has been widely utilized in the context of neural network representations to understand the alignment of features in different layers or networks.
In short, it gives you a similarity between two distributions of features.If the features are identical, then the score will be 1.0; the more the features' distributions deviate, the lower the score will go towards zero.Here we use CKA to compare the features between our two hyper-parameter selection strategies V E min and GSE min .
In Table 4, we show the results of our CKA analysis across the embeddings for all 10 output scalars from our leave-one-out experiments.Specifically, we apply CKA to compare the embeddings from pretrained embeddings to embeddings from V E min , and pretrained embeddings to GSE min .Our analysis demonstrates a clear pattern: the use of GSE min embeddings is consistently lower than the V E min embeddings.These lower scores indicate that GSE min consistently exhibits less simulation bias as compared to the embeddings obtained from V E min .
It is important to note the limitations of CKA, as discussed in recent literature [34], that performance can be influenced by outliers.This sensitivity to outliers implies that while CKA scores provide a useful comparative measure of feature similarity, they should be interpreted with caution.The observed differences in CKA scores, particularly those of minimal magnitude, should be considered indicative of a broader trend towards reduced similarity with the pretrained model rather than definitive evidence of the superiority of one method over another.Our findings suggests that the graph-based method might be a more robust and unbiased approach for generating embeddings in our context.

Discussion
In the current study, we advance the field of few-shot transfer learning in scientific contexts by introducing a novel approach that harnesses the versatility of Transformerbased architectures.Extending this versatility, our model is uniquely equipped to handle multi-modal data, incorporating both scalar and image formats seamlessly.This capability enables the model to predict complex physical systems with significantly less simulation bias.
A crucial part of our strategy is the innovative graph-based hyper-parameter optimization technique.Previous studies have explored few-shot learning and hyperparameter optimization from different angles.For instance, Franceschi et al. [35] introduced a bilevel programming framework for gradient-based hyper-parameter optimization and meta-learning, particularly for deep learning and few-shot learning scenarios.On the other hand, Mazumder et al. [36] developed a robust few-shot learning approach without specifically focusing on hyper-parameter optimization.
In contrast, while Van Rijn and Hutter [37] analyzed the importance of various hyper-parameters, they did not factor in the challenge of untrustworthy validation data, which our work addresses.Liang et al. [38] also recognized the issue of noisy labels in fewshot learning but diverged by choosing to incorporate sophisticated loss functions rather than emphasizing hyper-parameters.Our method, countering traditional challenges such as noisy validation error rates seen in prior work, leads to more reliable and generalizable hyper-parameter configurations that improve overall model performance.Furthermore, Muniraju et al. [39] presented parameterized coverage-based designs for superior sample mining and hyper-parameter optimization, indicating the increasing significance of these concepts in the scientific community.
Beyond optimization, our study's emphasis on surrogate modeling and addressing simulation bias stands on the shoulders of substantial previous research.Surrogate modeling, for example, has seen applications in varied scientific domains, from the rigorous optimization framework for expensive functions used in helicopter rotor blade design by Booker et al. [40] to Bayesian calibration techniques for computer models introduced by Kennedy and O'Hagan [41].In the specific arena of Inertial Confinement Fusion (ICF), the field has witnessed machine learning-driven efforts like that of Hatfield et al. [1], ensemble models from Nora et al. [2], and neural network-based approaches such as those by Kustowski et al. [4] and Kustowski et al. [5].These underline the persistent pursuit to address simulation bias and provide robust models, aligning with our work's objectives.
Building upon these foundations, our work further explores the frontier of predictive modeling within the ICF domain.A critical aspect of this exploration is the acknowledgment of potential radical changes in physical behavior in parts of the design space that remain unexplored experimentally.One such phenomenon, ignition, occurs when the energy generated within the fusion fuel surpasses the energy being lost, leading to a self-sustaining fusion reaction.This represents a drastic shift in the system's response and poses significant challenges for predictive modeling.The complexity of predicting events like ignition, particularly with simulation-based data, highlights the nonlinear and high-stakes nature of these transitions.Our approach, designed to enhance the predictive model's capability across a broad spectrum of conditions, aims to contribute a more comprehensive understanding and optimization of experimental yields in ICF research.By addressing these challenges, we pave the way for breakthroughs in fusion energy.
What sets our work apart is its potential for facilitating multi-modal transfer learning tasks in scientific domains.While the immediate impact of our contributions is evident, this work also lays the groundwork for more expansive research.Future sections will delve into the possibility of applying our methods to other disciplines, thereby widening the scope and impact of our findings.

Formal Definitions
We consider multi-modal physics simulation datasets given by D s = (X , O, I) consisting of input scalars X = {x 1 , x 2 , . . ., x N }, output scalars O = {o 1 , o 2 , . . ., o N }, and output images I = {I 1 , I 2 , . . ., I N }, where N denotes the size of the dataset and d j = (x j , o j , I j ).We also assume access to a "target" dataset D t which is ultimately the domain on which we want our model to be most accurate.We expect D s ̸ = D t , due to the known gap between them.Here, the source domain is typically a simulation dataset collected by sampling from a physics simulator, and the target dataset contains real experimental observations.Consequently, we assume that the number of available target samples is very small, N s >> N t .We use the superscript notation to denote the domain (source vs target) as required, and drop it otherwise for simplicity of notation.
Problem Setup Let us define a surrogate as f s θ : X s → (O s , I s ), where θ are its parameters to be learned.Due to the expected simulation-experiment gap, this model will likely perform poorly when tested directly on target data, i.e., we expect a large error in the prediction since f s θ (x t ) ̸ = (o t , I t ).This gap typically manifests as a task shift, i.e., where the input distribution X remains unchanged but the output distribution has changed significantly between source and target.As a result, the source model must be adapted or fine-tuned using a small number of training examples from D t so that this gap can be closed.

Fine-tuning and model adaptation
The biggest challenge in model adaptation in this context is the lack of sufficient training data.This makes the fine-tuning problem challenging due to two main reasons: (i) Risk of overfitting -While increasingly complex models with a large number of parameters can provide more useful inductive biases to ML surrogates, fine-tuning all the parameters on a very limited dataset will likely result in overfitting.To mitigate this issue, only a part of the network is adapted (typically the final few layers, though not always) while keeping the rest of the parameters fixed.In other words, we can split the parameters as θ s = [β s fixed , β s trainable ], indicating weights that are unchanged and weights that get updated.The fine-tuned model is typically of the form θ * = [β s fixed , β * trainable ], where * indicates the final, fine-tuned model that is used to make predictions.
(ii) Model selection with less validation data -Model selection is the problem of identifying the best set of hyper-parameters based on the performance on a held-out validation set (not seen during training).When the validation set is very small -as is likely the case when available labeled data for fine-tuning itself is very sparse -the best performing model on the validation set is unlikely to be the best performing model on the real test, due to very noisy estimates arising from very poor sampling of the validation set.As such, picking a model that is likely to generalize well is very challenging.
In the methods section, we outline our solution to both of these problems and show how the proposed transformer-based surrogate and model selection strategy are effective in addressing the simulation-experiment gap.

Masked training with Transformer Surrogates
Our first, and one of two main contributions, is the use of transformer models [10] as surrogates in the ICF application space.Transformers are a class of general-purpose learners that work on tokenized forms of data (such as patches or chunks) and learn attention across arbitrary data modalities [42].This enables them to capture important correlations on their own and, equally important, the architecture makes very few assumptions about the data.This phenomenon has led to successes in a variety of applications, such as computer vision [16] and other multi-modal data [43].In particular, we explore the use of masked training in transformers using the Masked Auto-Encoder (MAE) [29].Inspired by the successes of masked pre-training in language modeling, the MAE presents a pre-training strategy for image data that was a significant breakthrough in self-supervised representation learning for image data.We extend the MAE strategy from just one modality (text or image) to multiple modalities.
In order to effectively leverage masked autoencoding, we have to employ a deep transformer-based model.A diagram of our model is shown in figure 5.  deep transformer-based neural network, we can effectively capture these correlations, in addition to including X in the learned representation.Therefore, we introduce a more general version of f by incorporating multiple strategies from our novel masking framework which we define as follows.Let M = (M f orward , M random ) be a set of masking functions, each of which takes as input a data sample d j and returns only some element, i.e., o j , I j = M f orward (d j ) corresponding to a standard forward surrogate model o j , I j = M f orward (d j ) = f surr (x j ).We also note the inverse of a mask M to be the opposite of said mask, i.e., x j = Mforward (d j ).Our other masking strategies M random randomly selects from all elements of a data sample to mask at a fixed random rate (75% in our case).We emphasize that while our task only requires the two masking strategies, other strategies can be defined for other data representation tasks (such as an inverse mask), hence the flexibility of our framework.

Generalized
The general model f s θ is a deep transformer-based neural network that can take as input all scalars and images, correspondingly masked by a desired mask M , then, outputs all scalars and images for a given sample j: The mask enables flexible training of either: a standard surrogate style with only output prediction by using mask M f orward , or for standard masked auto-encoding training where inputs are randomly selected to be masked M random .
(i) We convert our data into embeddings, as all transformer-based operations deal with embeddings rather than raw data.
(ii) We encode the scalars into an embedding by simply multiplying a trainable embedding by the normalized (0-1) scalar.
(iii) We follow standard practice [16] by flattening image patches and learning a shared image embedding space by multiplying each patch by a learnable matrix W p .
(iv) For each embedding we add a positional encoding embedding.The image embeddings get a fixed 2d-sinusoidal encoding, whereas the scalars get a simple trainable encoding added.
(v) Our transformer model is split into two parts: the encoder and the decoder.
(vi) Each part is comprised of multiple transformer layers: Multi-Head Self-Attention, Layer Normalization [44], and a Feed-forward Neural Network.
(vii) The outputs of the encoder are combined with a series of mask tokens embeddings, depending on the masking strategy, and are fed into the decoder network.
(viii) The output of the decoder are prediction embeddings corresponding to all the data.These embeddings are either multiplied by an individual learnable prediction vector (for scalars) or by a shared prediction matrix (for images).
During both masked and surrogate forward passes, only the available data are embedded for the encoder to process.After being encoded, a "missing" data embedding is placed in the location of all the missing data.This embedding has a new positional encoding added to it (still fixed for the image embeddings).All those embeddings are passed through the decoder transformer layers to get output embeddings.A learnable inverse transformation is done on all the image patches, and each scalar has its own output embedding e k and a learnable output embedding vector map (e.g., ŷj = W k * e k ).

Simulation Pretraining
We investigate training our surrogate through two types of pretraining losses based on output prediction and masked prediction.The output prediction loss is a standard L2 loss on the outputs of a given example j when using M f orward : where γ i is a hyper-parameter tuned on the validation set of S and γ o = 1.
For masking loss, we convert the image into 16 equally-sized square embeddings, along with 19 scalar embeddings.We then remove 75% of those embeddings from the input to f s θ using M random and predict the values of the masked inputs, resulting in a mask loss defined as: The overall pretraining loss combines the prediction loss and masked auto-encoding loss, controlled by a hyper-parameter α: Here, α is a hyper-parameter tuned only on the simulation dataset.We found that setting α = 0 (corresponding to no prediction loss) produces consistently poor results during the fine-tuning stage.And we also found α = 1 to have inconsistent results, and therefore we treat α as a hyper-parameter passed down to our fine-tuning (either α = 1 or an optimized α of 0.02).

Experimental Data Fine-tuning
Due to the limited amount of data available, we must exercise caution when modifying the parameters of our pretrained model.We find that updating only a few parameters (i.e., layers) is effective.As discussed in Kustowski et al. [5], updating a single layer of the neural network, rather than all the parameters of the model, is essential to avoid overfitting.
To fine-tune our model f θ s on the experimental dataset R, we employ a leave-oneout cross-validation strategy, given the small size of our dataset D t which consists of N = 10 samples.In this process, we use 9 samples for training and 1 sample for testing.
During training, we compute a validation error by performing another round of leaveone-out validation, where we fine-tune a model on 8 of the 9 training points, and then evaluate on the held-out point.
As defined above, we specify a fully train model to be Where δ is the learning rate to update just the trainable parameters β j and we use the L pred loss function (2) and set γ o = 0 or γ i = 0.This zeroing out to focus on one modality is employed to avoid overfitting on scalars at the expense of images or vice versa, degrading the model's overall performance.By focusing on each modality individually, we ensure that the model can learn and capture the unique characteristics of each data type without being negatively influenced by the other.We investigated finetuning our model on both the images and scalars simultaneously.However, we found that this approach resulted in inferior performance compared to the separate training on images and scalars.
Finally, we repeat this process for all 9 training points, then average the error of the held-out validation points V = 1 N N j=1 V j .This approach allows us to systematically evaluate the model's performance across different experimental data splits while making the best use of the limited available data.
Hyper-parameter Grid Search During the fine-tuning process, we perform a grid search over a range of hyper-parameters.The aim of the grid search is to identify the optimal combination of hyper-parameters that yield the best performance on the validation set.Some of the hyper-parameters explored during the grid search include the learning rate, the number of fine-tuning epochs, and determining which layer to tune.By exhaustively searching the grid over the hyper-parameter space, we ensure an optimal model can be selected for a given training set.
Finally, due to the few-shot nature of our data, we fine-tune our model on both the training and validation data using the selected hyper-parameters.After the fine-tuning process is complete, we evaluate the performance of our model on the held-out test set.This provides us with an estimate of the model's generalization capability when applied to unseen experimental data.
Early Stopping Post-Hoc Correction As we often stop fine-tuning a model before it finds a local minimum of the loss function, we found these models to consistently underfit the training data.To counteract this deficiency in model fit, we propose a method that manually adjusts the bias and variance of the predictions in accordance with the training set.The primary idea behind this approach is to strike an optimal balance between overfitting (halting training when the prediction loss ceases to decrease) and underfitting (insufficient updates to the model weights to account for bias).We suggest a straightforward solution that involves manually modifying the model's final predictions using new variance and bias parameters.
We compute the average error from the training data for each predicted scalar , where f (x j ) k represents the k-th scalar output of the finetuned model f t θ , and adjust the final validation set predictions to account for this average error over the n training points: A similar approach is applied to the variance of the predictions.Let the average for scalar k be µ k = 1 N n j f surr (x j ) k and the variance be σ(y k ) = var(y k 0 , y k 1 , ..., y k n ):

Implementation and Dataset Details
In our implementation, we employ a variant of the Masked Autoencoder (MAE) that closely aligns with the popular architecture proposed by He et al. [29], albeit with modifications to suit our multi-modal dataset and computational constraints.Specifically, our MAE model is characterized by a reduced number of decoder blocks (6) and smaller embedding sizes, with 512 dimensions for the encoder and 256 dimensions for the decoder.Our decision to opt for a smaller model was chosen based on empirical evidence from preliminary experiments and is also informed by the broader observation in the field that, beyond a certain point, larger embedding sizes do not equate to significant performance improvements, particularly for datasets of moderate size and dimensionality.The hardware used for training comprised a single NVIDIA V100 GPU, with hyper-parameter tuning and experimentation facilitated by parallelization across a cluster of 64 V100 GPUs.The Adam optimizer is employed with a cosine annealed learning rate starting at 10 −3 (which gradually decreases to 0).The best model is selected based on the pre-training simulation test set average error rate (optimized over different hyperparameters: γ o , epochs, and learning rates).For each leave-one-out test set experiment, we select the best smoothed validation score as described above.
For the R dataset, the large simulation database, S, was created using the twodimensional radiation hydrodynamic code HYDRA [45].These simulations serve as an extensive sampling of the design space, permitting more robust predictive modeling.
Our second dataset, Y, is generated synthetically.It is designed to create a representative set of ICF experiments by employing an uncalibrated surrogate model.Instead of running new HYDRA simulations, which would be computationally expensive and time-consuming, Kustowski et al. [5] utilized their uncalibrated surrogate model to make predictions.This approach enabled them to create two lower-dimensional and physically inconsistent datasets for transfer learning, which are nearly equivalent to running a new set of simulations.To create the synthetic datasets, they fixed four of the nine input parameters and sampled the remaining five input parameters randomly within their original ranges.They used their uncalibrated surrogate model to predict the outputs and, then, perturbed the values of the asymmetry and preheat parameters to create 1,000 "experiments".
The pretraining simulation dataset comprises of 90000 training samples and 2000 test samples.Images are X-rays of 60x60 pixels and are self normalized, with each image's pixels divided by its own mean, as each image may span differing orders of magnitude.The experimental dataset consists of 10 samples, which are divided using leave-one-out for training.The synthetic dataset includes 1000 samples.To stay consistent with the experimental dataset, we only fine-tune with a few samples (7,9, or 50); and we report the average error over the other held-out points.

A novel Graph-Based approach for Robust Hyper-parameter Selection
Given a set of candidate hyper-parameter configurations, we construct a graph G = (V, E), where each node v i ∈ V represents a unique hyper-parameter configuration λ i , and an edge (v i , v j ) ∈ E exists if the corresponding configurations differ in exactly one hyper-parameter setting by a single step in that hyper-parameter.For example, an edge would exist between two hyper-parameter configures if they only differ in the learning rate by one step (e.g., 10 −3 or 10 −4 ).There would not be an edge between a two-step size difference, such as 10 −3 and 10 −5 .In addition, there would not be an edge if two hyper-parameters were changed; for example, if the learning rate and the epochs to train were different between two fine-tuning runs, then no edge would be between the two nodes corresponding to these two hyper-parameter configurations.This graph helps us understand the local structure of the hyper-parameter space and how small changes in the configurations are related.
The hyper-parameters we use are as follows: (i) Transformer decoder block to train (1-7) In our study, the determination of hyper-parameter grid points for exhaustive scans was initially guided by a trial-and-error approach, resulting in a comprehensive exploration across 6,048 hyper-parameter configurations for each experiment.Recognizing the potential inefficiencies of this method, we propose a more systematic approach for future work and practitioners aiming to optimize the hyper-parameter selection process.Specifically, employing Bayesian optimization offers a promising starting point for identifying promising regions within the hyper-parameter space.This probabilistic model-based approach can effectively suggest initial values that are likely to yield improved performance metrics.Following the identification of these regions, an exponential or binary search strategy could be implemented to refine the grid resolution.
Validation error rates are separately computed for both images and scalars.The error for an image is simply the MSE averaged over all the pixels, and the error for the scalars is the MSE averaged over the ten target scalars.The validation error is the average error from doing a leave-one-out cross validation on the training set.We separate the validation error rates between images and scalars to keep in line with our separate training process described above.For the sake of clarity, the following description only considers a single validation score (e.g., the image MSE).We assign node values based on the validation error rates, denoted by V = {V 1 , . . ., V n }, where V j = 1 N N j=1 V j corresponds to the validation error rate for the hyper-parameter configuration λ j averaged over the N leave-one-out experiments for a given training set.The minimum validation error rate configuration is defined as: Next, to exploit the graph structure for hyper-parameter optimization, we perform a simple smoothing on the graph G.This process updates the node values by considering both the original validation error rate and the average value of neighboring nodes.
Let A be the adjacency matrix of the graph G, and N (i) denote the set of neighbors of node i.We define the smoothed node value Ṽi as follows: where A ij denotes the element of the adjacency matrix at position (i, j).The first term in the equation represents half of the original validation error value, while the second term represents half of the average neighbor value.
After applying the smoothing, we select the hyper-parameter configuration corresponding to the node with the lowest smoothed value: The selected configuration GSE min represents an optimal choice that balances the original validation error rates and the information propagated from neighboring nodes.This graph-based approach is particularly beneficial in the context of few-shot learning, where the limited number of examples can lead to noisy estimates of model performance.By exploiting the structure of the hyper-parameter space, our method effectively identifies optimal hyper-parameter configurations and consistently improves the overall performance for our few-shot scenario.Our proposed design is based on the premise that we have a comprehensive grid search over the hyper-parameters of interest.This choice of exploration strategy lends itself naturally to the construction of the graph, where each node represents a unique hyper-parameter configuration and edges connect nodes that differ in exactly one dimension by a single parameter step.This approach results in a well-defined neighborhood structure that captures the local similarities between configurations.However, it is important to note that more complex neighboring strategies could be employed when dealing with more sophisticated hyperparameter sweep settings, such as random search or Bayesian Optimization [46].In such cases, alternative techniques for defining the connectivity between nodes might be required to capture the relationships between different configurations.
In our analysis, we focus on using the fewest neighbors possible in order to balance the exploitation of the graph structure and the preservation of the original validation error rates.This choice is motivated by the desire to avoid over-smoothing, which can lead to suboptimal hyper-parameter configurations.

Figure 1 .
Figure 1.Our method is separated into three distinct stages: First, pretraining on simulation data with masked autoencoding and surrogate losses.Second, finetuning our model on the experimental data with a hyper-parameter sweep.Finally, finding the best hyper-parameter settings using our novel graph-based selection.

Figure 2 A
Figure 2

Synthetic
Figure2.Our models' predictions on the held-out test X-ray images after fine-tuning on the real training data, compared to the baseline.Pixels here represent energy outputs of the experimental implosion.White pixels are high energy, purple are lower energy, and black are no energy.Zoom in to better see the results.The MSE for our method is lower than the baseline for the test predictions.While the image quality is not perfect, we find that our new model has modest improvements over the baseline both qualitatively and quantitatively.

Figure 3 .
Figure 3. Hyper-parameter graph smoothing ensures optimal model selection based on noisy validation error.(Left:) Validation versus test error rates for scalar predictions experiment data.(Right): Proposed graph-smoothed validation error vs test error.We highlight the minimum validation error and the minimum smoothed validation error finding that our smoothing removes the noisy data to find a robust, well performing hyper-parameter selection.

Figure 4 .
Figure 4. Detailed results comparing the masking strategies L masked and L pred , as well as using the smoothed graph validation error rate GSE min versus the non-graph minimum validation error rate V E min .We find the interesting result: masking is useful for the Y dataset, but not for R. Furthermore, using the graph is always an improvement.

Figure 5 .
Figure 5. Masked Pre-training: Our novel multi-modal architecture leverages both images and scalars as inputs and outputs for a transformer-based deep neural network.Transformers enable straightforward surrogate models as well as effective representation learning through masked autoencoding.

Table 3 Table 1 .
Table of Experiments.

Table 2 .
The

Table 3
. Graph smoothing converges to standard model selection when more data is available: Here, we use 50 synth train and 10 validation examples.Once again, our method outperforms the baseline significantly.Reassuringly, we note that with increased availability of training and validation data, our GSE min approach converges to standard model selection based on minimal val error V E min .
Surrogate Model with Flexible Masking Strategies While a traditional surrogate model is often defined as (o j , I j ) = f surr (x j ), in this work we explore a new formulation in order to capture correlations.Prior methods are designed around learning a representation that captures the correlations between Y and I by learning a compressed representation jointly.However, in this work, by utilizing a Transformer Decoder