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

arXiv: https://arxiv.org/abs/2302.14015 · HTML: https://mlsanity.com/paper/2302.14015/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 77 references that are in the map.
- First posted on arXiv 27 Feb 2023.

## Builds on (most cited first, 30 of 34)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [On Variational Bounds of Mutual Information](https://mlsanity.com/paper/1905.06922/) (2019)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Max-value Entropy Search for Efficient Bayesian Optimization](https://mlsanity.com/paper/1703.01968/) (2017)
- [Predictive Entropy Search for Efficient Global Optimization of Black-box Functions](https://mlsanity.com/paper/1406.2541/) (2014)
- [Variational Bayesian Optimal Experimental Design](https://mlsanity.com/paper/1903.05480/) (2019)
- [Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design](https://mlsanity.com/paper/2103.02438/) (2021)
- [Batch Bayesian Optimization via Local Penalization](https://mlsanity.com/paper/1505.08052/) (2015)
- [Deep End-to-end Causal Inference](https://mlsanity.com/paper/2202.02195/) (2022)
- [On Nesting Monte Carlo Estimators](https://mlsanity.com/paper/1709.06181/) (2017)
- [The Parallel Knowledge Gradient Method for Batch Bayesian Optimization](https://mlsanity.com/paper/1606.04414/) (2016)
- [Sequential Batch Learning in Finite-Action Linear Contextual Bandits](https://mlsanity.com/paper/2004.06321/) (2020)
- [Implicit Deep Adaptive Design: Policy-Based Experimental Design without Likelihoods](https://mlsanity.com/paper/2111.02329/) (2021)
- [A Unified Stochastic Gradient Approach to Designing Bayesian-Optimal Experiments](https://mlsanity.com/paper/1911.00294/) (2019)
- [Bayesian Experimental Design for Implicit Models by Mutual Information Neural Estimation](https://mlsanity.com/paper/2002.08129/) (2020)
- [Nearly Minimax-Optimal Regret for Linearly Parameterized Bandits](https://mlsanity.com/paper/1904.00242/) (2019)
- [Linear Bandits with Limited Adaptivity and Learning Distributional Optimal Design](https://mlsanity.com/paper/2007.01980/) (2020)
- [Interventions, Where and How? Experimental Design for Causal Models at Scale](https://mlsanity.com/paper/2203.02016/) (2022)
- [Variational Causal Networks: Approximate Bayesian Inference over Causal Structures](https://mlsanity.com/paper/2106.07635/) (2021)
- [Bayesian Algorithm Execution: Estimating Computable Properties of Black-box Functions Using Mutual Information](https://mlsanity.com/paper/2104.09460/) (2021)
- [Gradient-based Bayesian Experimental Design for Implicit Models using Mutual Information Lower Bounds](https://mlsanity.com/paper/2105.04379/) (2021)
- [Model-based Causal Bayesian Optimization](https://mlsanity.com/paper/2211.10257/) (2022)

## Built on by (most cited first, 4 of 4)

- [The Essential Role of Causality in Foundation World Models for Embodied AI](https://mlsanity.com/paper/2402.06665/) (2024)
- [Amortized Bayesian Experimental Design for Decision-Making](https://mlsanity.com/paper/2411.02064/) (2024)
- [Deep Adaptive Bayesian Screening](https://mlsanity.com/paper/2607.16927/) (2026)
- [Target-Dependent Limits of Causal Repair: A Leading-Log Frontier in a Gaussian Model](https://mlsanity.com/paper/2610.00424/) (2026)

## Abstract (arXiv)

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 employ black-box variational methods to simultaneously estimate it and optimize the designs in a single stochastic gradient scheme. In addition, to accommodate discrete actions within our framework, we propose leveraging continuous relaxation schemes, which can naturally be integrated into our variational objective. As a result, CO-BED provides a general and automated solution to a wide range of contextual optimization problems. We illustrate its effectiveness in a number of experiments, where CO-BED demonstrates competitive performance even when compared to bespoke, model-specific alternatives.