papers

Publications (75)

stat.CO2019

Hamiltonian Monte Carlo for Probabilistic Programs with Discontinuities

Bradley Gram-Hansen, Yuan Zhou, Tobias Kohn +3

Hamiltonian Monte Carlo (HMC) is arguably the dominant statistical inference algorithm used in most popular "first-order differentiable" Probabilistic Programming Languages (PPLs).…

stat.ML2020

A Unified Stochastic Gradient Approach to Designing Bayesian-Optimal Experiments

Adam Foster, Martin Jankowiak, Matthew O'Meara +2

We introduce a fully stochastic gradient based approach to Bayesian optimal experimental design (BOED). Our approach utilizes variational lower bounds on the expected information g…

stat.ML2018

Faithful Inversion of Generative Models for Effective Amortized Inference

Stefan Webb, Adam Golinski, Robert Zinkov +4

Inference amortization methods share information across multiple posterior-inference problems, allowing each to be carried out more efficiently. Generally, they require the inversi…

stat.ML2020

Variational Bayesian Optimal Experimental Design

Adam Foster, Martin Jankowiak, Eli Bingham +4

Bayesian optimal experimental design (BOED) is a principled framework for making efficient use of limited experimental resources. Unfortunately, its applicability is hampered by th…

stat.ML2017

Canonical Correlation Forests

Tom Rainforth, Frank Wood

We introduce canonical correlation forests (CCFs), a new decision tree ensemble method for classification and regression. Individual canonical correlation trees are binary decision…

cs.LG2025

Rethinking Aleatoric and Epistemic Uncertainty

Freddie Bickford Smith, Jannik Kossen, Eleanor Trollope +3

The ideas of aleatoric and epistemic uncertainty are widely used to reason about the probabilistic predictions of machine-learning models. We identify incoherence in existing discu…

cs.LG2022

Expectation Programming: Adapting Probabilistic Programming Systems to Estimate Expectations Efficiently

Tim Reichelt, Adam Goliński, Luke Ong +1

We show that the standard computational pipeline of probabilistic programming systems (PPSs) can be inefficient for estimating expectations and introduce the concept of expectation…

cs.LG2026

Efficient Adaptive Data Acquisition via Pretrained Belief Representations

Daolang Huang, Zhuoyue Huang, Conor Hassan +3

Learning effective policies for adaptive data acquisition remains challenging: posterior-based methods rely on surrogate models and posterior approximations that can be misspecifie…

cs.AI2025

Kosmos: An AI Scientist for Autonomous Discovery

Ludovico Mitchener, Angela Yiu, Benjamin Chang +34

Data-driven scientific discovery requires iterative cycles of literature search, hypothesis generation, and data analysis. Substantial progress has been made towards AI agents that…

stat.ML2025

A Geometric Approach to Optimal Experimental Design

Gavin Kerrigan, Christian A. Naesseth, Tom Rainforth

We introduce a novel geometric framework for optimal experimental design (OED). Traditional OED approaches, such as those based on mutual information, rely explicitly on probabilit…

stat.CO2018

Inference Trees: Adaptive Inference with Exploration

Tom Rainforth, Yuan Zhou, Xiaoyu Lu +4

We introduce inference trees (ITs), a new class of inference methods that build on ideas from Monte Carlo tree search to perform adaptive sampling in a manner that balances explora…

cs.LG2026

Active Learning with Task-Driven Representations for Messy Pools

Kianoosh Ashouritaklimi, Tom Rainforth

Active learning has the potential to be especially useful for messy, uncurated pools where datapoints vary in relevance to the target task. However, state-of-the-art approaches to…

stat.ML2019

Tighter Variational Bounds are Not Necessarily Better

Tom Rainforth, Adam R. Kosiorek, Tuan Anh Le +4

We provide theoretical and empirical evidence that using tighter evidence lower bounds (ELBOs) can be detrimental to the process of learning an inference network by reducing the si…

cs.LG2022

Active Surrogate Estimators: An Active Learning Approach to Label-Efficient Model Evaluation

Jannik Kossen, Sebastian Farquhar, Yarin Gal +1

We propose Active Surrogate Estimators (ASEs), a new method for label-efficient model evaluation. Evaluating model performance is a challenging and important problem when labels ar…

cs.LG2022

Amortized Rejection Sampling in Universal Probabilistic Programming

Saeid Naderiparizi, Adam Ścibior, Andreas Munk +9

Naive approaches to amortized inference in probabilistic programs with unbounded loops can produce estimators with infinite variance. This is particularly true of importance sampli…

stat.CO2018

On Nesting Monte Carlo Estimators

Tom Rainforth, Robert Cornish, Hongseok Yang +2

Many problems in machine learning and statistics involve nested expectations and thus do not permit conventional Monte Carlo (MC) estimation. For such problems, one must nest estim…

cs.LG2024

Making Better Use of Unlabelled Data in Bayesian Active Learning

Freddie Bickford Smith, Adam Foster, Tom Rainforth

Fully supervised models are predominant in Bayesian active learning. We argue that their neglect of the information present in unlabelled data harms not just predictive performance…

stat.ML2023

Deep Stochastic Processes via Functional Markov Transition Operators

Jin Xu, Emilien Dupont, Kaspar Märtens +2

We introduce Markov Neural Processes (MNPs), a new class of Stochastic Processes (SPs) which are constructed by stacking sequences of neural parameterised Markov transition operato…

cs.LG2023

Learning Instance-Specific Augmentations by Capturing Local Invariances

Ning Miao, Tom Rainforth, Emile Mathieu +4

We introduce InstaAug, a method for automatically learning input-specific augmentations from data. Previous methods for learning augmentations have typically assumed independence b…

stat.ML2021

Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design

Adam Foster, Desi R. Ivanova, Ilyas Malik +1

We introduce Deep Adaptive Design (DAD), a method for amortizing the cost of adaptive Bayesian experimental design that allows experiments to be run in real-time. Traditional seque…

stat.ML2026

Action-BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives

Tom Rossa, Angus Phillips, Tom Rainforth

Bayesian experimental design (BED) has traditionally been based on maximising expected uncertainty reductions from prior to posterior. A major shortfall of this approach is that it…

stat.ML2022

Online Variational Filtering and Parameter Learning

Andrew Campbell, Yuyang Shi, Tom Rainforth +1

We present a variational method for online state estimation and parameter learning in state-space models (SSMs), a ubiquitous class of latent variable models for sequential data. A…

cs.LG2022

Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning

Jannik Kossen, Neil Band, Clare Lyle +3

We challenge a common assumption underlying most supervised deep learning: that a model makes a prediction depending only on its parameters and the features of a single input. To t…

stat.ML2021

Improving VAEs' Robustness to Adversarial Attack

Matthew Willetts, Alexander Camuto, Tom Rainforth +2

Variational autoencoders (VAEs) have recently been shown to be vulnerable to adversarial attacks, wherein they are fooled into reconstructing a chosen target image. However, how to…

stat.ML2025

Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design

Marcel Hedman, Desi R. Ivanova, Cong Guan +1

We develop a semi-amortized, policy-based, approach to Bayesian experimental design (BED) called Stepwise Deep Adaptive Design (Step-DAD). Like existing, fully amortized, policy-ba…

stat.ML2024

Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design

Andrew Campbell, Jason Yim, Regina Barzilay +2

Combining discrete and continuous data is an important capability for generative models. We present Discrete Flow Models (DFMs), a new flow-based model of discrete data that provid…

cs.AI2023

SelfCheck: Using LLMs to Zero-Shot Check Their Own Step-by-Step Reasoning

Ning Miao, Yee Whye Teh, Tom Rainforth

The recent progress in large language models (LLMs), especially the invention of chain-of-thought prompting, has made it possible to automatically answer questions by stepwise reas…

cs.LG2021

Probabilistic Programs with Stochastic Conditioning

David Tolpin, Yuan Zhou, Tom Rainforth +1

We tackle the problem of conditioning probabilistic programs on distributions of observable variables. Probabilistic programs are usually conditioned on samples from the joint data…

stat.ML2023

Modern Bayesian Experimental Design

Tom Rainforth, Adam Foster, Desi R Ivanova +1

Bayesian experimental design (BED) provides a powerful and general framework for optimizing the design of experiments. However, its deployment often poses substantial computational…

cs.LG2019

LF-PPL: A Low-Level First Order Probabilistic Programming Language for Non-Differentiable Models

Yuan Zhou, Bradley J. Gram-Hansen, Tobias Kohn +3

We develop a new Low-level, First-order Probabilistic Programming Language (LF-PPL) suited for models containing a mix of continuous, discrete, and/or piecewise-continuous variable…

cs.LG2022

Capturing Label Characteristics in VAEs

Tom Joy, Sebastian M. Schmon, Philip H. S. Torr +2

We present a principled approach to incorporating labels in VAEs that captures the rich characteristic information associated with those labels. While prior work has typically conf…

stat.ML2022

On Incorporating Inductive Biases into VAEs

Ning Miao, Emile Mathieu, N. Siddharth +2

We explain why directly changing the prior can be a surprisingly ineffective mechanism for incorporating inductive biases into VAEs, and introduce a simple and effective alternativ…

cs.LG2022

Learning Multimodal VAEs through Mutual Supervision

Tom Joy, Yuge Shi, Philip H. S. Torr +3

Multimodal VAEs seek to model the joint distribution over heterogeneous data (e.g.\ vision, language), whilst also capturing a shared representation across such modalities. Prior w…

cs.CR2020

A note on blind contact tracing at scale with applications to the COVID-19 pandemic

Jack K. Fitzsimons, Atul Mantri, Robert Pisarczyk +2

The current COVID-19 pandemic highlights the utility of contact tracing, when combined with case isolation and social distancing, as an important tool for mitigating the spread of…

stat.ML2026

Bayesian Experimental Design via Score Matching

Angus Phillips, Gavin Kerrigan, Tom Rainforth

Policy-based approaches to Bayesian experimental design (BED) allow the learning of deep policy networks that adaptively make intelligent design decisions based on previously colle…

stat.CO2017

Interacting Particle Markov Chain Monte Carlo

Tom Rainforth, Christian A. Naesseth, Fredrik Lindsten +4

We introduce interacting particle Markov chain Monte Carlo (iPMCMC), a PMCMC method based on an interacting pool of standard and conditional sequential Monte Carlo samplers. Like r…

cs.LG2024

Incorporating Unlabelled Data into Bayesian Neural Networks

Mrinank Sharma, Tom Rainforth, Yee Whye Teh +1

Conventional Bayesian Neural Networks (BNNs) are unable to leverage unlabelled data to improve their predictions. To overcome this limitation, we introduce Self-Supervised Bayesian…

stat.ML2016

Probabilistic structure discovery in time series data

David Janz, Brooks Paige, Tom Rainforth +2

Existing methods for structure discovery in time series data construct interpretable, compositional kernels for Gaussian process regression models. While the learned Gaussian proce…

stat.ML2020

Divide, Conquer, and Combine: a New Inference Strategy for Probabilistic Programs with Stochastic Support

Yuan Zhou, Hongseok Yang, Yee Whye Teh +1

Universal probabilistic programming systems (PPSs) provide a powerful framework for specifying rich probabilistic models. They further attempt to automate the process of drawing in…

cs.LG2021

Test Distribution-Aware Active Learning: A Principled Approach Against Distribution Shift and Outliers

Andreas Kirsch, Tom Rainforth, Yarin Gal

Expanding on MacKay (1992), we argue that conventional model-based methods for active learning - like BALD - have a fundamental shortfall: they fail to directly account for the tes…

stat.ML2019

Hijacking Malaria Simulators with Probabilistic Programming

Bradley Gram-Hansen, Christian Schröder de Witt, Tom Rainforth +3

Epidemiology simulations have become a fundamental tool in the fight against the epidemics of various infectious diseases like AIDS and malaria. However, the complicated and stocha…

cs.CL2026

BED-LLM: Intelligent Information Gathering with LLMs and Bayesian Experimental Design

Deepro Choudhury, Sinead Williamson, Adam Goliński +5

We propose a general-purpose approach for improving the ability of large language models (LLMs) to intelligently and adaptively gather information from a user or other external sou…

cs.LG2025

Prediction-Oriented Subsampling from Data Streams

Benedetta Lavinia Mussati, Freddie Bickford Smith, Tom Rainforth +1

Data is often generated in streams, with new observations arriving over time. A key challenge for learning models from data streams is capturing relevant information while keeping…

stat.ML2023

CO-BED: Information-Theoretic Contextual Optimization via Bayesian Experimental Design

Desi R. Ivanova, Joel Jennings, Tom Rainforth +2

We formalize the problem of contextual optimization through the lens of Bayesian experimental design and propose CO-BED -- a general, model-agnostic framework for designing context…

stat.CO2016

On the Pitfalls of Nested Monte Carlo

Tom Rainforth, Robert Cornish, Hongseok Yang +1

There is an increasing interest in estimating expectations outside of the classical inference framework, such as for models expressed as probabilistic programs. Many of these conte…

stat.ML2017

Bayesian Optimization for Probabilistic Programs

Tom Rainforth, Tuan Anh Le, Jan-Willem van de Meent +2

We present the first general purpose framework for marginal maximum a posteriori estimation of probabilistic program variables. By using a series of code transformations, the evide…

cs.CL2025

Shh, don't say that! Domain Certification in LLMs

Cornelius Emde, Alasdair Paren, Preetham Arvind +6

Large language models (LLMs) are often deployed to perform constrained tasks, with narrow domains. For example, customer support bots can be built on top of LLMs, relying on their…

stat.ML2021

On Statistical Bias In Active Learning: How and When To Fix It

Sebastian Farquhar, Yarin Gal, Tom Rainforth

Active learning is a powerful tool when labelling data is expensive, but it introduces a bias because the training data no longer follows the population distribution. We formalize…

cs.CL2024

In-Context Learning Learns Label Relationships but Is Not Conventional Learning

Jannik Kossen, Yarin Gal, Tom Rainforth

The predictions of Large Language Models (LLMs) on downstream tasks often improve significantly when including examples of the input--label relationship in the context. However, th…

stat.ML2023

Trans-Dimensional Generative Modeling via Jump Diffusion Models

Andrew Campbell, William Harvey, Christian Weilbach +3

We propose a new class of generative models that naturally handle data of varying dimensionality by jointly modeling the state and dimension of each datapoint. The generative proce…

cs.LG2023

Rethinking Variational Inference for Probabilistic Programs with Stochastic Support

Tim Reichelt, Luke Ong, Tom Rainforth

We introduce Support Decomposition Variational Inference (SDVI), a new variational inference (VI) approach for probabilistic programs with stochastic support. Existing approaches t…

stat.ML2019

Disentangling Disentanglement in Variational Autoencoders

Emile Mathieu, Tom Rainforth, N. Siddharth +1

We develop a generalisation of disentanglement in VAEs---decomposition of the latent representation---characterising it as the fulfilment of two factors: a) the latent encodings of…

stat.ML2019

A Statistical Approach to Assessing Neural Network Robustness

Stefan Webb, Tom Rainforth, Yee Whye Teh +1

We present a new approach to assessing the robustness of neural networks based on estimating the proportion of inputs for which a property is violated. Specifically, we estimate th…

stat.ML2021

Statistically Robust Neural Network Classification

Benjie Wang, Stefan Webb, Tom Rainforth

Despite their numerous successes, there are many scenarios where adversarial risk metrics do not provide an appropriate measure of robustness. For example, test-time perturbations…

stat.ML2021

On Signal-to-Noise Ratio Issues in Variational Inference for Deep Gaussian Processes

Tim G. J. Rudner, Oscar Key, Yarin Gal +1

We show that the gradient estimates used in training Deep Gaussian Processes (DGPs) with importance-weighted variational inference are susceptible to signal-to-noise ratio (SNR) is…

cs.LG2021

Group Equivariant Subsampling

Jin Xu, Hyunjik Kim, Tom Rainforth +1

Subsampling is used in convolutional neural networks (CNNs) in the form of pooling or strided convolutions, to reduce the spatial dimensions of feature maps and to allow the recept…

cs.LG2019

On the Fairness of Disentangled Representations

Francesco Locatello, Gabriele Abbati, Tom Rainforth +3

Recently there has been a significant interest in learning disentangled representations, as they promise increased interpretability, generalization to unseen scenarios and faster l…

stat.ML2018

Auto-Encoding Sequential Monte Carlo

Tuan Anh Le, Maximilian Igl, Tom Rainforth +2

We build on auto-encoding sequential Monte Carlo (AESMC): a method for model and proposal learning based on maximizing the lower bound to the log marginal likelihood in a broad fam…

stat.ML2022

Certifiably Robust Variational Autoencoders

Ben Barrett, Alexander Camuto, Matthew Willetts +1

We introduce an approach for training Variational Autoencoders (VAEs) that are certifiably robust to adversarial attack. Specifically, we first derive actionable bounds on the mini…

stat.ML2026

Prediction-Powered Active Testing

Kianoosh Ashouritaklimi, Valentin Kilian, Daolang Huang +2

Active testing provides a label--efficient approach to risk estimation by adaptively selecting which test points should be labelled. However, existing estimators fail to exploit th…

stat.ML2024

On the Expected Size of Conformal Prediction Sets

Guneet S. Dhillon, George Deligiannidis, Tom Rainforth

While conformal predictors reap the benefits of rigorous statistical guarantees on their error frequency, the size of their corresponding prediction sets is critical to their pract…

cs.LG2021

Improving Transformation Invariance in Contrastive Representation Learning

Adam Foster, Rattana Pukdee, Tom Rainforth

We propose methods to strengthen the invariance properties of representations obtained by contrastive learning. While existing approaches implicitly induce a degree of invariance a…

stat.ML2019

Amortized Monte Carlo Integration

Adam Goliński, Frank Wood, Tom Rainforth

Current approaches to amortizing Bayesian inference focus solely on approximating the posterior distribution. Typically, this approximation is, in turn, used to calculate expectati…

cs.LG2025

Scaling Up Active Testing to Large Language Models

Gabrielle Berrada, Jannik Kossen, Freddie Bickford Smith +3

Active testing enables label-efficient evaluation of predictive models through careful data acquisition, but it can pose a significant computational cost. We identify cost-saving m…

cs.LG2024

Beyond Bayesian Model Averaging over Paths in Probabilistic Programs with Stochastic Support

Tim Reichelt, Luke Ong, Tom Rainforth

The posterior in probabilistic programs with stochastic support decomposes as a weighted sum of the local posterior distributions associated with each possible program path. We sho…

stat.ML2026

FairBED: A Bayesian Experimental Design Approach to Gathering Fairer Data

Marcel Hedman, Emily Alger, Brieuc Lehmann +2

Frameworks for ensuring fairness in machine learning typically focus on learning fair models from existing data. But this endeavor is often undermined by biases already present in…

cs.LG2023

Prediction-Oriented Bayesian Active Learning

Freddie Bickford Smith, Andreas Kirsch, Sebastian Farquhar +3

Information-theoretic approaches to active learning have traditionally focused on maximising the information gathered about the model parameters, most commonly by optimising the BA…

stat.ML2022

A Continuous Time Framework for Discrete Denoising Models

Andrew Campbell, Joe Benton, Valentin De Bortoli +3

We provide the first complete continuous time framework for denoising diffusion models of discrete data. This is achieved by formulating the forward noising process and correspondi…

stat.ML2021

Implicit Deep Adaptive Design: Policy-Based Experimental Design without Likelihoods

Desi R. Ivanova, Adam Foster, Steven Kleinegesse +2

We introduce implicit Deep Adaptive Design (iDAD), a new method for performing adaptive experiments in real-time with implicit models. iDAD amortizes the cost of Bayesian optimal e…

stat.ML2018

Nesting Probabilistic Programs

Tom Rainforth

We formalize the notion of nesting probabilistic programming queries and investigate the resulting statistical implications. We demonstrate that while query nesting allows the defi…

stat.ML2021

Towards a Theoretical Understanding of the Robustness of Variational Autoencoders

Alexander Camuto, Matthew Willetts, Stephen Roberts +2

We make inroads into understanding the robustness of Variational Autoencoders (VAEs) to adversarial attacks and other input perturbations. While previous work has developed algorit…

cs.LG2026

Loss-Driven Bayesian Active Learning

Zhuoyue Huang, Freddie Bickford Smith, Tom Rainforth

The central goal of active learning is to gather data that maximises downstream predictive performance, but popular approaches have limited flexibility in customising this data acq…

stat.ML2021

Active Testing: Sample-Efficient Model Evaluation

Jannik Kossen, Sebastian Farquhar, Yarin Gal +1

We introduce a new framework for sample-efficient model evaluation that we call active testing. While approaches like active learning reduce the number of labels needed for model t…

stat.ML2018

On Exploration, Exploitation and Learning in Adaptive Importance Sampling

Xiaoyu Lu, Tom Rainforth, Yuan Zhou +2

We study adaptive importance sampling (AIS) as an online learning problem and argue for the importance of the trade-off between exploration and exploitation in this adaptation. Bor…

cs.LG2023

Do Bayesian Neural Networks Need To Be Fully Stochastic?

Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick +1

We investigate the benefit of treating all the parameters in a Bayesian neural network stochastically and find compelling theoretical and empirical evidence that this standard cons…