FIRE: Multi-fidelity Regression with Distribution-conditioned In-context Learning using Tabular Foundation Models

Rosen Ting-Ying Yu; Nicholas Sung; Faez Ahmed

cs.LGJan 29, 2026

FIRE: Multi-fidelity Regression with Distribution-conditioned In-context Learning using Tabular Foundation Models

Authors: Rosen Ting-Ying Yu, Nicholas Sung, Faez Ahmed

Organizations: Center for Computational Engineering, Massachusetts Institute of Technology, Cambridge, MA, USA · Department of Mechanical Engineering, Massachusetts Institute of Technology, Cambridge, MA, USA

Abstract

Multi-fidelity (MF) regression often operates in regimes of extreme data imbalance, where the commonly-used Gaussian-process (GP) surrogates struggle with cubic scaling costs and overfit to sparse high-fidelity observations, limiting efficiency and generalization in real-world applications. We introduce FIRE, a training-free MF framework that couples tabular foundation models (TFMs) to perform zero-shot in-context Bayesian inference via a high-fidelity correction model conditioned on the low-fidelity model's posterior predictive distributions. This cross-fidelity information transfer via distributional summaries captures heteroscedastic errors, enabling robust residual learning without model retraining. Across 31 benchmark problems spanning synthetic and real-world tasks (e.g., DrivAerNet, LCBench), FIRE delivers a stronger performance-time trade-off than seven state-of-the-art GP-based or deep learning MF regression methods, ranking highest in accuracy and uncertainty quantification with runtime advantages. Limitations include context window constraints and dependence on the quality of the pre-trained TFM's.

Figures & tables

Figure 1: Overview of FIRE. The method decomposes multi-fidelity regression into a low-fidelity base inference and a residual correction step, conditioned on the base model’s statistical and distributional information.

Figure 2: Overall performance. Elo rating and average rank for (a) accuracy performance (NRMSE) and (b) uncertainty calibration (NLL) across 31 problems. FIRE outperforms the other baselines in both accuracy and uncertainty.

Figure 3: Pareto frontier analysis. Pareto plots of accuracy (Left) and uncertainty (Right) vs. runtime. FIRE dominates the Pareto frontier for both metrics, being the fastest and most accurate.

Figure 4: Ablation study of FIRE components. We compare the Elo rating of various configurations against NARGP, the strongest GP baseline. The full FIRE framework (Top), which combines residual learning with statistical augmentation, significantly outperforms partial implementations.

Figure 5: Performance (NRMSE and NLL) for representative synthetic (Currin, HD 50), physics-based (Car, Concrete), and HPO (Higgs) tasks. FIRE maintains top-ranked accuracy and uncertainty calibration in the extreme low-data regime (2–5% HF ratio), where baseline performance degrades.

Figure 6: Performance under varying data imbalance. Aggregated results for problem domain as a function of the High-Fidelity to Low-Fidelity sample ratio (

N_{\text{HF}}/N_{\text{LF}}

%): (a)Overall (all 31 problems), (b)Synthetic, (c)Engineering, (d)HPO. FIRE demonstrates superior robustness compared to other baselines when HF data is extremely scarce (

2\%-5\%

), confirming the efficacy of TFM’s zero-shot adaptation, especially for the real-world physics simulation and HPO tasks.

Figure 7: Heteroscedastic cross-fidelity discrepancy on the Forrester(

d=1

) benchmark (right). Regions of larger LF uncertainty

\sigma_{\theta}^{2}

coincide with larger and more variable HF residuals

r

, motivating uncertainty-conditioned residual modeling (left).

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Ablation on the number of surrogates. We compare the standard bi-level FIRE (using 2 TFMs) against a fully recursive variant modeling every fidelity step (using $T-1$ TFMs) and the NARGP baseline. The bi-level approach achieves superior trade-offs between accuracy and efficiency, avoiding the prohibitive runtime costs of the recursive strategy.

Figure 9: Impact of surrogate choice on performance. Replacing the TFM with a standard GP (FIRE(GP)) still yields state-of-the-art results compared to existing GP baselines, confirming that our residual learning framework drives significant gains. However, the performance gap between FIRE(TFM) and FIRE(GP) demonstrates the additional value provided by the TFM’s in-context prior.

Figure 10: Comparison against specialized HPO TFM surrogates. We benchmark FIRE against state-of-the-art PFNs explicitly designed for learning curve extrapolation and Bayesian optimization (FT-PFN, LC-PFN) on the LCBench dataset. Despite FIRE using a general-purpose tabular foundation model, it significantly outperforms these specialized architectures in NRMSE and ranks among the top on the 6 HPO datasets in this study.

Figure 11: Benchmark of FIRE instantiated with different TFM. We compare FIRE using RealTabPFN, TabPFNv2.5, TabPFNv2, and Mitra against the strongest GP-baseline (NARGP). While all TFM-based variants outperform the GP baseline, TabPFN-based backbones yield the strongest results.

Figure 12: Robustness to non-nested data. We compare FIRE against autoregressive GP baselines (NARGP, AR(1), ResGP) on nested ( $X_{HF}\subseteq X_{LF}$ ) versus disjoint ( $X_{HF}\not\subseteq X_{LF}$ ) input designs. Elo ratings (left) are calibrated to FIRE=3000. While baselines benefit from nesting, FIRE achieves the highest Elo and lowest rank in both settings, demonstrating superior generalization to disjoint data.

Figure 13: Impact of surrogate choice for base $f_{\theta}$ and residual $\delta_{\phi}$ modeling. We compare combinations of GP and TFM backbones on their accuracy (NRMSE). Red bars indicate a TFM residual model, while navy bars indicate a GP. Hatching indicates a mixture of models in the FIRE framework. The TFM’s ability to model residuals via ICL provides a distinct advantage over standard kernels.

Figure 14: Impact of post-hoc ensembling on FIRE. Ensembling the residual model improves predictive accuracy (Elo and rank of NRMSE) but increases inference time significantly.

Figure 15: Residual model choice for $\delta_{\phi}$ . With the TabPFNv2.5 base fixed, we replace the TFM residual with scratch-trained GP, RF, XGBoost, and LightGBM models (Elo calibrated to FIRE=1000). The TFM residual achieves the best accuracy, confirming that training the residual from scratch does not improve FIRE’s Pareto frontier.

Figure 16: Uncertainty quantification analysis. (a) Reliability diagram (empirical vs. nominal coverage; the dashed diagonal denotes perfect calibration) and (b) average rank of Interval Score (lower is better) across all benchmarks. Both FIRE variants are well-calibrated and top-ranked on IS, with FIRE-GP achieving the best calibration and FIRE-TabPFNv2.5 remaining competitive.

Figure 17: Overall aggregate performance across 31 problems. Performance is aggregated via Elo, Normalized Score, Average Rank, and raw averages for NRMSE, NLL, and $R^{2}$ . FIRE dominates all aggregation schemes, confirming its status as the state-of-the-art method for accuracy and uncertainty quantification.

Figure 18: Critical Difference (CD) diagrams based on the Nemenyi post-hoc test (

\alpha=0.05

) for (a) NRMSE, (b) NLL, and (c)

R^{2}

. Methods connected by a horizontal bar are not statistically significantly different. FIRE ranks first statistically across all three metrics, specifically forming a distinct significance group from the GP baselines in accuracy (NRMSE).

Figure 19: Performance vs. Data Imbalance. Aggregate metrics (Elo, Normalized Score, Rank, Raw) plotted against the HF/LF sample ratio. FIRE exhibits high stability in low-data regimes (2-5%), significantly outperforming baselines which require larger HF budgets to converge.

Figure 20: Runtime scaling with input dimensionality. This figure reports wall-clock runtime on the HD benchmark suite as the input dimension

d

increases from 10 to 50. FIRE has the lowest runtime at all dimensions. GP and deep learning baselines take longer because they rely on kernel fitting or extensive training.

Figure 21: Pairwise win-rate comparison for all problems. The heatmap displays the percentage of tasks where the algorithm on the Y-axis outperforms the algorithm on the X-axis. FIRE consistently beats all baselines across the majority of tasks (higher numbers are better).

Figure 22: Synthetic Benchmarks: Performance vs. data imbalance. Aggregated metrics vs. HF/LF ratio for 18 synthetic tasks. While GP-based methods like AR(1) remain competitive, FIRE maintains top-tier performance, particularly in the low-data regime (

2-5\%

) for NLL and

R^{2}

Figure 23: Physics-based Simulations Benchmarks: Performance vs. data imbalance. Aggregated results for the 7 engineering and physics problems. FIRE demonstrates a clear advantage here, significantly outperforming all baselines in NRMSE and NLL.

Figure 24: Hyperparameter Optimization (HPO) tasks: Performance vs. data imbalance. Aggregated results for the 6 LCBench learning curve datasets. FIRE outperforms all baselines, validating its effectiveness in modeling the irregular response surfaces typical of hyperparameter optimization landscapes.

Figure 25: Visualizations of classic heteroscedastic regression benchmarks.

	Goldberg et al. (1997)	Yuan and Wahba (2004)	Williams (1996)
FIRE-GP without distribution info $\sigma^{2},\{q^{(\tau)}\}$	0.1933 $\pm$ 0.0289	0.1844 $\pm$ 0.0328	0.1110 $\pm$ 0.0328
FIRE-GP	0.1867 $\pm$ 0.0341	0.1834 $\pm$ 0.0328	0.1082 $\pm$ 0.0323
FIRE without distribution info $\sigma^{2},\{q^{(\tau)}\}$	0.1710 $\pm$ 0.0317	0.1707 $\pm$ 0.0117	0.1076 $\pm$ 0.0320
FIRE	0.1699 $\pm$ 0.0322	0.1696 $\pm$ 0.0105	0.1073 $\pm$ 0.0321

Appendix

Table 1: Predictive Accuracy (NRMSE) on Heteroscedastic Benchmarks (lower the better, bold is the best).

	Goldberg et al. (1997)	Yuan and Wahba (2004)	Williams (1996)
FIRE-GP without distribution info $\sigma^{2},\{q^{(\tau)}\}$	5.4922 $\pm$ 2.9140	9.2186 $\pm$ 5.435	21.797 $\pm$ 18.610
FIRE-GP	3.2145 $\pm$ 0.7244	6.9235 $\pm$ 3.456	17.366 $\pm$ 12.476
FIRE without distribution info $\sigma^{2},\{q^{(\tau)}\}$	1.5759 $\pm$ 0.1027	1.6433 $\pm$ 0.1705	2.3486 $\pm$ 3.3993
FIRE	1.5631 $\pm$ 0.1031	1.6359 $\pm$ 0.1884	1.5039 $\pm$ 2.3797

Appendix

Table 2: Uncertainty Calibration (NLL) on Heteroscedastic Benchmarks (lower the better, bold is the best).

Name	Type	Vaules
batch_size	integer	$[2^{4},2^{9}]$ (log)
learning_rate	continuous	$[10^{-4},10^{-1}]$ (log)
max_dropout	continuous	$[0,1]$
max_units	integer	$[2^{6},2^{10}]$ (log)
momentum	continuous	$[0.1,0.99]$
max_layers	integer	$[1,5]$

Appendix

Table 3: The hyperparameters for LCBench datasets.

Figure 26: Quasi-2D low-fidelity pipeline for the CAR dataset. Top-left: projected mid-plane view of the 3D DrivAer-style geometry after removing wheels. Top-right: cleaned 2D silhouette obtained from the projected points. Bottom-right: Gmsh-generated quasi-2D mesh (shown as a zoom-in around the car; the full domain extends several car lengths upstream/downstream and vertically) with refinement concentrated near the body. Bottom-left: sample incompressible RANS solution from the simpleFoam run used to compute the LF drag coefficient.

Dataset	Data Source	Dimension (Features, $d$ )	Fidelity Level ( $T$ )	$\{N_{t=1},\,N_{t=2},\,$ $\cdots,\,N_{t=T-1}\}$	$N_{T}$
Bohachevsky	Synthetic (MF2)	2	2	{200}	{4, 8, 10, 20, 40, 50}
Booth	Synthetic (MF2)	2	2	{200}	{4, 8, 10, 20, 40, 50}
Borehole	Synthetic (MF2)	8	2	{200}	{4, 8, 10, 20, 40, 50}
Branin	Synthetic (MF2)	2	2	{200}	{4, 8, 10, 20, 40, 50}
Currin	Synthetic (MF2)	2	2	{200}	{4, 8, 10, 20, 40, 50}
Forrester	Synthetic (MF2)	1	2	{200}	{4, 8, 10, 20, 40, 50}

Appendix

Table 4: List of datasets

Figure 27: Result of Physics-based simulation problems (Car, Concrete, BWB).

Figure 28: Result of Physics-based parametric problems (Wing, Beam, HOIP).

Figure 29: Result of Hyperparameter Optimization (HPO) problems from LCBench.

Figure 30: Result of synthetic problems (Bohachevsky, Booth, Borehole, Branin, Branin3f, Currrin).

Figure 31: Result of synthetic problems (Forrester, Hartmann3f, Hartmann6, Himmelblau, Park91a, Park91b).

Figure 32: Result of synthetic problems (HD 10

\sim

HD 50, sixhump-camelback).

Multifidelity Modeling

Multifidelity modeling combines inexpensive approximate models with costly accurate simulations or measurements to improve prediction and decision-making.

See all 36 papers

May 31, 2026

On the Uncertainty Quantification Ability of Tabular Foundation Models

Foundation models (FMs) have achieved substantial success in generalizing across tasks without problemspecific training or fine-tuning. However, many critical applications in mechanics and computational science require not only accurate predictions but also reliable uncertainty quantification (UQ). Herein we investigate the UQ capabilities of tabular FMs in regression tasks through a comprehensive empirical study comparing Tabular Prior-Data Fitted Networks (TabPFN) against Gaussian processes (GPs). We systematically evaluate these two methods across a host of regression problems with varying complexity, dataset sizes, and input dimensionalities. We use a default setting to build all the GPs and for a fair comparison against TabPFN v2.5. Our findings highlight an important trade-off between explicit and learned priors: while TabPFN achieves highly competitive performance for complex, high-dimensional problems with sufficient data, GPs often provide superior predictive accuracy and UQ in data-scarce settings. Moreover, when the chosen kernel constitutes a good prior for the underlying function, GP performance can substantially exceed that of TabPFN. Our results can be reproduced from https://github.com/kianswarehouse/GPvsPFN.

Tyler R. Johnson, Kian Ben-Jacob, Nima Negarandeh +2

Department of Mechanical and Aerospace Engineering, University of California, Irvine · Department of Civil and Environmental Engineering, University of California, Irvine

Tabular Foundation Models Uncertainty Quantification Prior-Data Fitted Networks Tabular Regression Gaussian Process Regression

Ask PDF

Aug 19, 2026

GEAR: Generative Expansion and Real Anchoring for Two-Stage Distillation of Tabular Foundation Models

Tabular foundation models (TFMs) achieve strong performance through in-context learning, but context-dependent inference imposes substantial latency and memory costs, hindering large-scale deployment. We propose GEAR (Generative Expansion and Real Anchoring), a modular two-stage framework that distills TFMs into lightweight MLP or tree-based predictors that can be deployed on commodity CPUs. Stage 1 uses synthetic covariates solely as teacher-query locations and trains the student on soft TFM targets, expanding coverage beyond observed rows. Stage 2 re-anchors the student to the target distribution using real labels and out-of-fold teacher predictions, whitch avoids self-labeling leakage. We further derive a risk certificate characterizing the trade-off between generated-query volume and generator fidelity. Experiments on TALENT and TabArena demonstrate the broad applicability of GEAR. Two-stage MLPs outperform supervised MLPs by 1.81--2.00 AUC points on binary tasks and 1.19--1.35 points on multiclass tasks, with additional gains over real-data-only distillation of 1.76--2.19 and 2.09--2.40 points, respectively. On binary tasks, the gains also transfer to LightGBM and XGBoost, and all three student families outperform CatBoost, the strongest non-TFM baseline, in mean AUC. Ablations show gains beyond longer training or alternative warm starts, greater stability from staged than mixed optimization, and generator-dependent diminishing returns as query volume increases. Finally, GEAR reduces median inference time by 57--2866 times and peak prediction memory by 1.9--3.3 times, while retaining higher AUC than matched supervised baselines.

Qi Qin, Jiajie Zhu, Dali Chen +6

Renmin University of China, Center for Applied Statistics and School of Statistics, Beijing, China · Ant Digital Technologies, Ant Group, Hangzhou, China · Nanjing University, Nanjing, China

Knowledge Distillation Tabular Foundation Models Cost-Aware Inference Tabular ICL

Ask PDF

May 13, 2026

VIP-COP: Context Optimization for Tabular Foundation Models

Tabular foundation models (TFMs) have emerged as a powerful paradigm for in-context learning on structured data, enabling direct prediction on new tabular tasks without task-specific training. However, their effectiveness is constrained by context length limits, restricting application to medium-scale data and degrading performance when inference-time data exceed pretraining size distributions. Our work introduces VIP-COP, estimating the Value of Importance for Prediction of training examples and features for hard Context OPtimization for TFMs. Its explicit selection mechanism suppresses noise and isolates influential data, enabling the model to also benefit from data augmentation by prioritizing high-value augmented samples and features. VIP-COP is (i) fast, boosting performance often within minutes of optimization, based on an online KernelSHAP-based regression with iterative refinement, value-guided context sampling, and multi-fidelity pruning; (ii) budget-aware and any-time, improving with additional test-time compute unlike heuristics that produce fixed contexts; (iii) model-aware yet fully black-box, requiring no access to model internals, making it compatible with both proprietary and open-source TFMs; (iv) interpretable, identifying discrete ``Very Important Predictors'' (samples and features) that maximize signal-to-noise, which makes it (v) robust, isolating high-value data from noise. In contrast, soft-prompt optimization requires model gradients, produces abstract latent tokens, and lacks explicit signal discrimination. Extensive experiments show that VIP-COP consistently outperforms heuristic and optimized baselines across large-scale high-dimensional testbeds, including data augmentation and data-noise settings, establishing a new state of the art in test-time context refinement for TFMs.

Yilong Chen, Xueying Ding, Leman Akoglu

Carnegie Mellon University

In-Context Example Selection Tabular Foundation Models In-Context Learning Feature Selection Tabular ICL

Ask PDF

Abstract

Multifidelity Modeling

Explore similar work

On the Uncertainty Quantification Ability of Tabular Foundation Models

GEAR: Generative Expansion and Real Anchoring for Two-Stage Distillation of Tabular Foundation Models

VIP-COP: Context Optimization for Tabular Foundation Models