NAISS
SUPR
NAISS Projects
SUPR
Probabilistic methods for large-scale machine learning models
Dnr:

NAISS 2025/5-586

Type:

NAISS Medium Compute

Principal Investigator:

Martin Trapp

Affiliation:

Kungliga Tekniska högskolan

Start Date:

2026-01-01

End Date:

2026-07-01

Primary Classification:

10210: Artificial Intelligence

Allocation

Abstract

This research project aims to advance the trustworthiness and reliability of large-scale machine learning (ML) models, such as Large Language Models (LLMs) and vision-language models, through basic research in machine learning and the development of scalable and effective probabilistic methods for uncertainty quantification and reduction. In safety-critical domains, such as healthcare diagnostics or agentic systems, ML models must be able to make decisions under uncertainty and accurately capture and reflect their inherent uncertainties. Unfortunately, contemporary machine learning models are often overconfident in their predictions and lack the ability to "know when they don't know" [1]. Moreover, uncertainty quantification in contemporary machine learning models faces additional challenges, as these models often contain millions or billions of parameters or are pre-trained on proprietary data sets. This project will build on the applicant's expertise in probabilistic machine learning, Bayesian deep learning, and tractable models, with a strong publication record in top-tier machine learning venues, to tackle the aforementioned challenges. The research will build on the recent works of the application [2-4] and explore tractable representations [5,6] for efficient and scalable representation, propagation, and reduction of uncertainties. Achieving these goals will require significant computational resources to work with large language models, LLM agents, and other large-scale modelling families (e.g., vision-language models). The research output aims to have a positive impact by studying and improving the robustness and trustworthiness of contemporary machine learning models, resulting in open-source software to ensure a meaningful, broader impact of the project. --- References: [1] E. Nalisnick, A. Matsukawa, Y. Teh, D. Gorur & B. Lakshminarayanan, "Do deep generative models know what they don't know?", in ICLR 2019. [2] R. Li, M. Klasson, A. Solin, and M. Trapp, "Streamlining prediction in Bayesian deep learning", in ICLR 2025. [3] A. Sladek, M. Trapp, and A. Solin, “Approximate Bayesian inference via bitstring representations”, in UAI 2025. [4] A. Baumann, R. Li, M. Klasson, S. Mentu, S. Karthik, Z. Akata, A. Solin, and M. Trapp, “Post-hoc Probabilistic Vision-Language Models”, arXiv:2412.06014, 2024. [5] L. Loconte, A. M. Sladek, S. Mengel, M. Trapp, A. Solin, N. Gillis, and A. Vergari, “Subtractive mixture models via squaring: Representation and learning,” in ICLR 2024 (spotlight). [6] Z. Yu, M. Trapp, and K. Kersting, “Characteristic circuits,” in NeurIPS 2023 (oral).