AIConfPaper
← Search

Tom Rainforth

54 accepted papers

2026

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

Deepro Choudhury, Sinead Williamson, Adam Golinski, Ning Miao, Freddie Bickford Smith, Michael Kirchhof +2
ICLR 2026poster

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 source using the framework of sequential Bayesian experimental design (BED). This enables LLMs to act as effective multi-tu…

Cited by 0Source ↗Scholar ↗
2026

Scaling Bayesian Experimental Design to High-Dimensions with Information-Guided Diffusion

Adhi Saravanan, Rik Knowles, Gavin Kerrigan, Tom Rainforth
ICLR 2026poster

We present DiffBED, a Bayesian experimental design (BED) approach that scales to problems with high-dimensional design spaces. Our key insight is that current BED approaches typically cannot be scaled to real high--dimensional design problems because of the need to specify a likelihood model that re…

Cited by 0Source ↗Scholar ↗
2025

Rethinking Aleatoric and Epistemic Uncertainty

Freddie Bickford Smith, Jannik Kossen, Eleanor Trollope, Mark van der Wilk, Adam Foster, Tom Rainforth
ICML 2025poster

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 discussions of these ideas and suggest this stems from the aleatoric-epistemic view being insufficiently expressive to captur…

Cited by 2Source ↗PDF ↗Scholar ↗
2025

Scaling Up Active Testing to Large Language Models

Gabrielle Berrada, Jannik Kossen, Freddie Bickford Smith, Muhammed Razzak, Yarin Gal, Tom Rainforth
NeurIPS 2025poster

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 measures that enable active testing to be scaled up to large language models (LLMs). In particular we show that the surro…

Cited by 0Source ↗Scholar ↗
2025

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

Cornelius Emde, Alasdair Paren, Preetham Arvind, Maxime Guillaume Kayser, Tom Rainforth, Thomas Lukasiewicz +2
ICLR 2025poster

Large language models (LLMs) are often deployed to do constrained tasks, with narrow domains. For example, customer support bots can be built on top of LLMs, relying on their broad language understanding and capabilities to enhance performance. However, these LLMs are adversarially susceptible, pote…

Cited by 1Source ↗PDF ↗Scholar ↗
2025

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

Marcel Hedman, Desi R. Ivanova, Cong Guan, Tom Rainforth
ICML 2025poster

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-based BED approaches, Step-DAD trains a design policy upfront before the experiment. However, rather than keeping this pol…

Cited by 6Source ↗PDF ↗Scholar ↗
2024

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

Tim Reichelt, Luke Ong, Tom Rainforth
AISTATS 2024poster

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 show that making predictions with this full posterior implicitly performs a Bayesian model averaging (BMA) over paths. This…

Cited by 1Source ↗PDF ↗code ↗Scholar ↗
2024

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

Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, Tommi Jaakkola
ICML 2024poster

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 provides the missing link in enabling flow-based generative models to be applied to multimodal continuous and discrete data pr…

Cited by 112Source ↗PDF ↗code ↗Scholar ↗
2024

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

Jannik Kossen, Yarin Gal, Tom Rainforth
ICLR 2024poster

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, there is currently no consensus about how this in-context learning (ICL) ability of LLMs works. For example, while Xie et a…

Cited by 37Source ↗PDF ↗code ↗Scholar ↗
2024

Making Better Use of Unlabelled Data in Bayesian Active Learning

Freddie Bickford Smith, Adam Foster, Tom Rainforth
AISTATS 2024poster

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 but also decisions about what data to acquire. Our proposed solution is a simple framework for semi-supervised Bayesian…

Cited by 5Source ↗PDF ↗code ↗Scholar ↗
2024

On the Expected Size of Conformal Prediction Sets

Guneet S. Dhillon, George Deligiannidis, Tom Rainforth
AISTATS 2024poster

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 practical utility. Unfortunately, there is currently a lack of finite-sample analysis and guarantees for their prediction set…

Cited by 11Source ↗PDF ↗code ↗Scholar ↗
2024

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

Ning Miao, Yee Whye Teh, Tom Rainforth
ICLR 2024poster

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 reasoning. However, when faced with more complicated problems that require non-linear thinking, even the strongest LLMs make…

Cited by 114Source ↗PDF ↗code ↗Scholar ↗
2023

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

Desi R. Ivanova, Joel Jennings, Tom Rainforth, Cheng Zhang, Adam Foster
ICML 2023poster

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 contextual experiments using information-theoretic principles. After formulating a suitable information-based objective, we empl…

Cited by 4Source ↗PDF ↗code ↗Scholar ↗
2023

Deep Stochastic Processes via Functional Markov Transition Operators

Jin Xu, Emilien Dupont, Kaspar Märtens, Tom Rainforth, Yee Whye Teh
NeurIPS 2023poster

We introduce Markov Neural Processes (MNPs), a new class of Stochastic Processes (SPs) which are constructed by stacking sequences of neural parameterised Markov transition operators in function space. We prove that these Markov transition operators can preserve the exchangeability and consistency o…

Cited by 7Source ↗PDF ↗Scholar ↗
2023

Do Bayesian Neural Networks Need To Be Fully Stochastic?

Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, Tom Rainforth
AISTATS 2023poster

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 construction may be unnecessary. To this end, we prove that expressive predictive distributions require only small amounts o…

Cited by 69Source ↗PDF ↗code ↗Scholar ↗
2023

Learning Instance-Specific Augmentations by Capturing Local Invariances

Ning Miao, Tom Rainforth, Emile Mathieu, Yann Dubois, Yee Whye Teh, Adam Foster +1
ICML 2023poster

We introduce InstaAug, a method for automatically learning input-specific augmentations from data. Previous methods for learning augmentations have typically assumed independence between the original input and the transformation applied to that input. This can be highly restrictive, as the invarianc…

Cited by 12Source ↗PDF ↗code ↗Scholar ↗
2023

Prediction-Oriented Bayesian Active Learning

Freddie Bickford Smith, Andreas Kirsch, Sebastian Farquhar, Yarin Gal, Adam Foster, Tom Rainforth
AISTATS 2023poster

Information-theoretic approaches to active learning have traditionally focused on maximising the information gathered about the model parameters, most commonly by optimising the BALD score. We highlight that this can be suboptimal from the perspective of predictive performance. For example, BALD lac…

Cited by 46Source ↗PDF ↗code ↗Scholar ↗
2023

Trans-Dimensional Generative Modeling via Jump Diffusion Models

Andrew Campbell, William Harvey, Christian Dietrich Weilbach, Valentin De Bortoli, Tom Rainforth, Arnaud Doucet
NeurIPS 2023spotlight

We propose a new class of generative model that naturally handles data of varying dimensionality by jointly modeling the state and dimension of each datapoint. The generative process is formulated as a jump diffusion process that makes jumps between different dimensional spaces. We first define a di…

Cited by 11Source ↗PDF ↗code ↗Scholar ↗
2022

A Continuous Time Framework for Discrete Denoising Models

Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, Arnaud Doucet
NeurIPS 2022accept

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 corresponding reverse time generative process as Continuous Time Markov Chains (CTMCs). The model can be efficiently trained using…

Cited by 169Source ↗PDF ↗code ↗Scholar ↗
2022

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

Jannik Kossen, Sebastian Farquhar, Yarin Gal, Tom Rainforth
NeurIPS 2022accept

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 are expensive. ASEs address this active testing problem using a surrogate-based estimation approach that interpolates the…

Cited by 20Source ↗PDF ↗code ↗Scholar ↗
2022

Amortized Rejection Sampling in Universal Probabilistic Programming

Saeid Naderiparizi, Adam Scibior, Andreas Munk, Mehrdad Ghadiri, Atilim Gunes Baydin, Bradley J. Gram-Hansen +6
AISTATS 2022poster

Naive approaches to amortized inference in probabilistic programs with unbounded loops can produce estimators with infinite variance. This is particularly true of importance sampling inference in programs that explicitly include rejection sampling as part of the user-programmed generative procedure.…

Cited by 7Source ↗PDF ↗code ↗Scholar ↗
2022

Expectation programming: Adapting probabilistic programming systems to estimate expectations efficiently

Tim Reichelt, Adam Goliński, Luke Ong, Tom Rainforth
UAI 2022poster

We show that the standard computational pipeline of probabilistic programming systems (PPSs) can be inefficient for estimating expectations and introduce the concept of expectation programming to address this. In expectation programming, the aim of the backend inference engine is to directly estimat…

Cited by 3Source ↗PDF ↗Scholar ↗
2022

Learning Multimodal VAEs through Mutual Supervision

Tom Joy, Yuge Shi, Philip Torr, Tom Rainforth, Sebastian M Schmon, Siddharth N
ICLR 2022spotlight

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 work has typically combined information from the modalities by reconciling idiosyncratic representations directly in the…

Cited by 29Source ↗PDF ↗code ↗Scholar ↗
2022

On Incorporating Inductive Biases into VAEs

Ning Miao, Emile Mathieu, Siddharth N, Yee Whye Teh, Tom Rainforth
ICLR 2022poster

We explain why directly changing the prior can be a surprisingly ineffective mechanism for incorporating inductive biases into variational auto-encoders (VAEs), and introduce a simple and effective alternative approach: Intermediary Latent Space VAEs (InteL-VAEs). InteL-VAEs use an intermediary set…

Cited by 10Source ↗PDF ↗code ↗Scholar ↗
2022

Rethinking Variational Inference for Probabilistic Programs with Stochastic Support

Tim Reichelt, Luke Ong, Tom Rainforth
NeurIPS 2022accept

We introduce Support Decomposition Variational Inference (SDVI), a new variational inference (VI) approach for probabilistic programs with stochastic support. Existing approaches to this problem rely on designing a single global variational guide on a variable-by-variable basis, while maintaining th…

Cited by 3Source ↗PDF ↗code ↗Scholar ↗
2021

Active Testing: Sample-Efficient Model Evaluation

Jannik Kossen, Sebastian Farquhar, Yarin Gal, Tom Rainforth
ICML 2021spotlight

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 training, existing literature largely ignores the cost of labeling test data, typically unrealistically assuming large te…

Cited by 71Source ↗PDF ↗code ↗Scholar ↗
2021

Capturing Label Characteristics in VAEs

Tom Joy, Sebastian Schmon, Philip Torr, Siddharth N, Tom Rainforth
ICLR 2021poster

We present a principled approach to incorporating labels in variational autoencoders (VAEs) that captures the rich characteristic information associated with those labels. While prior work has typically conflated these by learning latent variables that directly correspond to label values, we argue t…

Cited by 55Source ↗PDF ↗code ↗Scholar ↗
2021

Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design

Adam Foster, Desi R Ivanova, Ilyas Malik, Tom Rainforth
ICML 2021oral

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 sequential Bayesian optimal experimental design approaches require substantial computation at each stage of the experiment. T…

Cited by 111Source ↗PDF ↗code ↗Scholar ↗
2021

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

Desi R. Ivanova, Adam Foster, Steven Kleinegesse, Michael U. Gutmann, Tom Rainforth
NeurIPS 2021poster

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 experimental design (BOED) by learning a design policy network upfront, which can then be deployed quickly at the time of…

Cited by 64Source ↗PDF ↗code ↗Scholar ↗
2021

Improving Transformation Invariance in Contrastive Representation Learning

Adam Foster, Rattana Pukdee, Tom Rainforth
ICLR 2021poster

We propose methods to strengthen the invariance properties of representations obtained by contrastive learning. While existing approaches implicitly induce a degree of invariance as representations are learned, we look to more directly enforce invariance in the encoding process. To this end, we firs…

Cited by 27Source ↗PDF ↗Scholar ↗
2021

Improving VAEs' Robustness to Adversarial Attack

Matthew JF Willetts, Alexander Camuto, Tom Rainforth, S Roberts, Christopher C Holmes
ICLR 2021poster

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 defend against such attacks remains an open problem. We make significant advances in addressing this issue by introduci…

Cited by 36Source ↗PDF ↗Scholar ↗
2021

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

Tim G. J. Rudner, Oscar Key, Yarin Gal, Tom Rainforth
ICML 2021spotlight

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) issues. Specifically, we show both theoretically and via an extensive empirical evaluation that the SNR of the gradient es…

Cited by 5Source ↗PDF ↗code ↗Scholar ↗
2021

Online Variational Filtering and Parameter Learning

Andrew Campbell, Yuyang Shi, Tom Rainforth, Arnaud Doucet
NeurIPS 2021oral

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. As per standard batch variational techniques, we use stochastic gradients to simultaneously optimize a lower bound on the…

Cited by 23Source ↗PDF ↗code ↗Scholar ↗
2021

Probabilistic Programs with Stochastic Conditioning

David Tolpin, Yuan Zhou, Tom Rainforth, Hongseok Yang
ICML 2021spotlight

We tackle the problem of conditioning probabilistic programs on distributions of observable variables. Probabilistic programs are usually conditioned on samples from the joint data distribution, which we refer to as deterministic conditioning. However, in many real-life scenarios, the observations a…

Cited by 6Source ↗PDF ↗code ↗Scholar ↗
2021

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

Jannik Kossen, Neil Band, Clare Lyle, Aidan Gomez, Tom Rainforth, Yarin Gal
NeurIPS 2021poster

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 this end, we introduce a general-purpose deep learning architecture that takes as input the entire dataset instead of pro…

Cited by 164Source ↗PDF ↗code ↗Scholar ↗
2021

Towards a Theoretical Understanding of the Robustness of Variational Autoencoders

Alexander Camuto, Matthew Willetts, Stephen Roberts, Chris Holmes, Tom Rainforth
AISTATS 2021poster

We make inroads into understanding the robustness of Variational Autoencoders (VAEs) to adversarial attacks and other input perturbations. While previous work has developed algorithmic approaches to attacking and defending VAEs, there remains a lack of formalization for what it means for a VAE to be…

Cited by 44Source ↗PDF ↗Scholar ↗
2020

A Unified Stochastic Gradient Approach to Designing Bayesian-Optimal Experiments

Adam Foster, Martin Jankowiak, Matthew O’Meara, Yee Whye Teh, Tom Rainforth
AISTATS 2020poster

We introduce a fully stochastic gradient based approach to Bayesian optimal experimental design (BOED). Our approach utilizes variational lower bounds on the expected information gain (EIG) of an experiment that can be simultaneously optimized with respect to both the variational and design paramete…

Cited by 85Source ↗PDF ↗code ↗Scholar ↗
2020

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

Yuan Zhou, Hongseok Yang, Yee Whye Teh, Tom Rainforth
ICML 2020poster

Universal probabilistic programming systems (PPSs) provide a powerful framework for specifying rich probabilistic models. They further attempt to automate the process of drawing inferences from these models, but doing this successfully is severely hampered by the wide range of non–standard models th…

Cited by 24Source ↗PDF ↗Scholar ↗
2019

A Statistical Approach to Assessing Neural Network Robustness

Stefan Webb, Tom Rainforth, Yee Whye Teh, M. Pawan Kumar
ICLR 2019poster

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 the probability of the event that the property is violated under an input model. Our approach critically varies from the f…

Cited by 105Source ↗PDF ↗code ↗Scholar ↗
2019

Disentangling Disentanglement in Variational Autoencoders

Emile Mathieu, Tom Rainforth, N Siddharth, Yee Whye Teh
ICML 2019oral

We develop a generalisation of disentanglement in variational autoencoders (VAEs)—decomposition of the latent representation—characterising it as the fulfilment of two factors: a) the latent encodings of the data having an appropriate level of overlap, and b) the aggregate encoding of the data confo…

Cited by 361Source ↗PDF ↗code ↗Scholar ↗
2019

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

Yuan Zhou, Bradley J. Gram-Hansen, Tobias Kohn, Tom Rainforth, Hongseok Yang, Frank Wood
AISTATS 2019poster

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 variables. The key success of this language and its compilation scheme is in its ability to automatically distinguish parameters…

Cited by 33Source ↗PDF ↗code ↗Scholar ↗
2018

Faithful Inversion of Generative Models for Effective Amortized Inference

Stefan Webb, Adam Golinski, Rob Zinkov, Siddharth N, Tom Rainforth, Yee Whye Teh +1
NeurIPS 2018poster

Inference amortization methods share information across multiple posterior-inference problems, allowing each to be carried out more efficiently. Generally, they require the inversion of the dependency structure in the generative model, as the modeller must learn a mapping from observations to distri…

Cited by 57Source ↗PDF ↗Scholar ↗
2018

Tighter Variational Bounds are Not Necessarily Better

Tom Rainforth, Adam Kosiorek, Tuan Anh Le, Chris Maddison, Maximilian Igl, Frank Wood +1
ICML 2018oral

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 signal-to-noise ratio of the gradient estimator. Our results call into question common implicit assumptions that tighter E…

Cited by 246Source ↗PDF ↗Scholar ↗
2016

Bayesian Optimization for Probabilistic Programs

Tom Rainforth, Tuan Anh Le, Jan-Willem van de Meent, Michael A Osborne, Frank Wood
NeurIPS 2016poster

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 evidence of any probabilistic program, and therefore of any graphical model, can be optimized with respect to an arbitrary su…

Cited by 34Source ↗PDF ↗code ↗Scholar ↗
2016

Interacting Particle Markov Chain Monte Carlo

Tom Rainforth, Christian Naesseth, Fredrik Lindsten, Brooks Paige, Jan-Willem Vandemeent, Arnaud Doucet +1
ICML 2016poster

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 related methods, iPMCMC is a Markov chain Monte Carlo sampler on an extended space. We present empirical results that sho…

Cited by 44Source ↗PDF ↗code ↗Scholar ↗