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

The gist The FIRE framework introduces a training-free multi-fidelity regression method that couples tabular foundation models with distribution-conditioned in-context learning to perform zero-shot Bayesian inference via a high-fidelity correction model conditioned on the low-fidelity model’s posterior predictive distributions.

Multi-Fidelity Regression Challenges

Multi-fidelity (MF) regression is the problem of learning to predict an expensive, high-accuracy target function using a small number of high-fidelity observations together with a larger set of cheaper but less accurate low-fidelity observations. These practical applications often reside in an extreme data imbalance regime, where HF measurements are so scarce (less than 10% of the total dataset), that common probability regression models struggle to maintain reliable generalization. Classical MF regression methods typically rely on Gaussian Process (GP) surrogates using autoregressive (AR) couplings and residual decompositions to transfer knowledge across fidelities. However, these methods face significant computational and statistical barriers, such as GP’s cubic scaling costs making it less effective as LF datasets grow. Furthermore, most GP-based MF methods rely on nested input sets (XHF ⊆ XLF) to enable stable learning.

FIRE Framework Decomposition

FIRE decomposes zero-shot multi-fidelity regression learning into three stages: (i) low-fidelity inference, (ii) residual modeling conditioned on augmented statistical information, and (iii) additive uncertainty propagation, using frozen TFMs for zero-shot regression. The algorithm involves fitting a low-fidelity base model fθ on the aggregated low-fidelity dataset DLF. This stage yields the predictive mean µθ(x), variance σ2θ(x), and conditional quantiles q(τ)θ(x) for high-fidelity inputs.

Distribution-Conditioned Residual Learning

FIRE introduces distribution-conditioned residual transfer for MF regression, a two-stage inference procedure that conditions the HF correction on LF predictive distributions (mean/variance/quantiles), enabling heteroscedastic and non-Gaussian discrepancy handling without retraining. For each high-fidelity training point xHF, the residual is defined as r = yHF − µθ(xHF). To model this residual, an augmented input vector zaug = [xHF, µθ(xHF), σ2θ(xHF), q(τ)θ(xHF)] is constructed. A second TFM δϕ is then used to perform in-context learning on Daug to produce a predictive distribution for the residual: δϕ(zaug) ∼ qϕ(r zaug, Daug).

Theoretical Justification

The theoretical analysis confirms that conditioning HF residuals on LF base model’s distributional summaries is theoretically principled. Theorem 7.2 (Bayes Risk Monotonicity) confirms that adding variance and quantiles does not increase Bayes risk under the squared error loss. Furthermore, Proposition 7.6 (Quantile Conditioning Reduces Risk Under Misspecification) shows that if the cross-fidelity residual r is non-Gaussian given Zmv, then R(Zquant) < R(Zmv), capturing shape information that variance cannot represent. This demonstrates that conditioning on variance and quantiles allows the residual model to adapt the magnitude of its correction to the local epistemic uncertainty of the base model.

Empirical Performance and Robustness

Across 31 benchmark problems, FIRE improves the accuracy–uncertainty–runtime tradeoff and remains robust in extreme HF scarcity (2–5%), where standard GP-based methods struggle. FIRE achieves the highest Elo rating and lowest average rank of both accuracy (NRMSE) and uncertainty (NLL), consistently outperforming both GP-based and deep learning baselines. The performance gap is largest on real-world physics simulations and HPO tasks, confirming FIRE’s ability to effectively transfer information without overfitting. FIRE also demonstrates superior generalization to disjoint data structures, maintaining superior performance in the non-nested setting compared to autoregressive GP baselines.

Efficiency and Model Choice

FIRE is faster and more accurate than current baselines, standing out when high-fidelity data is very scarce (2-5%). FIRE’s zero-shot paradigm still offers a decisive speed advantage for high-dimensional regression tasks due to its zero-shot nature. Ablation studies show that using a TFM for the residual model consistently outperforms those using a GP, suggesting the TFM’s flexible prior handles complex, non-Gaussian structures better than a GP. The framework is also shown to be robust across different TFMs, with TabPFNv2.5 and RealTabPFN achieving the highest Elo ratings.

Conclusion

FIRE presents a training-free multi-fidelity regression framework that integrates tabular foundation models with residual and distribution-based uncertainty conditioning to tackle problems with very scarce high-fidelity data. It consistently achieves SOTA performance while maintaining a favorable runtime profile, often outperforming extensive GP or deep learning model retraining from the existing methods. FIRE is faster and more accurate than current baselines, standing out when high-fidelity data is very scarce (2-5%). It suggests that large-scale pre-training on diverse synthetic priors provides a strong inductive bias that compensates for HF data scarcity. The method aims to accelerate scientific discovery and engineering design by improving the efficiency of multi-fidelity regression.

Impact Statement

This work aims to accelerate scientific discovery and engineering design by improving the efficiency of multi-fidelity regression. By reducing the reliance on expensive high-fidelity simulations (e.g., in computational fluid dynamics or hyperparameter optimization), our method with faster runtime has the potential to lower the computational energy associated with data-intensive scientific modeling. We do not foresee any immediate negative ethical or societal consequences beyond the environmental impact of running inference on foundation models and standard considerations in automating engineering workflows.

References

Abrams, D. A. Design of Concrete Mixtures, volume 1. Structural Materials Research Laboratory, Lewis Institute, 1919.

Angrist, J., Chernozhukov, V., and Fernandez-Val, I. Quantile regression under misspecification, with an application to the us wage structure. Econometrica, 74(2):539–563.

Aretz, N., Gunzburger, M., Morlighem, M., and Willcox, K. Multifidelity uncertainty quantification for ice sheet simulations. Computational Geosciences, 29(1):5.

Bonfiglio, L., Perdikaris, P., Brizzolara, S., and Karniadakis, G. Multi-fidelity optimization of super-cavitating hydrofoils. Computer Methods in Applied Mechanics and Engineering, 332:63–85.

Casella, G. and Berger, R. Statistical inference (2nd ed.). Chapman and Hall/CRC, 2024.

Colombo, P., Miller, C., Yang, X., O’Donnell, R., and Maranzano, P. Warped multifidelity gaussian processes for data fusion of skewed environmental data. Journal of the Royal Statistical Society Series C: Applied Statistics, 74(3):844–865.

Cutajar, K., Pullin, M., Damianou, A., Lawrence, N., and Gonzalez, J. Deep gaussian processes for multi-fidelity ´ modeling. arXiv preprint arXiv:1903.07320.

Eggensperger, K., Muller, P., Mallik, N., Feurer, M., ¨Sass, R., Klein, A., Awad, N., Lindauer, M., and Hutter, F. Hpobench: A collection of reproducible multifidelity benchmark problems for hpo. arXiv preprint arXiv:2109.06716.

Elo, A. E. The proposed uscf rating system, its development, theory, and applications. Chess life, 22(8):242–247.

Elrefaie, M., Dai, A., and Ahmed, F. Drivaernet: A parametric car dataset for data-driven aerodynamic design and prediction. Journal of Mechanical Design, 147(4):041712.

Erickson, N., Purucker, L., Tschalzev, A., Holzmuller, D., ¨Desai, P. M., Salinas, D., and Hutter, F. Tabarena: A living benchmark for machine learning on tabular data. In The Thirty-ninth Annual Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2025.

Feldstein, A., Lazzara, D., Princen, N., and Willcox, K. Multifidelity data fusion: Application to blended-wingbody multidisciplinary analysis under uncertainty.

1. textbfDistribution-conditioned residual learning for heteroscedastic errors: Introduce a two-stage inference procedure that conditions HF corrections on LF predictive distributions (mean/variance/quantiles), enabling heteroscedastic and non-Gaussian discrepancy handling without retraining. This allows the system to capture input-dependent variability, where predictive uncertainty varies across the domain which standard mean-only couplings obscure.

2. textbfZero-shot Bayesian inference via Tabular Foundation Models: Leverage pre-trained 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 enables prediction in a single forward pass without local parameter updates, achieving SOTA performance even with extreme HF scarcity (2–5%) where other methods struggle.

3. textbf Quantile embedding for non-Gaussian residual modeling: Condition the residual model on a set of predictive quantiles Q = [0.1, 0.2,..., 0.9] to exploit distributional mismatch when residuals are non-Gaussian, as shown by Proposition 7.6 (Quantile Conditioning Reduces Risk Under Misspecification). This reduces the Bayes risk compared to conditioning only on means and variances when the cross-fidelity residual exhibits skewness, kurtosis != 3, or or heavy-tailed features.

4. textbfRobustness to non-nested data structures: Maintain superior generalization on disjoint input sets by avoiding assumptions of nested designs, as FIRE maintains superior performance and stability in the disjoint regime, demonstrating that it is robust to the irregular, non-nested data structures common in real-world applications.