Organizations: AOT-IC Group, Los Alamos National Laboratory, Los Alamos, NM, USA · CAI Division, Los Alamos National Laboratory, Los Alamos, NM, USA · HPC Division, Los Alamos National Laboratory, Los Alamos, NM, USA
PDE foundation models are typically pretrained on large, diverse corpora of PDE datasets and can be adapted to new settings with limited task-specific data. However, most downstream evaluations focus on forward problems, such as autoregressive rollout prediction. In this work, we study an inverse problem in inertial confinement fusion (ICF): estimating system parameters (inputs) from multi-modal, snapshot-style observations (outputs). Using the open JAG benchmark, which provides hyperspectral X-ray images and scalar observables per simulation, we finetune the PDE foundation model and train a lightweight task-specific head to jointly reconstruct hyperspectral images and regress system parameters. The fine-tuned model achieves accurate hyperspectral reconstruction (test MSE 1.2e-3) and strong parameter-estimation performance (up to R^2=0.995). Data-scaling experiments (5%-100% of the training set) show consistent improvements in both reconstruction and regression losses as the amount of training data increases, with the largest marginal gains in the low-data regime. Finally, finetuning from pretrained MORPH weights outperforms training the same architecture from scratch, demonstrating that foundation-model initialization improves sample efficiency for data-limited inverse problems in ICF.
Figures & tables
Fig. 1 : Schematic illustrating how the PDE foundation model (MORPH) latent representation is coupled with a lightweight task-specific head. Hyperspectral images are provided as input to the foundation model, which is fine-tuned for reconstruction. The transformer-block outputs are passed to a task-specific head (a dense neural network). In parallel, the task-specific head (TSH) also ingests 15 scalar observables/diagnostics and is trained to predict a 5D parameter output. The foundation model and TSH are trained jointly end-to-end, using separate loss functions with independent optimizers and learning-rate schedulers.
Fig. 2 : Ridge-regression sensitivity map for predicting the five parameters from multi-modal features. Columns show standardized image PCA scores (PC1–PC32) and scalar diagnostics (scalar0–scalar14), separated by the vertical line; colors denote signed coefficient magnitude. The strongest dependencies are scalar-driven for param1 and param2 , while param0 and param3 exhibit near-zero coefficients, indicating poor identifiability.
Fig. 3 : Reconstruction results on four held-out test samples (a-d). For each example, the top row shows the ground-truth four-channel hyperspectral image, and the bottom row shows the corresponding reconstructed (predicted) image. Across the full test set, the recorded MSE between true and reconstructed images is 0.0012.
Fig. 4 : Parameter estimation results on the test set. Scatter plots compare ground-truth versus predicted values for each of the three predicted parameters. We achieve strong agreement for three parameters: Param1 attains R2=0.975 with L2=0.035 , Param2 attains R2=0.995 with L2=0.013 , and Param4 attains R2=0.990 with L2=0.022 .
Fig. 5 : Data-scaling study showing training and validation loss curves (log-log scale) as the fraction of available training data is varied (5%, 10%, 25%, 50%, 75%, and 100 %). (a) MORPH reconstruction training loss, (b) task-specific head (TSH) parameter-regression training loss, (c) MORPH reconstruction validation loss, and (d) TSH parameter-regression validation loss, all plotted versus epochs (log-scaled axes).
Fig. 6 : fine-tuned model vs trained-from-scratch: Test-set parameter-regression loss of the TSH versus training-data fraction, comparing MORPH finetuning from pretrained weights to training the same architecture from scratch. Finetuning achieves lower loss, with the largest gains in the low-data regime (e.g., 10%), and the gap narrows as more training data becomes available.
Department Modeling and Artificial Intelligence Fraunhofer Institute for Integrated Systems and Device Technology IISB Schottkystrasse 10, 91058 Erlangen, Germany