Research of the Probabilistic Machine Learning Group

We develop new methods for probabilistic modeling, Bayesian inference and machine learning. Our current focuses are:

  • Amortized Experimental Design, Optimization and Simulation-based Inference
  • Generative Modeling
  • ATOM (Artificial agents with Theory Of Mind) and Human-In-The-Loop
  • Robustness to Distribution Shifts in Deployment
  • Privacy-preserving and Secure AI
  • AI For Research (e.g. BioDesign, Neuroscience, Health, Drug Design, Economics)


    • These topics match the research programs of FCAI which we contribute to.

    Amortized Experimental Design, Optimization and Simulation-based Inference

    Image 1 Image 2

    Bayesian experimental design (BED), Bayesian optimization (BO), and simulation-based inference (SBI) are driving forces in science and engineering as principled frameworks for collecting high-quality data and parameter estimation, respectively. A fundamental challenge with these frameworks is computational scalability due to the need for repeated inference. BED and BO involve repeating a lengthy process for estimating and optimizing the information gain or the acquisition function, respectively, for each new query, while in SBI, running costly posterior sampling methods (e.g., MCMC) for every new observation is a bottleneck. We tackle these challenges by developing methods that are “amortized”, i.e., methods that pay an upfront cost of training a deep learning model to circumvent a recurring computationally expensive step. Amortized experimental design, optimization, and SBI methods therefore accelerate scientific discovery and data analysis in many fields of science and engineering, making Bayesian inference scalable for complex problems.

    Reference

    Generative Modeling

    Generative models offer powerful tools to represent complex data distributions, with recent advances, such as flow-matching and diffusion models, providing expressive, high-dimensional priors. Together, these approaches form two key pillars for scientific applications. In generative modeling, our research focuses on improving model reliability when data is scarce or when prior expert knowledge is valuable. We develop methods for reusing and adapting pretrained generative models, integrating human-in-the-loop feedback, and reducing epistemic uncertainty in generated outputs. Our applications include molecular and image generation as well as natural language-based scientific interaction.

    Reference

    ATOM (Artificial agents with Theory Of Mind) and Human-In-The-Loop

    Critical problems tackled by machine learning involve human users in the loop. Typically, their solutions do not account for the users as individuals and hence fail to adapt to the tacit and evolving needs of their users. Addressing this caveat requires reasoning about the implications of human behavior and, reversely, taking the user's understanding of their own process (i.e., the machine learning system) into account. This includes e.g. inferring tacit and changing goals and knowledge, and modeling how humans strategically adapt to the system based on how they interpret its behavior.

    Our research aims to move away from systems in which users are treated as passive sources of data, towards expecting users to be strategic and collaborative. Our methods combine Bayesian inference and decision making, multi-agent reasoning, and ideas in cognitive science such as theory of mind and computational rationality. Output of this research are scalable posterior approximation methods, for example inference over the user, novel frameworks for human-AI decision making, and user models for realistic and sophisticated settings, such that our AI can truly collaborate with humans.
    Keywords: human-AI collaboration, Bayesian inference, sequential (multi-agent) decision making, user modeling, theory-of-mind.

    Reference

    Robustness to Distribution Shifts in Deployment

    The immense promise of machine learning is often challenged by its practical deployment, which faces a critical hurdle: distribution shifts. A core assumption in machine learning is that future data will follow the same statistical distribution as the training data. However, this assumption is frequently violated in the real world, since the underlying data-generating processes can shift over time or differ across contexts. Such shifts typically cause a sharp decline in model performance, due to issues like model misspecification, unrepresentative data sampling, or spurious correlations. Our work focuses on developing principled methods for building machine learning systems that are robust to unexpected confounders, making them apt for deployment in critical applications such as personalized medical treatment recommendation and drug discovery.

    To this end, our research addresses the core technical hurdles that prevent robust learning. We have developed techniques to improve the transfer of knowledge between different but related tasks, such as mitigating the effects of prior misspecification when adapting probabilistic models to new data (Sloman, et al., 2025). Furthermore, to ensure models generalize effectively across diverse environments and avoid the pitfalls of negative transfer, we have designed novel meta-learning approaches that capture the similarity of underlying causal mechanisms across datasets (Wharrie, et al., 2024).

    Reference

    Privacy-preserving and Secure AI

    Machine learning models are capable of unintentionally capturing sensitive information underlying the training data, leading to the disclosure of sensitive personal information. We combine differential privacy (DP), the prevailing formalism for anonymity, with Bayesian inference to develop powerful techniques for probabilistic modeling with rigorous privacy guarantees. We study the theoretical foundations and provide practical implementations for modern statistical inference under differential privacy, such as the software package Twinify, for creating and releasing anonymised synthetic twins of sensitive data sets. We also work on making DP easier to apply to new problems, supporting researchers in identifying optimal models and training configurations under DP.

    Reference

    AI For Research (e.g. BioDesign, Neuroscience, Health, Drug Design, Economics)

    While applying AI to individual scientific domains has produced techniques such as Physics-Informed Machine Learning, we focus on AI For Research, a paradigm that enhances the research workflow itself, by optimizing Design-Build-Test-Learn (DBTL) cycles, experimental design and data acquisition, and human-in-the-loop discovery.

    We work closely with researchers from other fields and industry partners to pursue breakthroughs in challenging real-world applications. Methodologically, many of our projects center on sequential decision-making problems, spanning healthcare, synthetic biology, drug design, cognitive science, materials science, robotics, and autonomous systems. Aligned with the FCAI Virtual Labs initiative, our recent efforts mainly split into the following strands:

  • BioDesign and BioFoundry: In collaboration with FCAI and VTT, we build virtual laboratories for enzyme engineering and biomanufacturing. We automate and accelerate the DBTL cycles by combining Bayesian experimental design and other machine learning techniques with domain expertise to improve enzyme and strain engineering.
  • Healthcare and personalized medicine: We develop causal machine learning and graph neural network (GNN) models to analyze large-scale datasets from nationwide biobanks and electronic health records. Our research focuses on estimating treatment effects, improving risk predictions, and uncovering complex patterns within healthcare systems. We work closely with medical experts and researchers from diverse institutions to advance these goals.
  • Drug discovery: We aim to accelerate drug discovery by optimizing experimental design, identifying toxic compounds early in the pipeline, and ultimately reducing clinical trial failure rates, in collaboration with pharma partners. By integrating machine learning with a human-in-the-loop framework, we improve molecule generation, molecular property prediction, and optimization, enabling more efficient and safer therapeutic development.
  • Neuroscience: Our work on identifying features in EEG and MEG data with Bayesian methods helps diagnose different levels of cognitive decline, such as mild cognitive impairment (MCI) and dementia, which we investigate with our partners in the AI-Mind consortium.

  • Reference

    • Masood, M.A., Kaski, S. & Cui, T. Molecular property prediction using pretrained-BERT and Bayesian active learning: a data-efficient approach to drug design. J Cheminform 17, 58 (2025). [Link].
    • Nahal, Y., Menke, J., Martinelli, J. et al. Human-in-the-loop active learning for goal-oriented molecule generation. J Cheminform 16, 138 (2024). [Link].
    • Klami A, Damoulas T, Engkvist O, Rinke P, Kaski S. Virtual laboratories: transforming research with AI. Data-Centric Engineering (2024) [Link].
    • Sophie Wharrie, Zhiyu Yang, Andrea Ganna, Samuel Kaski (2023). Characterizing personalized effects of family information on disease risk using graph representation learning. In Machine Learning for Healthcare Conference (pp. 824-845). PMLR. [Link].
    • Eemeli Leppäaho, Hanna Renvall, Elina Salmela, Juha Kere, Riitta Salmelin, Samuel Kaski (2019). Discovering heritable modes of MEG spectral power. Human Brain Mapping, 40(5):1391-1402. [Link].