Andreas BedorfAt the prestigious 2025 Artificial Intelligence and Statistics (AISTATS) conference, held from May 3–5 at the Splash Beach Resort in Mai Khao, Thailand, the van der Schaar lab presents five papers pushing the boundaries in key fields such as causal inference, language model auditing, symbolic regression, feature acquisition, and clinical trial design.
These works reflect the innovative, impactful research our lab is known for, driven by the exceptional efforts of our students and collaborators. From advancing personalised medicine to redefining fairness and efficiency in AI, these papers demonstrate our commitment to shaping the future of AI and statistics at one of the field’s premier events.
Our Showcase:
Revolutionising Personalised Treatment: AFA4CATE Delivers Precision with Efficiency
We are unveiling AFA4CATE, a transformative framework reshaping how we approach personalised treatment decisions. Tackling the challenges of Conditional Average Treatment Effect (CATE) estimation head-on, AFA4CATE introduces a smarter way to acquire only the most critical features for each individual – saving time, reducing costs, and boosting accuracy. By addressing complex issues like confounding bias and overlap violations, this framework sets a new standard for making causal inference practical and impactful in real-world settings.
Whether in healthcare or other high-stakes fields, this work bridges cutting-edge theory with actionable solutions, paving the way for a more efficient, data-driven future in personalised decision-making.
Active Feature Acquisition for Personalised Treatment Assignment
J Piskorz, N Astorga, J Berrevoets, M van der Schaar
AISTATS 2025
Abstract
Making treatment effect estimation actionable for personalized decision-making requires overcoming the costs and delays of acquiring necessary features. While many machine learning models estimate Conditional Average Treatment Effects (CATE), they mostly assume that all relevant features are readily available at prediction time – a scenario that is rarely realistic.
In practice, acquiring features, such as medical tests, can be both expensive and time-consuming, highlighting the need for strategies that select the most informative features for each individual, enhancing decision accuracy while controlling costs. Existing active feature acquisition (AFA) methods, developed for supervised learning, fail to address the unique challenges of CATE, such as confounding, overlap, and the structural similarities of potential outcomes under different treatments.
To tackle these challenges, we propose specialised feature acquisition metrics and estimation strategies tailored to the CATE setting. We demonstrate the effectiveness of our methods through experiments on synthetic datasets designed to reflect common biases and data issues. In doing so, this work aims to bridge the gap between cutting-edge CATE estimation techniques and their practical, cost-efficient application in personalised treatment assignment.
Breaking New Ground in Causal Discovery: Identifiability Meets Differentiable Learning
This paper introduces NOTIME, the first differentiable Directed Acyclic Graph (DAG) learning algorithm with provable identifiability guarantees for Linear Non-Gaussian Additive Models (LiNGAM). By leveraging the d-dimensional Hilbert Schmidt Independence Criterion (dHSIC), NOTIME tackles the longstanding limitations of existing algorithms like NOTEARS, which fail under variable scaling and heteroscedastic noise. The approach finally connects theoretical causal discovery and practical applications, ensuring reliable identification of true causal structures regardless of data normalisation. Empirical results demonstrate NOTIME’s robustness and accuracy, outperforming state-of-the-art alternatives across a variety of settings.
This work paves the way for scalable, practical, and reliable causal inference in complex real-world scenarios.
Differentiable Causal Structure Learning with Identifiability by NOTIME
J Berrevoets, J Raymaekers, M van der Schaar, T Verdonck, R Yao
AISTATS 2025
Abstract
The introduction of the NOTEARS algorithm resulted in a wave of research on differentiable Directed Acyclic Graph (DAG) learning. Differentiable DAG learning transforms the combinatorial problem of identifying the DAG underlying a Structural Causal Model (SCM) into a constrained continuous optimization problem. Being differentiable, these problems can be solved using gradient-based tools which allow integration into other differentiable objectives.
However, in contrast to classical constrained-based algorithms, the identifiability properties of differentiable algorithms are poorly understood. We illustrate that even in the well-known Linear Non-Gaussian Additive Model (LiNGAM), the current state-of-the-art methods do not identify the true underlying DAG. To address the issue, we propose NOTIME, the first differentiable DAG learning algorithm with provable identifiability guarantees under the LiNGAM by building on a measure of (joint) independence.
With its identifiability guarantees, NOTIME remains invariant to normalization of the data on a population level, a property lacking in existing methods. NOTIME compares favourably against NOTEARS and other (scale-invariant) differentiable DAG learners, across different noise distributions and normalization procedures. Introducing the first identifiability guarantees to general LiNGAM is an important step towards practical adoption of differentiable DAG learners.
Shining a Light on Language Models: Auditing LLMs with Distribution-Based Sensitivity Analysis
In this work, we present DBSA (Distribution-Based Sensitivity Analysis), a revolutionary approach to auditing black-box language models (LLMs) for real-world, high-stakes applications. Unlike traditional methods that narrowly focus on biases or require access to model internals, DBSA provides a model-agnostic, plug-and-play tool for understanding how specific input tokens affect output distributions. By reframing auditing as a hypothesis-testing problem, DBSA captures the stochastic behaviour of LLMs and offers interpretable, token-level insights across diverse domains like legal, medical, and customer support.
This first-of-its-kind framework ensures greater accountability, highlights critical model sensitivities, and empowers practitioners with actionable insights into LLM behaviour.
Visualizing token importance for black-box language models
P Rauba*, Q Wei*, M van der Schaar
AISTATS 2025
Abstract
We consider the problem of auditing black-box large language models (LLMs) to ensure they behave reliably when deployed in production settings, particularly in high-stakes domains such as legal, medical, and regulatory compliance. Existing approaches for LLM auditing often focus on isolated aspects of model behavior, such as detecting specific biases or evaluating fairness.
However, these methods lack the flexibility, generality, and ease-of-use required by practitioners who want to audit their language models. We are interested in a more general question—can we understand how the outputs of black-box LLMs depend on any perturbation? There is a critical need to have such tools in real-world applications that rely on inaccessible API endpoints to language models. However, this is a non-trivial problem, as LLMs are stochastic functions (i.e. two outputs will be different by chance), while computing prompt-level gradients to approximate sensitivity is infeasible.
To address this, we propose Distribution-Based Sensitivity Analysis (DBSA), a lightweight model-agnostic procedure to evaluate the sensitivity of the output of a language model for any perturbation, without making any distributional assumptions about the LLM. DBSA is developed as a practical tool for practitioners, enabling quick, plug-and-play visual exploration of LLMs reliance on specific input tokens. Through illustrative examples, we demonstrate how DBSA can enable users to inspect LLM inputs and find sensitivities that may be overlooked by existing LLM interpretability methods.
Beyond Size: Redefining Complexity in Symbolic Regression
Here we are challenging traditional notions of complexity in symbolic regression (SR) by introducing a task-specific metric centred on Single-Feature Global Perturbation Analysis (SGPA). Traditional size-based measures like expression tree depth or term count often fail to capture the nuanced difficulty of performing analytical tasks. This work proposes a unified mathematical framework that quantifies the complexity of SGPA, offering a clearer understanding of how changes to input features affect global outputs. By aligning complexity measures with practical tasks like debugging and risk scoring, this approach reshapes the criteria for interpretable and efficient symbolic regression models.
With potential applications across physics, medicine, and other safety-critical domains, this research sets the stage for more meaningful and actionable metrics in the future of SR.
Beyond Size-Based Metrics: Measuring Task-Specific Complexity in Symbolic Regression
K Kacprzyk, M van der Schaar
AISTATS 2025
Abstract
Symbolic regression (SR) is a machine learning approach aimed at discovering mathematical closed-form expressions that best fit a given dataset. Traditional complexity measures in SR, such as the number of terms or expression tree depth, often fail to capture the difficulty of specific analytical tasks a user might need to perform.
In this paper, we introduce a new complexity measure designed to quantify the difficulty of conducting single-feature global perturbation analysis (SGPA)—a type of analysis commonly applied in fields like physics and risk scoring to understand the global impact of perturbing individual input features.
We present a unified mathematical framework that formalizes and generalizes these established practices, providing a precise method to assess how challenging it is to apply SGPA to different closed-form equations. This approach enables the definition of novel complexity metrics and constraints directly tied to this practical analytical task. Additionally, we establish a reconstruction theorem, offering potential insights for developing future optimization techniques in SR.
Reimagining Clinical Trials: RFAN Bridges Regulatory Standards and Real-World Impact
Randomize First Augment Next (RFAN) is an innovative framework that reshapes Phase III clinical trials by integrating regulatory compliance with real-world treatment effectiveness and fairness. RFAN pioneers two novel objectives: Post-Trial Mean Benefit (PTMB) and Post-Trial Fairness (PTF), ensuring that trial designs account for treatment policy value and equitable outcomes for underrepresented populations. By combining a randomised stage with an adaptive learning phase, RFAN utilises causal and Bayesian active learning techniques to optimise patient recruitment and treatment allocation.
With robust empirical results on synthetic and real-world datasets, this framework marks a transformative step toward clinical trials that not only meet approval standards but also maximise benefits across diverse patient groups, setting a new benchmark for fairness and efficacy in medical research.
Towards Regulatory-Confirmed Adaptive Clinical Trials: Machine Learning Opportunities and Solutions
O Noy Klein, A Hüyük, R Shamir, U Shalit, M van der Schaar
AISTATS 2025
Abstract
Randomized Controlled Trials (RCTs) are the gold standard for evaluating the effect of new medical treatments. Treatments must pass stringent regulatory conditions in order to be approved for widespread use, yet even after the regulatory barriers are crossed, real-world challenges might arise: Who should get the treatment? What is its true clinical utility? Are there discrepancies in the treatment effectiveness across diverse and under-served populations?
We introduce two new objectives for future clinical trials that integrate regulatory constraints and treatment policy value for both the entire population and under-served populations, thus answering some of the questions above in advance. Designed to meet these objectives, we formulate Randomize First Augment Next (RFAN), a new framework for designing Phase III clinical trials.
Our framework consists of a standard randomized component followed by an adaptive one, jointly meant to efficiently and safely acquire and assign patients into treatment arms during the trial. Then, we propose strategies for implementing RFAN based on causal, deep Bayesian active learning. Finally, we empirically evaluate the performance of our framework using synthetic and real-world semi-synthetic datasets.
Marika Niihori
