Bayesian Experimental Design Via Contrastive Diffusions
Authors: Jacopo Iollo
Christophe Heinkelé
Pierre Alliez
Florence Forbes
Venue: ICLR 2025 Spotlight (2025)
Keywords: Bayesian Optimal Experimental Design
Conditional Diffusion Models
score based sampling
Bayesian Inverse Problems
Experimental Design
Sampling as Optimization
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: The paper introduces an efficient BOED method leveraging diffusion-based samplers and bi-level optimization ideas to jointly sample the introduced pooled posterior and maximize Expected Information Gain, enabling larger-scale applications.

Abstract: Bayesian Optimal Experimental Design (BOED) is a powerful tool to reduce the cost of running a sequence of experiments.
When based on the Expected Information Gain (EIG), design optimization corresponds to the maximization of some intractable expected  *contrast* between prior and posterior distributions.
Scaling this maximization to high dimensional and complex settings has been an issue due to BOED inherent computational complexity.
In this work, we introduce an *pooled posterior* distribution with cost-effective sampling properties and provide a tractable access to the EIG contrast maximization via a new EIG gradient expression. Diffusion-based samplers are used to compute the dynamics of the pooled posterior and ideas from bi-level optimization are leveraged to derive an efficient joint sampling-optimization loop, without resorting to lower bound approximations of the EIG. The resulting efficiency gain allows to extend BOED to the well-tested generative capabilities of diffusion models. 
By incorporating generative models into the BOED framework, we expand its scope and its use in scenarios that were previously impractical. Numerical experiments and comparison with state-of-the-art methods show the potential of the approach.

Published as a conference paper at ICLR 2025
BAYESIAN EXPERIMENTAL DESIGN VIA CONTRASTIVE DIFFUSIONS
Jacopo Iollo!.2,
Christophe Heinkele
Pierre Alliez?
Florence Forbes
1: Université Grenoble Alpes, Inria, CNRS, G-INP, France, nane. surnamc@inria.
2: Cerema, Andsum-Sttasbouun France, christophe wew.eeeeeeeeeeeeeeeeee.e. fr
3: Université Côte d' Azur, Inria, France, name. , stresssurnnnn surnamelinria. fr
ABSTRACT
Bayesian Optimal Experimental Design (BOED) is a powerful tool to reduce the cost of running a sequence of experiments. When based on the Expected Information Gain (EIG), design optimization corresponds to the maximization of some intractable expected contrast between prior and posterior distributions. Scaling this maximization to high dimensional and complex settings has been an issue due to BOED inherent computational complexity. In this work, we introduce a pooled posterior distribution with cost-effective sampling properties and provide a tractable access to the EIG contrast maximization via a new EIG gradient expression. Diffusion- samplers are used to compute the dynamics of the pooled posterior and ideas from bi-level optimization are leveraged to derive an efficient joint sampling-optimizttinn to extend BOED to the well-tested generative capabilities of diffusion models. By incorporating generative models into the BOED framework, we expand its scope and its use in scenarios that were previously impractical. Numerical experiments and comparison with state-of-the-art methods show the potential of the approach.
loop. The resulting efficiency gain allows
INTRODUCTION
Designing optimal experiments can be critical in numerous applied contexts where experiments are constrained in terms of resources or more generally costly and limited. In this work, design is assumed to be characterized by some continuous parameters & E E C Rd, which refers to the experimental part, such as the choice of a measurement location, that can be controlled to optimize the experimental outcome. We consider a Bayesian setting in which the parameters of interest is E θ C RTH and design is optimized to maximize the information gain on O. Bayesian optimal experimental design (BOED) is not a new topic in statistics, see e.g. Chaloner and Verdinelli (1995); Sebastiani and Wynn (2000); Amzal et al. (2006) but has recently gained new interest with the use of machine learning techniques, see Rainforth et al. (2024); Huan et al. (2024) for recent reviews. The most common approach consists of maximizing the so-called expected information gain (EIG), which mutual information criterion that accounts for information via the Shannon's entropy. Let p(0) denote a prior probability distribution and p( Pgg/0.E) likelihood defining the observation y ) generating process. The prior is assumed be independent on & and p(y 0 {) available in closed-form. To our knowledge, all previous BOED approaches also assume that the prior is available in closed-form a setting that we refer to as density- BOED. In this work, by making BOED more computationally efficient, we open the first access to diffusion-based generative models and introduce based BOED when the prior is only available through samples. This broadens the scope of problems that can be tackled to wide range of inverse problems (Daras et al., 2024).
The EIG, denoted below by L, admits several equivalent expressions, see e.g. Foster et al. (2019). It can be written the expected loss in entropy when accounting for observation y at E (eq. (1))
or as a mutual information (MI) or expected Kullback-Leibhe (KL) divergence (eq. (2)). Denoting Pe (e(OO) y) p(0, ÿ the joint distribution of (0, Y) and using p00, y = p(0 y,, () <) (0), it comes,
ply 0, (O),
T(E) = Epis E) (p))))) H(p(O|Y,E)
(1)
Eppi E) [KL(((OY.EE.P((0) MI(pe)
(2)

---

Published as a conference paper at ICLR 2025
where random variables are indicated with uppercase letters, Epoa EpI-] denotes the expectation with respect to P and H(p(O)) = - ee [log p(0) (log)(0) is the entropy of p. The joint distribution Pe completely determines all other distributions, marginal (prior) and conditional (posterior) distributions, so that the mutual information, which is the KL between the joint and the product of its marginal distributions, can be written as a function of P5 E P(@ x )) only. In the following P(O x 3), resp. P()), resp. P(Y), denotes the set of probability measures on @ x V resp. O resp. y/
In BOED, we look for 5" satisfying
E" E arg max I(E) = arg max MI(pg). EER
ECR-
The above optimization is usually referred to static design optimization. The main challenge in ElG-based BOED is that both the EIG and its gradient with respect to & are doubly intractable. Their respective expressions involve an expectation of an intractable integrand over a posterior distribution which is itself not straightforward to sample from. The posterior distribution is generally only accessible through an iterative algorithm providing approximate samples. practice, the inference problem is further complicated as design optimization is considered in sequential context, in which a series of experiments is planned sequentially and each successive design has to be accounted for. In order to remove the integrand intractability issue, solutions have been proposed which optimize an EIG lower bound (Foster et al., 2019). This lower bound can be expressed as an expectation of a tractable integrand and becomes tight with increased simulation budgets. The remaining posterior sampling issue has then been solved in different ways. A set of approaches consists of approximating the problematic posterior distribution, either with variational techniques (Foster et al., 2019) with efficient sequential Monte Carlo (SMC) sampling (lollo et al., 2024; Drovandi et al., 2013). Other approaches avoid posterior estimation, using reinforcement learning (RL) and off-line policy learning to bypass the need for sampling (Foster et al., 2021; Ivanova et al., 2021; Blau et al., 2022). However, some studies have shown that estimating the posterior was beneficial, egg. Iollo et al. (2024) and Ivanova et al. (2024), which improves on Foster et al. (2021) by introducing posterior estimation steps in order to refine the learned policy. In addition, one should keep in mind that posterior inference is central in BOED as the ultimate goal is not design per se but to gain information on the parameter of interest. This is challenging especially in a sequential context. Previous attempts that provide both candidate design and estimates of the posterior distribution, such as Foster et al. (2019); Iollo et al. (2024), are thus essentially in alternating stages, approximate design optimization being dependent on approximate posterior sampling and vice-versa.
In this work, we propose a novel 1-stage approach which leverages a sampling- sanpling-as-optinng ostomiimiz setting (Korba and Salim, 2022; Marion al., 2025) where sampling is seen an optimization task over the space of probability distributions. We introduce a new EIG gradient expression (Section 3), which highlights the EIG gradient as a function of both the design and some sampling outcome. This fits into bi-level optimization framework adapted to BOED in Section 4. So doing. at each step. both an estimation of the optimal design and samples from the current posterior distribution can be provided in a single loop described in Section 5. It results an efficient procedure that can handle both traditional density-based samplers and data-based samplers such as provided by the highly successful diffusion-based generative models. The resulting efficiency gain enables BOED applications at significantly larger scales than previously feasible, including inpainting problems ranging from computer vision to protein engineering and MRI (Quan et al., 2024; Yang et al., 2019; Aali et al., 2023). Figure is an illustration on a 28 x 28 image O reconstruction problem from 7 x 7 sub-images centered at locations to be selected, details in Section 6. For simpler notation, first present our approach in the static design case. Adaptation to the sequential case is specified in Section 6 and all numerical examples are in the sequential setting.
3
3
8
3
3
3
3
3
3
3 3
8
3
3
3
3
9
9
3
9
9
7
Q
9
9
9
9
9
3
9
Figure 1: 28 Image (1st column) reconstruction from seven 7 sub-images y AçO + n centered at
seven central pixels & (designs) selected sequentially. Optimized vs. random designs: measured outcome (2nd vs. 3rd column) and parameter 0 estimates (rcconstruction) with highest weights (upper vs. lower sub-row).
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
We focus on gradient-based BOED for continuous problems. Applying a first-order method to solve (3) requires computing gradients of the EIG which are no more tractable than itself. Gradient-based BOED is generally based on stochastic gradient-type algorithms (see Section 4.3.2. in Huan et al. (2024)). This requires in principle unbiased gradient estimators, although stochastic approximation solutions using biased oracles have also been investigated, see e.g. Demidovich et al. (2023); Liu and Tajbakhsh (2024). To meet this requirement, most stochastic gradient-based approaches start from an EIG lower bound that yields tractable unbiased gradient estimators. More specifically, EIG lower bounds have usually the advantage to remove the nested expectation issuc, see e.g. Foster et al. (2019). In contrast, very few approaches focus on direct EIG gradient estimators. To our knowledge, this is only the case in Goda et al. (2022) and Ao and Li (2024), Goda et al. (2022) propose unbiased estimator of the EIG gradient using randomized version of multilevel nested Monte Carlo (MLMC) estimator from Rhee and Glynn (2015). A different estimator is proposed by Ao and Li (2024), who use MCMC samplers leading to biased estimators, for which the authors show empirically that the bias could be made negligible. In this work, we first show, in Section 3, that their two apparently different solutions actually only differ in the way the intractable posterior distribution is approximated. We then propose a third way to compute EIG gradients that is more computationally efficient and scales better larger data volumes and sequential design contexts. This new expression makes use of a distribution that we introduce and name the pooled posterior distribution. This latter distribution has interesting sampling features that allow us to leverage score-based sampling techniques and connect to the so-called implicit diffusion framework of Marion et al. (2025). Our single loop procedure in Section is inspired by Marion et al. (2025) and other recent developments in bi-level optimization (Yang et al., 2021; Dagréou et al., 2022; Hong et al., 2023). However, these latter settings do not cover doubly intractable objectives such as the EIG, which requires both appropriate gradient estimators and sampling operators, see our Sections 3 and 4. In BOED, efficient single loop procedures have been proposed by Foster et al. (2020) but they rely heavily on variational approximations, which may limit accuracy in scenarios with complex posterior distributions.
POOLED-POSTERIOR ESTIMATION OF THE EIG GRADIENT
Efficient EIG gradient estimators are central for accurate scalable BOED, Gradients derived from the reparameterization trick are often preferred, over the ones obtained with score-based techniques, as they have been reported to exhibit lower variance (Xu al., 2019).
EIG gradient via a reparametrization trick. Assuming 00 {) is such that Y can be rewritten as Y = TTe with Tg.0 invertible so that U = Te (Y) and U is a random variable independent on θ and with tractable distribution pt (U). The existence of Tq.e is straightforward if the direct model corresponds to an additive Gaussian noise as the transformation is then linear in U. Results exist guarantee the existence of such a transformation in more general situations (Papamakarios et al., 2021). I). Using this change of variable, two expressions of the EIG gradient, (5) and (6) below, can be derived. Detailed steps are given in Appendix A. With pę denoting the joint distribution p(0, ylE). g quantity related to the score g(E y. 0, 0') Ve
(u))0' ,))
( (==Tdd()) and
denoting (i(E, y, 0, 000") Vep(Te,o(u)) 00,, .E) (zdd)))) a first expression is
Epior [M Y,,C,,Y,,,') Ep(or |p 10', [
Ve [(E) (())=Fpp g Y,
Considering importance sampling formulations for the second term of (4), with an importance distribution g P()), potentially depending y. 0 and &, further leads to
p(0') g0 ..E) h Y 0. 0'
E 0.E)
TEE) = Epg g(E, 0,00)
(5)
E 0.0
p(0')
Pi
1
In Goda et al. (2022), this latter expression is used a randomized MLMC procedure with q set to a Laplace approximation of the posterior distribution, without justifiction for this specific choice of g.
3

---

Published as a conference paper at ICLR 2025
It results an estimator which is not unbiased but can be de-biased following Rhee and Glynn (2015). Alteratively, a second expression the EIG gradient is the starting point of Ao and Li (2024),
VET(E) pp [9(s. Y, 0, 0) E [g(E: 0,0,0000 e']]
follows a nested Monte Carlo estimator (30) given in Appendix A. using samples { I=)NN from the joint Pe and for each y/s, samples (0 (= from an MCMC procedure approximating the intractable posterior (())))) yy: E). Interestingly, expression (6) can also be recovered by setting the importance proposal q(0" l. 0, {) to p(6' ly, <) in (5), which provides a clear justification of why the choice of y made Goda al. (2022) is relevant. Approaches by Goda et al. (2022) and Ao and Li (2024) thus mainly differ in their choice of approximations for the posterior distribution. Using Laplace approximation as in Goda et al. (2022) is relevant only the posterior unimodal, which may not be the case in practice. The MCMC version of Ao and Li (2024) is then potentially more general but also more costly as it requires running N times a MCMC sampler, targeting each time a different posterior p( ye Q). In the next paragraph, introduce the pooled posterior distribution and derive another, more computationally efficient, gradient expression.
Importance sampling EIG gradient estimator with pooled posterior proposal. In their work, Ao and Li (2024) consider only static design, which hides the fact that for more realistic sequential design contexts, their solution is not tractable due to its computational complexity. Their solution faces the standard issue of nested estimation (Rainforth et al., 2018). To avoid this issue we propose to use an importance sampling expression for the second term in (6), which has the advantage to move the dependence on y (and 0) from the sampling part to the integrand part. We consider proposal distribution q E P()) that does not depend on Y nor θ. It comes,
((('( Y.YY
V I = Epe g Y. 0. -
9 Y. 0 0,0
q(0'lE)
and an approximate gradient can be obtained as
pp00' yi E) <(0')E)
g yi 02 07]
g ys, 0;,0 0.) 1.yer15))
(8)
N
The second term in (8) still requires N importance sampling approximations whose quality depends on the choice of the proposal distribution q. The ideal proposal q easy to simulate, with computable weights at least up to a constant, and so that and the multiple target distributions p(-ly E) are not too far apart. Given N samples {(O,, y: i-N:N from Pe, we propose thus to take q = gg.N where ag.N is the following logarithmic pooling or geometric mixture, with E
(0) x E) p()) Ily
We refer to Ig.N the pooled posterior distribution, defined in more general way in (13). It allows to assess the effect of a candidate design & on samples from the prior. It differs from a standard posterior as real data y obtained by running the experiment is available during the optimization. We only have access to samples {O0, (1))aa >=:::N from the joint Pç. The pooled posterior can be seen as distribution that takes into account all possible outcomes of candidate experiment & given the samples 1= from the prior. This choice of g..v is justified in Appendix B, using Lemma 2, proved therein. Lemma shows that, for EN Σi=1% 1, ge.N is the distribution g that minimizes the weighted sum of the KLs against each posterior p(0 ys, E), i.e. Ei" 1 PgKL[q p( (0 (0lyr, E)), leading to an efficient importance sampling proposal. It follows our new gradient estimator,
M
VeI(E)
1 N
g(E, y: 0,, 0,)
1
Σ
wij 9(E. yi, O 07)
(10)
M
where (0,,y:))==a) follow Pe, (0)h= follow gENv and wis p(e' ()
denotes the importance sampling weight. When this fraction can only be evaluated up constant, we consider
(0,)
4

---

Published as a conference paper at ICLR 2025
self normalized importance sampling (SNIS) using P. 9g.N the unnormalized versions of p and qg.N, Wi.j gg.N (0')
p(0 19. E)
ply ply.
W
and Wir
(11)
I P( E)
2j=1 Wid
Although with reduced computational cost, computing gradients with (10) still requires an iterative sampling algorithm ideally run for a large number of iterations to reach satisfying approximations of the joint Pe and the pooled posterior gg.v. In static design, sampling from the joint not generally difficult as the prior and the likelihood are assumed available but this becomes problematic in sequential design, as further detailed in Section 6. Sequential design the setting to be kept in mind in this paper and in practice, the exact distributions are rarely reached. To assess the impact on gradient approximations, it is convenient to introduce, as in Marion et al. (2025), gradient operators. In the next section, we show how to adapt the formalism of Marion et al. (2025) to our BOED task.
4 EIG OPTIMIZATION THROUGH SAMPLING
To maximize the EIG using its gradient estimator (10), samples are needed from both the joint distribution Pe and the pooled posterior proposal Ne If handled naively, it results computationally expensive nested sampling-optimization optimization loop where new samples from both distributions need to be generated at every update of the design parameter Ç. To derive more efficient procedures, we propose to adapt to BOED the framework of Marion et al. (2025) that integrates sampling and optimization into a single bi-level optimization loop. To do so, the EIG gradient V. (g) has first to be expressed as function I' of three key components: the joint distribution pę. the proposal distribution q. and the design parameter& & itself. Our choice of the pooled posterior as proposal distribution q is then justified for its interesting sampling properties and the concept of sampling operator of Marion et al. (2025) is generalized to efficiently generate the samples needed to estimate our EIG gradient via I.
Estimation of gradients through sampling. to IRd, defined as,
Denote by I' function from
x Rd
[((0())) ((0)) g SiS, 1...Y00001]
[(p. E) = Ep g Y. 0000)- - -By
(12)
Expression (7) shows that Vel(E) = I(pe. q. E). where g is a distribution q(0'(E) on 0' possibly depending E. The gradient estimator (10) corresponds then to VET(E) [(pe- (g, +). where Pg. = E ((0..))) and 0g.v = Σド 50;- In general, sampling from Pę. or its sequential counterpart, and Vg.Nv is challenging and only possible through an iterative procedure. However, an interesting feature of our pooled posterior is that it does not add additional sampling diffieulties.
Pooled posterior distribution.
More generally (details in Appendix B), we define,
( (0) x exp
(( E))
(13)
where p is a measure on Y When p(y) = i vg,v(() in (9). The special structure of the pooled posterior allows to sample from it using the same algorithmic structure to sample from a single posterior p(0 y. E). Indeed, the score of (0) is linked to the posterior score,
1 ( with Σ; 1; 1, we recover ge,p
Vo log (0) = [, logp(0) (
(14)
which for ge,v simplifies into ΣM Vo p(0 y., sampling as the output of stochastic process iterating a so-called sampling operator.
E). In practice, we consider the operation of
Iterative sampling operators. Iterative sampling operators, as introduced in Marion et al. (2025), are mappings to space of probabilities. In our BOED setting, we consider two such operators. The first one is defined, for each &. through a sequence over s of functions from p(O x )) to p(O x )) and denoted by Σ (p. E). Sampling is defined as the outcome in the limit s + co or for some
finite s S of the following process starting from (0) E p(@ x y) and iterating
p(o+1)= ΣY (p(+),))
(15)
5

---

Published as a conference paper at ICLR 2025
Fasted Postetio 12.0%
Figure 2: Source localisation example. Prior (left) and pooled posterior (right) samples at experiment Final S (orange cross) at the end of the optimization sequence Ea- Er (blue crosses). This optimization "contrasts" the two distributions by making the pooled posterior "as different as possible" from the prior.
where pl*) can be explicit, egg. pls) is a Gaussian distribution with parameters depending on 8. or represented by a random variable X. ~ pl.). For example, in density-basch BOED, for X, -Y,, Os), consider the Euler discretization of the Langevin diffusion converging P
y ((+ = yy(0) 7 g g(*) 7 VoV ()0)))))=)))))
(6+) 060) + 2 y..:
(16)
() + / Bo..
where By and Bg., are realizations of independent standard Gaussian variables, 7 is step-size and V Vyy 0, E) = p()) log p(y) <) the pe potential, pe 0) x (-V 0, E)). The dynamics induced lead to samples from Pe for 2c. In the following, we will thus use the notation 2y to mean that we have access to samples from p(st1), which is equivalent to apply E to an )i=INN. Similarly, we can produce samples 1j=i:M Ij=1.M from the pooled posterior 9E,N, using its score expression, via the updating,
empirical version of p(*) built from samples {(y) (*) e
(0) (0)(0+1))
0'+++)= ((()) % v. VoV 00()) .) + 2
(17)
For sampling operator of the pooled posterior general form (13), need extend the definition in Marion et al. (2025) by adding dependence on some distribution p E P(Y) for the conditioning part. The second sampling operator defined, for some given & and p. through a sequence over of parameterized functions from P(O) to P(O) and denoted by Σ (g. p). The sampling operator defined as the outcome of the following process starting from ( P(@) and iterating
(e+1) I
(18)
For instance, when p
Em Uidgo g(o+1) Ig' (g()), P) can then be shorthand for (17).
5 SINGLE LOOP CONTRASTIVE EIG OPTIMIZATION
The perspective of optimization through sampling leads naturally to a nested loop procedure. An inner loop is performed reach good approximations Pę and qg.v using two samplers as specified in Section 4 and summarized in the nested loop Algorithm 1. Considering sampling as an optimization over the space of distributions (Marion al., 2025), a more efficient single loop procedure can be derived. As illustrated in the single loop Algorithm 2, at each optimization step over &, the sampling operators are applied only once using the current ę. which is updated turn, etc. Sampling operators can be derived from traditional density-based sampling, like in (16) and (17), where an expression of the target distribution is required to compute the score, and also from data-based sampling where only training samples are available. In the latter case, conditional score- based generative models have emerged as a very active field of research. We explicit below how a recent such framework proposed by Dou and Song (2024) can be used in our setting.

---

Published as a conference paper at ICLR 2025
Data-based samplers. In BOED, we are interested in sampling from a conditional distribution p(Oly, 5) with the following objectives. First we need to sample from the pooled posterior 9g.v (0) which requires conditioning section 6.1 and Appendix D), we need to condition on the history of observations D and produce samples from pk0) Both these issues can be tackled in the framework of diffusion models for inverse problems (Daras et al., 2024). When the likelihood corresponds to a linear measurement Y with Y = A + m and (0, E), inspiring recent attempts, such as (Corenflos al., 2025; Cardoso et al., 2024), have addressed the problem of sampling efficiently from p(0) g. E) using only the pre-trained score (0, t) of a diffusson model, without the need for any kind of retraining (see Appendix C for details). For conditional sampling, this would mean running an SDE with conditional score Ve log (c y. E) which is intractable. See (43) and Appendix C.2. As solution, Dou and Song (2024) propose method named FPS that approximates pr (r((() y. E) by pe (t) (+), E) with ( the noised observation at time As the score V Volog pt (0 (y(() E) can be written as Volog Pa (e(*))
the observation y. Second, in sequential design problems (see
(((())+ Veloggg Pr (t) (((*)) E), we can leverage the learned score s& (e t) and
the closed form of Ve (gy))) 1 <) sample approximately from p(O yy, {) using backward SDE with the approximate score, see (45) in Appendix C.2. This allows to sample efficiently from Dx-1) and, using Vo log N Σim U Vo log : E), from the pooled posterior with (t) is the noised yı at time of the forward SDE,
the extension of (45) below, where y
d0 (4) = 3(t) (1) 3(t) Σ V.Vo Po (4) 1y ,)) dt + BB())dB. 2
(2)
(19)
where above now flowing backwards from infinity to 1-0. In practice, (19) is solved approximately using a numerical discretization and an initialization of the process with (T) NoO, I) for some large finite An additional resampling SMC-like step can also be added as explained in Appendix E. The approach allows to handle new sequential data-based BOED tasks as illustrated in Section 6.3.
Algorithm 1: Nested-loop optimization
Algorithm 2: Single loop optimization
Result: Optimal design E
Result: Optimal design &*
Initialisation: & ERd
Initialisation: E (( po. (0) + qo
for t=0:T- (outer optimization loop) do
for 1=0:=1 (issmpling-optimiiaa
loop) do
(0) Po and ( qo
((++1) ()(())
(p()),))
Pi
for s=0:S- / (pg inner sampling) do
(++1) (n((+1) marginal over y)
Di+1 Pv
[++1)
=Y (p: (x) E.)
Pt
((++1) = ". ((:Pe1))))))
end
Compute
pg. + Pt (8)
Ve[(E))=I(((tt) (+1))
2, E.) in (12)
P: (y) (pe, marginal over y)
Update &, with SGD another optimizer
for 1a' (1)-1) Lg. (d)) Pc)
s'=:S' (qg.p inner sampling) do
end return ETi
end
qc. q: (S') Compute ) [f qg, E.) in (12) Update& &, with SGD another optimizer
Measure
1
2
3
4
5
6
CoDiff
.227
.338
.528
.673
.789
.826
Random
.168
.275
.350
.391
.421
.463
end return Exi
Table 1: CoDiff and random reconstruction quality comparison with SSIM, in [-1,1], the higher the better.
Dbased samplers. Among density-based samplers, we can mention score- score-based MCMC samplers, including Langevin dynamies via the Unadjusted Langevin Algorithm (ULA) and Metropolis Adjusted Langevin Algorithm (MALA) (Roberts and Tweedie, 1996), Hamiltonian Monte Carlo (HMC) samplers (Hoffman and Gelman, 2014). Section 6.2, an illustration is given with Langevin and sequential Monte Carlo (SMC) to handle sequential density-based BOED task.
Contrastive Optimization. Optimizing & using the gradient expression (10) encourages to select & that gives either high probability p(y ( :..&$ {) samples (0:, y:) from Pe low probability
7

---

Published as a conference paper at ICLR 2025
p(y, E) to samples 8 from gg.N. This contrastive behaviour is also visible in (2) where the EIG is defined as the mean over the experiment outcomes of the KL between posterior and prior distributions. The pooled posterior qg.v then used a proxy to the intractable posterior, to perform this contrastive optimization. Figure 2 provides a visualization of this contrastive behavior in the source localization example of Section 6.2. It corresponds set the next design & to a value that climinates the most parameter 0 values (right plot) among the possible ones a priori (left plot). This is analogous Noise Constrastive Estimation (Gutmann and Hyvärinen, 2010) methods where model parameters are computed so that the data samples are as different as possible from the noise samples. Additional illustrations are given in Appendix Figure 6.
NUMERICAL EXPERIMENTS
Two sequential Censity-based (Section 6.2) and data-based (Section 6.3) BOED examples are considered to illustrate that our method extends to the sequential case in both settings.
6.1 SEQUENTIAL BAYESIAN EXPERIMENTAL DESIGN
In the sequential setting, a sequence of K experiments is planned while gradually accounting for the successively collected data. At step k wish to pick the best design &. given previous outcomes The expected information gain in this scenario is given by:
Dx-1 = ((11.1)) E.), [E Da-1) = Pp (L((p(O &, Dx-)), p(0) Dk-1))]
where p(0 and p(Oly, DR-1) act respectively as prior and posterior analogues to the static case (2). See Appendix D for more detailed explanations. The main difference is that we no longer have direct access to samples from the step k prior p(0)Dx-1). However, as D(0)Dx-1) x p(0) I P(u,,0.. E.,) and p(0 y. &. DR-1) α p(0)p(yr)0 <) I[==1 P0y.,00.4.0.0000000 E.) we can still compute the score of these distributions and run sampling operators similar to (15) and (18). To emphasize their dependence on D they are denoted by Σ (p(x))((()) and >.g"
|D, (g-11((()), (g(*), p). Examples of these operators are provided in (20) and (21) in Section 6.2.
Evaluation metrics and comparison. We refer to our method as CoDiff. In Section 6.2, comparison is provided with other recent approaches, namely a reinforcement learning- based approach RL-BOED from Blau et al. (2022), the variational prior contrastive estimation VPCE of Foster et al. (2020) and a recent approach named PASOA (lollo et al., 2024) based on tempered sequential Monte Carlo samplers. We also compare with non tempered version of this latter approach (SMC) and with random baseline, where the observations (y1,- yyK} } are simulated with designs generated randomly, More details about these methods are given Appendix F.2. To compare methods in terms of information gains, we use the sequential prior contrastive estimation (SPCE) and sequential nested Monte Carlo (SNMC) bounds introduced in Foster et al. (2021) and used in Blau et al. (2022). These quantities allow to compare methods on the produced design sequences only, via their [SPCE, SNMCJ intervals which contain the total EIG. Their expressions are given Appendix F. We also provide the 12 Wasserstein distance between the produced samples and the true parameter O. For methods that do not provide posterior estimations or poor quality ones (RL-BOED, VPCE, Random), we compute Wasserstein distances posterior samples obtained by using tempered SMC on their design and observation sequences. In contrast, SMC Wasserstein distances are computed on the SMC posterior samples. In Section 6.3, our evaluation is mainly qualitative. The previous methods do not apply and we are not aware of existing attempts that could handle such generative setting.
6.2 SOURCES LOCATION FINDING
We present a source localization example inspired by Foster et al. (2021); Blau al. (2022). The setup
involves C' sources in RR?, with unknown positions 0 = {01.... 3c). The challenge is to determine optimal measurement locations to accurately infer the sources positions. When a measurement is
a- taken at location & E R.2 the signal strength defined as 1a(0 )) = b Σ m+100--51 where aee b, and m predefined constants. We assume standard Gaussian prior for each source location, 0, ~ (0, I2), and model the likelihood as nor-nolmal: (log y 0, E) N Nloog (O,E)) a). with o representing the standard deviation. For this experiment, we set C 2, on 02 1, 10-1
8

---

Published as a conference paper at ICLR 2025
b -1001. o- 0.5, and plan K 30 sequential design optimizations. In the notation of the single loop Algorithm 2, we consider Σ S.y. (p(t), .x)) and Σ, 1 ( Ee: pc+1) operators that correspond respectively to the update of batch samples of size N 200 and M 200
{(y (1) 0 I and {0, (1) JJ=I.M with ui+1 = Ei= P:O y using Langevin diffusions.
Making use of the availability of the likelihood in this example, sampling from it, is straightforward and sampling operator iterations simplify into, for i == 1 : .N and j = 1: M
1 0
+ Vo log (( D + Bot and y (e+1) ~ p(ylo (1) .g.)
(20)
N
0 e 50 + th Σ us Vo log (( ly.
(c+1) E... Dx-1) + v Bo
(21)
In practice, Langevin diffusion can get trapped in local minima and causes the sampling to be too slow to keep pace with the optimization process. To address this, we augment the Langevin diffusion with the Diffusive Gibbs (DiGS) MCMC kernel proposed by Chen et al. (2024). DiGS is an auxiliary variable MCMC method where the auxiliary variable X a noisy version of the original variable X. DiGS enhances mixing and helps escape local modes by alternately sampling from the distributions p(ā]z), which introduces noise via Gaussian convolution, and p(a[ā), which denoises the sample back to the original space using a score- update (here Langevin diffusion). With 400 total samples, each measurement step takes 2.9 s. This number of samples is insightful as it is usually the amount of samples one can afford to compute in the diffusion models of Section 6.3. The whole experiment is repeated 100 times with random source locations each time. Figure 3 shows, with respect k. the median for SPCE, the L2 Wasserstein distances between weighted samples and the true source locations and SNMC. CoDiff' clearly outperforms all other methods, with significant improvement, both in terms of information gain and posterior estimation. It improves by 30% the non-myopic RL-BOED results on SPCE and provides much higher SNMC. The L2 Wasserstein distance is two order of magnitude lower, suggesting the higher quality of our measurements.
LACS varit repeiteect
W. over sepetrents
QUAE cape mwneet
VFCC
Figure 3: Source location. Median and standard error over 100 rollouts for SPCE, L2 Wasserstein distance (log-scale), SNMC with respect to number of experiments k. Number of samples N+M=400.
6.3 IMAGE RECONSTRUCTION WITH DIFFUSION MODELS
We build an artificial experimental design task to illustrate the ability of our method to handle design parameters related to inverse problems with high dimensional parameter O. We consider the task of recovering an hidden image from only partial observations of its pixels. The image to be recovered is denoted by 0. An experiment corresponds to the choice of pixel around which an observation mask is centered and the image becomes visible. The measured observation y is then a masked version of O. The likelihood derives from the model Y Ag0 + 77 where Aş is square mask centered at E and 77 some Gaussian variable. For the image prior, we consider diffusion model trained for generation of the MNIST dataset (LeeCun et al., 1998). The goal is thus to select sequentially the best central pixel locations for 7 x 7 masks so as to reconstruct an entire 28 x 28 MNIST image in the smallest number of experiments. The smaller the mask the more interesting it becomes to optimally select the mask centers. Algorithm 2 is used with diffusion-basel sampling operators specified in Appendix F.2.2. The gain in optimizing the mask placements is illustrated in Figure 1 and Appendix Figure 9. It is confirmed quantitatively in Table 1. which reports reconstruction quality as measured by the structural similarity index measure (SSIM) ( (Wang et al., 2004), details in Appendix F.2.2. Progressive reconstructions are shown in Figures 4, 7 and 8. The digit to be recovered is shown in the

---

Published as a conference paper at ICLR 2025
1st column. The successively selected masks are shown (red line squares) in the 2nd column with the resulting gradually discovered part of the image. The reconstruction per se can be estimated from the posterior samples shown in the last 16 columns. At each experiment, the upper sub-row shows the 16 most-likely reconstructed images, while the lower sub-row shows the 16 less probable ones. As the number of experiments increases the posterior samples gradually concentrate on the right digit.
51
2
7
1
1
1 Dnteaniaent
9
7
7
5
3
161
3
3
3
3
3 3
3
3
3
8
3
L
O
8
/
&
S
8
3
5 3 8 5
3
5
5
8
5
3
5
5
O
6
O
6
5
7
2 O 7
4
O
7
7
L
Q
3
3
9
3
4 3
7
/
3
3
/
3
9
S
3
HGT METTTI WAOT CHECCAAR TOM
2
3
3
6
O
/
//
O
3
7
4
O
2
4
4
9
4
4
3
4
4
2
4
4
8
9
/
9
3
0
6
2
7
L
CA
2
2
2
7
7
2
/
*
9
3
9
5
3
3
0
2
9
6
9
3
9
3
3
3
2
5
7
2
C
2
a
2
S
9
2
2
2
6
3
8
3
8
3
3
8
3
3
3
3
9
3
3
3
3 3
3
HGt migets kndky wwwpt///ote....... Detcc1
2
3
0
2
2
O
2
2
2
2
2
2
2
C
Figure 4: Image reconstruction. First 6 experiments (rows): image ground truth, measurement at experiment k samples from current prior p(0) Dx-1). with best (resp. worst) weights in upper (resp. lower) sub-row. The samples incorporate past measurement information as the procedure advances. Each design steps takes 7.3s
CONCLUSION
We presented new approach, CoDiff, to gradient-based BOED that allows very efficient implementations. The performance was illustrated in a traditional density-based setting with superior accuracy and lower computational cost compared to state-of-the-art methods. In addition, the possibility of our method to also handle data-based sampling represents, to our knowledge, the first extension of BOED to diffusion- framework of diffusion models for our sampling operators, we were able to optimize a design parameter & concurrently with the diffusion process. This was illustrated in BOED involving high dimensional image parameters. The foundation of our approach lies on a new EIG gradient estimator, bi-level optimization, conditional diffusion models and their application to inverse problems. Thanks to this advancement, there are as many new potential applications of BOED as there are trained diffusion models for specific inverse problem tasks. Current limitations include that CoDiff remains a greedy approach, that it requires an explicit expression of the likelihood and that when using diffusions to address inverse problems only linear forward models are currently handled. However, the non-linear setting is an active field of research, and advancements in this area could be directly applied to our framework. The applicability of our method could also be extended by considering settings with no explicit expression of the likelihood and investigating inference such as developed by Ivanova et al. (2021); Kleinegesse and Gutmann (2021); Kleinegesse et al. (2020). In addition, although in censity-basee BOED, we have shown that greedy approaches could outperform be interesting to investigate an extension to non myopic approaches such as Iqbal et al. (2024).
generative models. By integrating the highly successful
application for
comulation-base
reinforcement learning procedures, in a Mata-based setting, it would
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors thank the reviewers and area chair for their interesting and useful comments and the Inria Challenge project ROAD-AI for partial funding. This work was performed using HPC//I resources from GENCI-IDRIS (Grant
Pierre Alliez is supported by the French government, through the 31A Côte d' Azur Investments in the Future project managed by the National Research Agency ANR-19-P3IA-0002
REFERENCES
Aali, A., Arvinte, M., Kumar, S., and Tamir, I. (2023). Solving inverse problems with score-based generative priors learned from noisy data. In 2023 57th Asilomar Conference on Signals, Systems, and Computers, pages 837-843, IEEE.
Alquier, P. (2024). User-friendly Introduction to PAC-Bayes Bounds. Foundations and Trends in Machine Learning, 17(2):174-333.
Amari, S.-i. (2007). Integration of Stochastic Models by Minimizing a-Divergence. Neural Computation, 19(10):770-2796.
Amzal, B., Bois, F., Parent, E., and Robert, C. P. (2006). Bayesian-Optimal Design via Interacting Particle Systems. Journal of the American Statistical Association, 101(474):773-755
Anderson, B. D. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):33--326.
Ao. Z. and Li, J. (2024). On Estimating the Gradient of the Expected Information Gain in Bayesian Experimental Design. Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages
Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, Hessel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I., King, M., Kunesch, M., Martens, L.,, Merzie, H., Mikulik, V., Norman, T., Papamakarios, G., Quan, J., Ring. R., Ruiz, F,, Sanchez, A., Schneider, R., Sezener, E., Spencer. S., Srinivasan, S., Stokowiec, W., Wang, L., Zhou, G., and Viola, F. (2020). The DeepMind JAX Ecosystem.
Blau, T., Bonilla, E. V., Chades, I., and Dezfouli, A. (2022). Optimizing sequential experimental design with deep reinforcement learning. In Proceedings of the 39th International Conference on Machine Learning (ICML), volume 162, pages 2107-2128. PMLR.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q (2020). JAX: composable transformations of Python+NumPP programs.
Cardoso, G., el idrissi, Y. J., Corff, S. L., and Moulines, E. (2024). Monte Carlo guided Denoising Diffusion models for Bayesian linear inverse problems. In The Twelfth International Conference on Learning Representations (ICLR).
Carvalho, L., Villela, D., Coelho, F., and Bastos, L. (2022). Bayesian Inference for the Weights in Logarithmic Pooling. Bayesian Analysis, 1(1):1-29.
Chaloner, K. and Verdinelli, I. (1995). Bayesian experimental design: A review. Statistical Science, 10(3):273-304
Chatterjee, S. and Diaconis, P. (2018). The sample size required in importance sampling. The Annals of Applied Probability, 2892):099--1135
Chen, W., Zhang, M., Paige, B., Hernández-Lobate J. M., and Barber, D. (2024). Diffusive Gibbs sampling. In Proceedings of the 41st International Conference on Machine Learning (ICML), volume 235, pages 7731-7747. PMLR.
11

---

Published as a conference paper at ICLR 2025
Corenflos, A., Zhao, Z., Särkkä, S., Sjölund, J., and Schön, T. B. (2025). Conditioning diffusion models by explicit Intelligence and Statistics (AISTATS).
bridging. The 28th International conference on Artificial
Dagréou, M., Ablin, P., Vaiter, S., and Moreau, (2022). A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. In Advances in Neural Information Processing Systems.
Daras, G. Chung, H., Lai, C.-H., Mitsufuji, Y.. Milanfar, P., Dimakis, A. G., Ye, C., and Delbracio, M. (2024). survey on diffusion models for inverse problems. githtb. gaituubiiopubbiccttossdffffoooo oytp
Demidovich, Y., Malinovsky, G., Sokolov, I., and Richtárik, P. (2023). A Guide Through the Zoo of Biased SGD. In Thirty-seventh Conference on Neural Information Processing Systems.
Dhariwal, P. and Nichol, A. (2021). Diffusion models beat GANx on image synthesis. Advances in neural information processing systems, 34:8780-8794
Donsker, M. and Varadhan, S. (1976). Asymptotic evaluation of certain Markov process expectations for large time- III. Communications on Pure and Applied Mathematics, 29(4))389 461.
Dou, Z. and Song, Y. (2024). Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference Learning Representations (ICLR).
Drovandi, C. C., McGree, J., and Pettitt, A. N. (2013). Sequential Monte Carlo for Bayesian sequentially designed experiments for discrete data. Computational Statistics & Data Analysis, 57(1):320-333
Efron, B. (2011). Tweedie's formula and selection bias. Journal of the American Statistical Association, 106(4966;100-16611
Foster, A., Ivanova, D. R., Malik, I., and Rainforth, T. (2021). Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 161, pages 3384-3395, PMLR.
Foster, A., Jankowiak, M., Bingham, E., Horsfall, P., Teh, Y. W., Rainforth, T., and Goodman, N. (2019). Variational Bayesian Optimal Experimental Design. In Advances in Neural Information Processing Systems, pages 14059-14077 14070.
Foster, A., Jankowiak, M., O' Meaaa, M., Teh, Y. W. and Rainforth, T. (2020). A Unified Stochastic Gradient Approach to Designing Bayesian-Optinn Experiments. In Proceedings of the 23rd International Conference in Artificial Intelligence and Statistics, volume 108, pages 2959-2969.
Goda, T., Hironaka, T.. Kitade, W., and Foster, A. (2022), Unbiased MLMC Stochastic Gradient- Based Optimization of Bayesian Experimental Designs. SIAM Journal on Scientific Computing. 14(1):A286AA11
Gumbel, E. (1961). Bivariate logistic distribution. Journal American Statistical Association, pages 335-349.
Gutmann, M. and Hyvärinen, A. (2010). Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statisties, pages 297-304. JMLR Workshop and Conference Proceedings.
Hoffman, M. and Gelman, A. (2014). The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(47):153-11623
Hong, M., Wai, H., Wang, Z.,, and Yang, Z. (2023). A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actor-critic. SIAM Journal on Optimization, 18:147-180.
Huan, X., Jagalur, J., and Marzouk, Y. (2024). Optimal experimental design: Formulations and computations. Acta Numerica, 33:715-840.
12

---

Published as a conference paper at ICLR 2025
Hyvärinen, A. (2005). Estimation of Normalzzd Statistical Models by Score Matching. Journal of Machine Learning Research, 6(24):695
Iollo, J., Heinkelé, C., Alliez, P., and Forbes, F. (2024). PASOA- PArticle baSed Bayesian Optimal Adaptive design. Proceedings of the 41st International Conference on Muchine Learning (ICML), volume 235, pages 11020-21046 PMLR.
Iqbal, S., Corenflos, A., Särkkä, S., and Abdulsamad, H. (2024). Nesting particle filters for experimental design in dynamical systems. In Proceedings of the 41st International Conference Machine Learning (ICML), volume 235. PMLR.
Ivanova, D. R., Foster, A., Kleinegesse, S., Gutmann, M. U.. and Rainforth, T. (2021). Implicit deep adaptive design: Policy-based experimental design without likelihoods. Advances in neural information processing systems, 34:25785-25798
Ivanova, D. R., Hedman, M., Guan, C., and Rainforth, (2024). Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design. ICLR 2024 Workshop on Data-centric Machine Learning Research (DMLR).
Kingma, D. and Ba, J. (2015). Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diega, CA. USA.
Kleinegesse, S., Drovandi, C., and Gutmann, M. U. (2020). Sequential Bayesian Experimental Design for Implicit Models via Mutual Information. Bayesian Analysis, 16:773-802.
Kleinegesse, S. and Gutmann, M. U. (2021). Gradient-based Bayesian Experimental Design for Implicit Models using Mutual Information Lower Bounds. ArXiv, abs/2105500437.
Knoblauch, J., Jewson, J., and Damoulas, T. (2022). An Optimization-Centir View on Bayes' Rule: Reviewing and Generalizing Variational Inference. Journal of Machine Learning Research, 23(1).
Korba, A. and Salim, A. (2022). Sampling as first-order optimization over a space of probability measures. Tutorial at the 39th International Conference Machine Learning (ICML).
Kullback, S. (1959). Information Theory and Statistics. Wiley.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998). Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):278-2324.
Liu, Y. and Tajbakhsh, S. D. (2024). Stochastic Optimization Algorithms for Problems with Controllable Biased Oracles. https:/arriv.orgabb/2220...77.
Malik, H. J. and Abraham, B. (1973). Multivariate logistic distributions. Annals of Statistics, 3:588-590.
Marion, P., Korba, A., Bartlett, P., Blondel, M., Bortoli, V. D., Doucet, A., Llinares-Lóper F., Paquette, C., and Berthet, Q. (2025). Implicit diffusion: Efficient optimization through stochastic sampling. In The 28th International conference on Artificial Intelligence and Statistics (AISTATS).
Minka, T. (2005). Divergence measures and message passing. Technical report, Research, Microsoft.
Papamakarios, G., Nalisnick, E.,, Rezende, D. J., Mohamed, S., and Lakshminarayanan. B. (2021). Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(1).
Quan, W., Chen, J., Liu, Y., Yan, D.-M., and Wonka, P. (2024). Deep learning- baasdd image and video inpainting: A survey. International Journal of Computer Vision, 132(7):237-22400.
Rainforth, T., Cornish, R., Yang, H., Warrington, A., and Wood, F. (2018). On nesting Monte Carlo
estimators. In International Conference on Machine Learning (ICML), pages 4267-4276. PMLR.
Rainforth, T., Foster, A., Ivanova, D. R., and Bickford Smith, F. (2024). Modern Bayesian Experimental Design. Statistical Science, 39(1):000-114
13

---

Published as a conference paper at ICLR 2025
Rhee, C.-h. and Glynn, P. W. (2015). Unbiased estimation with square root convergence for SDE models. Operations Research, 63(5):1026-106
Roberts, G. O. and Tweedie, R. L. (1996). Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, 2(4):341-363.
Sebastiani, P. and Wynn, H. (2000). Maximum entropy sampling and optimal Bayesian experimental design. Journal of the Royal Statistical Society: Series B (Statistical Methodology).
Song. Y.. Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole. B. (2021). Score- based generative modeling through stochastic differential equations. In The Ninth International Conference Learning Representations (ICLR).
Wang, Z., Bovik, A. C., Sheikh, H. R., and Simoncelli, E. P. (2004). Image Quality Assessment: From Error Visibility to Structural Similarity. IEEE Transactions on Image Processing. 13(4):000-61.
Xu, M., Quiroz, M., Kohn, R., and Sisson, S. A. (2019). Variance reduction properties of the reparameterization trick. In The 22nd international conference on artificial intelligence and statistics, pages 2711-2720. PMLR.
Yang. J., Ji, K., and Liang, Y. (2021). Provably Faster Algorithms for Bilevel Optimization. In Advances in Neural Information Processing Systems.
Yang, K. K., Wu, Z., and Amold, F. H. (2019). Machine-learmingguudee directed evolution for protein engineering. Nature methods, 16(8):687-694.
14

---

Published as a conference paper at ICLR 2025
A TWO EXPRESSIONS FOR THE EIG GRADIENT
Both approaches presented below, that of Goda et al. (2022) and Ao and Li (2024), start from EIG gradient expressions derived using reparameterization trick. Using the change of variable Y = Te (U)), we can derive the following expression for the EIG gradient,
(E) pu (w))(() [Ve (ogr(Te..(U)) <)] (w) [Ve (ogg((T,((UU)E))
(22)
The first term in (22) involves only the known likelihood and is generally not problematic. For the second term, use,
V Tglog((T,.o()1e Vep(Te,o(U)) P(Te,O(U)(E)
(23)
with p(TT, = En(r) (p(Tg,O(U)10', and (eN(TE,((U)IE) = VeRr(e)) (V(TE,(U)1(0'E
r [ Ve
(24)
E [ x) Velogp(Te, (U))0',E)
(25)
Subsequently, two expressions of the EIG gradient can be derived depending on which of (24) or (25) is used. Using (24) and p(T = Eor)r (p(T :)] it comes
Epery Ve (TT.T
(E((E) po (u)p(0)
V log Logg((Te.oo (
Epcor) [((E,.(U))0'EE)
Eo ( (,,Y00,0))
-pp
9(E, 0(,,Y,
(26)
Epcer) [p(Y((',E)) [
with
- g plogp) (aa_a)))00))))=_p) {)
and
h(E, y. = V (u)
Considering, in the second term, an additional importance distribution g(0'ly. 0, .) leads to the expression used in Goda et al. (2022),
p(0') a(0' Y.O,E) (g. Y 0 0') p(0')
T(((S) Epp 9 Y 0,
(27)
1g 0,E))
pl Y (0
( 0.
It can be used to derive estimators of the form,
N
M M
p(0
h(E, y: 0,, 0/)
M j= a0
1v. E.
Vel(E) ~
g(E. y: O,,
(28)
M
plo )
p( yi 1 E
( v.
where (3), i.) simulated from the joint distribution pe and for each i [0% is a sample from g(-1y,, +,,E). Goda et al. (2022) use (28) with N = 1. Even with perfect sampling, this estimator is not unbiased due to ratio in the second term but can be de-biased following Rhee and Glynn (2015). The randomized MLMC procedure of Rhee and Glynn (2015) is post-hoc general procedure that can be more generally applied to de-bias a sequence of possibly biased estimators, provided the estimators consistent.
N.
Alternatielly, using (25) instead, another expression of the EIG gradient can derived. Replacing (25) in (23), it comes,
(P(T 8) Vglogp(TT p(Teo(U)]E)
VE logg(TT. = Ep(or)
(E(U)))) [Ve (elogr(((.o(())00..))
15

---

Published as a conference paper at ICLR 2025
which, with the definition of above, leads to
(E) =Epg [9(5. Y, 0, E [g(E, Y, 0, 0')]]
(29)
This alternative expression (29) is the starting point of Ao and Li (2024), who subsequently use the following estimator,
Vel(E) Σ lo(e, y. 0.) Becorlu.s)) [o(E, ye, O,, 0)]]
where {3) 1=IN is as before sample from the joint distribution pę and where for each ys- (0'' [y.: {) is tractable approximation the intractable posterior p(0" 13. <). More specifically, Ao and (2024) propose to approximate each posterior distribution by q(0 ly. {) = , Σド 60,,' using sample (0',, h 1:M from MCMC procedure. It follows the nested Monte Carlo estimator below,
Vel(E)
1 N
g y1. 0.)
1 M
Σ gE ys Og 04.))
(30)
B LOGARITHMIC POOLING AS A GOOD IMPORTANCE SAMPLING PROPOSAL
When considering importance sampling with proposal distribution q and target distribution p. Chatterjee and Diaconis (2018) proved that under certain conditions, the number of simulation draws required for both importance sampling and self normalized importance sampling (SNIS) estimators to have small Lı error with high probability was roughly exp(KL(p, g)), see Theorem 1.2 in Chatterjee and Diaconis (2018) for SNIS. Similarly, selecting a proposal distribution which minimizes the importance sampling estimator variance is equivalent to finding distribution with small x² -distance to p. see e.g. Appendix of Minka (2005). More generally, finding a good proposal q is linked to the problem of minimizing a-divergences f-divergence between p and q. which are jointly convex in and see Minka (2005). In this work, we consider KL((g, p) as measure of proximity between and q. This choice ultimately arbitrary but has the advantage of leading to interpretable proposal with interesting sampling properties. To justify the pooled posterior QENN in (9) and its use in (10), then use Lemma below show that for Σ V 1, the distribution q that minimizes the weighted sum the KL against each posterior p(0 yo, E), i.e. E (g. p(0 y. E)) is
N
(0( & p(0) II 0 E)
(31)
Il ((0y,.E))
(32)
which the logarithmic pooling (or geometric mixture) of the respective posterior distributions (Ally: E). Lemma 2 results from application of lemma mentioned by Alquier (2024) (Lemma 2.2 therein), and recalled below in Lemma 1. This Lemma has been known since Kullback (Kullback, 1959) in the case of finite parameter space O but the general case is due Donsker and Varadhan (Donsker and Varadhan, 1976). Recall that P(0) denotes the set of probability measures on θ and p a given probability measure in P(O).
Lemma (Donsker and Varadhan's variational formula) For any measurable, funcctioon f: + + F. the supremum with respect to g E P(O) of
Ey [f( KL(q.p)
is the following Gibbs measure py defined by its density with respect to
(xp(((0))
dpş = E. [exp(f(())] dp.
The following Lemma 2 is application of Lemma 1.
16

---

Published as a conference paper at ICLR 2025
Lemma 2 For given probability measure P()) and measure p (not necessarily probability measure). define for any probability measure q E P(@)
{(g)
0)]] Bo
Kl(g,P)
(33)
results from the Donsker and Varadhan's 's variational formula Lemma 1 that the supremum of {(g)
with respect to q is reached for the Gibbs measure q" defined by its density with respect to p.
q (0) p(0) exp(By=ollogp({|g
In addition maximizing is equivalent to minimising
By [KL(((())
which means that q" the measure that minimizes the KL each p(0ly) on average with respect to
Proof of Lemma The expression of q" results from direct application of Lemma 1 f(0) Ey [log po)p(()))) assuming it is measurable and bounded function of 0 (to be checked in practice). The second part results from rewriting as
log(p( Y 0 (()
((g) Ey-p
Ea-g
log (())
= EE [KL Ey~p log)(Y))
Example: as already mentioned our pooled posterior gg.N corresponds the application of this result to p Eis pgóy, with E = 1. This latter result with E = can also be recovered from more general result by Amari (2007), which stated for any a-divergence (see Theorem of Amari (2007) with a = 1). The continuous weight version is also mentioned by Amari (2007) (Theorem 4), referring to other papers for the proof. We provided here simple proof for the KL (a 1) case.
Remark 1: p by. or p ΣN dy., we recover the standard variational formulation of the posterior distribution (see e.g. Table I in (Knoblauch et al., 2022)). The posterior distribution p(Olyi..... yN) differs from the logarithmic pooling (for which the weights sum to 1) in the relative weight given to the prior. The result valid for very general not necessarily expressed as an expectation.
Remark 2: Regarding logarithmic pooling, the result is similar to a result in Carvalho et al. (2022) (Remark 3.1 therein) by showing that, in the case of the sum of the KL, ΣM KL(g, p(0lp.))))))))) the optimal pooling weights are equal, v; = -
The pooled posterior distribution solution. Indeed, is easy to show that q" is also the measure that minimizes the KL between the joint distribution and a product form approximation where one of the factor fixed p(y).
Remark 3:
also be recovered
constrained mean field
q" = arg min (L(((())(()),p(())))) Y)). (EP(()
C DIFFUSION-BASED GENERATIVE MODELS
C.I DENOISING DIFFUSION MODELS
Given a distribution Po E P()) only available through a set of samples of 9, diffusion models are based the addition of noise to the available samples in such manner that allows to learn the reverse process that "denoises" the samples. This learned process can then be exploited to generate new samples by denoising random noise samples until get back to the original data distribution. As an appropriate noising process, in our experiments we ran the Variance Preserving SDE from Dhariwal and Nichol (2021):
d0 (t)
3(t)
(0)
dt + BB(t)dB
(34)
2
17

---

Published as a conference paper at ICLR 2025
where B(t) >> O is a linear noise schedule that controls the amount of noise added at time t. Solving SDE (34) leads to
N 0 exp( 2 /// B(s)ds), (1- exp(- S. B(s)ds))I)
1
(35)
which can be written as
(4) va. + + āre with 0e - exp(
/ B(s)ds) and
(36)
where E ~ Ao I) is a standard Gaussian random variable. Samples from are transformed to samples approximately from a standard Gaussian distribution after some large time
The reverse denoising process can then be written as the reverse of the diffusion process (34), which as stated by Anderson (1982) is:
[ ((t) g - V log (o dt + B(t)dB..
d
(37)
2
where flows backwards from infinity t-0 and Pt is the distribution of (0) from (34), In practice, the process is started at some large finite T assuming that ((2) NO I). In this Appendix, we rather consider increasing time from 0 to T, using that (37) can be equivalenty written as,
(0 3(T - ( + B(T - Vo log pr dt - -B
(38)
2
which is now initialized with ( (0, I). Solving this reverse SDE, the distribution of ( closed to po for large T, which allows approximate sampling from po-
is
The score function Vo P: (0) of the noisy data distribution time t is intractable and is then estimated by learning a neural network (0, t) with parameters Φb Score matching (Hyvärinen,
2005) is method to train $6 by minimizing the following loss:
(p) (0, () Vo log Pc (0)1
(39)
As Pe (0) is still unknown and only samples from pr (0 in (35) are available, Song et al. (2021) rewrite this loss function as:
E U, (0 (e (())) ((t))l=>(0,t) V p. (0 |2]
(40)
where A(t) > 0 is a weighting function that allows to focus more on certain timesteps than others. It common to take X(t) inversely proportional the variance of (35) at time
Once the neural network so has been trained by minimizing (40), it can be used to generate new samples approximately distributed as the target distribution Po by running numerical scheme on the reverse SDE (38). By running for example the Euler-Marnyama scheme on (38), get the following update step for the reverse process:
BTT t) (()) At 3(T
((++A))
(+))
VB(T t)AtE.
(41)
We can then generate samples approximately from Po by running the reverse process (41) with small enough At.
C.2 CONDITIONAL DIFFUSION MODELS
Conditional diffusion models arise when, for some measurement y, we want produce samples from some conditional distribution po(g))) Sampling from conditional distributions is a problem that arises in inverse problems. When using diffusion models, numerous solutions have been investigated as mentioned in a very recent review (Daras et al., 2024). We specify in this section the approach adopted for our applications. With the application to experimental design in mind, assume here that
Y A
(42)
18

---

Published as a conference paper at ICLR 2025
where 77 ~ (0 is the measurement noise, Ag is the operator that represents the experiment at ç.
Sampling from the conditional distribution p(0) y. E) can be done by running the reverse diffusion process on the conditional SDE:
do (1) = 3(T - 0 + 3(T - V) logpr- (. dt + -B(T-t)dB -o---ddD t)dB,.
(43)
2
with the usual score Ve log (( replaced by the conditionnal score Ve P y. E The main objective of conditional SDE is generate samples from the conditional distribution p( o(Oly.& E) without retraining new neural network s8 for the new conditional score. Writing in terms of the forward process, the conditional score can be written using:
V á po y. E) Vg log .E) + Ve log
(44)
and we can leverage a pre-computed neural network S6 that was trained to estimate the score Vo lo (1)
in the unconditional case. If we know how to evaluate the first term V log (y 0 E), conditional distribution p(0 y, E). Unfortunately, this term does not have a closed form expression.
then the reverse process (43) to generate samples from the
As solution, Dou and Song (2024) propose to approximate the intractable Va log Pt (y 0 .E) by
the tractable Vo Pg (t)
E) where (*) is a noisy version of y at time 1. Then, the following backward SDE can run to generate samples from the conditional distribution p(Oly, E):
3(T t) (1)
(t) (+)
+ B(T (ogr) (T-0 3(T (0 dt
2
+ VB(T-I)dBB
(45)
The sequence of noisy (() can be generated with noising process like (36),
1 Vary v1 ai Age with a₂ = exp( '/' 3(s)ds),
(46)
which using the forward model (42) can be written as:
y (() Ago
+ van.
(47)
We can then evaluate V py ( 0 E as
Va log
E)
A (2
A²O"M)
(48)
σ²āi
C.3 GRADIENT ESTIMATION
When sampling is performed with a finite time horizon diffusion model, the first iterations of the corresponding sampling operator may provide too noisy samples which would result in gradient estimations with little information. One solution proposed by Marion et al. (2025) is to use a queuing trick which requires to store in memory a queue of samples. The memory burden is high and only acceptable for a limited number of particles. We propose a simpler solution, which consists in using Tweedie's 's formula (Efron, 2011) perform one-shot backward step, replacing potentially too
10 (f-(t) o())), which can be interpreted its prediction at time 0. The Tweedie's 's formula provides this prediction in close-form,
noisy g( by the conditional mean E
(U)
o( (t)) + (1 aT Vo log pr-,
(49)
Vār-t
Gradients are then computed using the 0 (t) values, while (()) is updated into (++1) from the backward SDE, as mentioned above. This is only really impactful for small as for large f, 0 (t)
and
g(*) get closer. This extra computation does not add cost as (49) uses a score value that is already computed for (45).
19

---

Published as a conference paper at ICLR 2025
D SEQUENTIAL BAYESIAN EXPERIMENTAL DESIGN
In this framework, experimental conditions are determined sequentially, making use of measurements that are gradually made. This sequential view is referred to as sequential iterated design. In sequential setting, we assume that we plan a sequence of К experiments. For each experiment, we wish pick the best Eş using the data that has already been observed D = (1), (1
Given this design, we conduct an experiment using &.. and obtain outcome yi- Both E₁ and Yı. are then added D for new set D; D (y:: Ex). After each step, our belief about 0 is updated and summarised by the current posterior px0) which acts as the next prior at step comes,
When the observations are assumed conditionally independent,
k
pk)) x p(0) II Pyyn E...)
(50)
and
ply 0IE, Dx-1) x p(0) ply E) II plyn 0. c.)
(51)
A greedy design can be seen choosing each design E if it the last one. This means that &x is chosen as ET. the value that maximizes
c = arg max (E Dx-1)
&
where
(0,
(E 1 Epp
log p(Y E. D pl D 1
MI(PE)
(52)
with y denoting the joint distribution ply. 0|c. Dx-1) = p(y10, Dx-)) Distribution p involves the current prior p(0) (x-). which is not available in closed-form and is not straightforward sample from. Distribution y can be written Gibbs distribution by defining the potential Va as
pe (y 0) exp (-V E,E)
with 0, = log ( logrly
log por((y E..)
= Vyy. 0, (0),
where V(y, 0,
0,g) E) has been already defined in Section 4. Note that the marginal in O of pe is the posterior at step ls or equivalently the current prior p D and the marginal in y is
plyle, D Ep [p(y ( 0, Dx-1)]
Once new Er is computed and new observation ys is performed, the posterior at step k is p(Oly), Ex. Dk-1) which the conditional distribution of plys, 0|c,- DA-1).
E SEQUENTIAL MONTE CARLO (SMC)-STYLE RESAMPLING
SMC an essential addition when dealing with sequential BOED. In density-based BOED, has been already exploited in the sequential context showing real improvement in the quality of the generated samples (Iollo et al., 2024). SMC is also useful in simpler static cases as can improve the quality of the generated 0 and contrastive O' samples, that in turn improves the accuracy of the gradient estimator (30). A particularly central step in SMC the resampling step, first recalled below in the density-based case. Using framework, is it also possible to derive SMC-style resampling scheme the data- basea case. This is becoming a popular strategy in the context generative models (Dou and Song, 2024; Cardoso al., 2024).
20

---

Published as a conference paper at ICLR 2025
Density-based BOED. In static density- BOED, the prior p()) and the likelihhod p(y) E) are available in closed-form. In the sequential experiment context, we want to generate N samples ....... N from the current prior p(0) D and M samples oi..... M from the pooled posterior gg.N As both p(0) Dx-1) and (g.v(0') x II= p(0'ly: E. 1 can be evaluated up to normalizing constant, it is straightforward to extend the sampling operators of Section 4 and add ,
resampling step to the samples .l,..., ON and .i...., O's with weights w; and w/z:
W/
Σ
N
with w; = ((,,)
i=
w/ where ÿ and i are the unnormalized versions of p(0) and (0") respectively.
with w = q(0')
j==1
Data-based BOED. In the setting of data-based BOED, we assume access conditional diffusion model that allows to generate samples from px0 and (0'). The resampling scheme proposed in (Dou and Song, 2024) can be used as is, improve the quality of the samples from (x-1) as this is usual conditional distribution. The resampling scheme is
based on the FPS update: ( first moved using the backward SDE into of(+1) according to ((0((+1))0()) y (-1), Dx-), which satisfies
(53) where (+++)) 1009 DR-1) is given in closed form by the unconditional diffusion model and ((y(T-t-1)) ((++1)
p(0(+11) ( (T ,,
Pa (0((+)))0()) DR-1) (0y(T=-)))) (((+1)) .
E) is given by (47). As both these distributions Gaussian, p(0(t+)) (()), (T (T-111) 1), E. 1 can be written closed form and resampling weights can be
written as:
Wi with 3, = p (-
Wi
(54)
1==1 W;
where P .E) tractable (see Dou and Song (2024) for more details).
For the pooled posterior (0') x II P(O (yi &, Dx-1)" update (53) takes the form:
N
N
II 1=1
p (0 y. (T-t-1), ,, 1 x (t+ @ D Il
ply -T- 1 (((t))
N
1(0
( p(y! (T- 10(0)))))))))))) ,x)
(55)
which leads to the following resampling weights:
w/
w
with
=
I1 ply (T (T--1)
10,00
F NUMERICAL EXPERIMENTS
F.1 SEQUENTIAL PRIOR CONTRASTIVE ESTIMATION (SPCE) AND SEQUENTIAL NESTED MONTE CARLO (SNMC) CRITERIA
The SPCE introduced by Foster et al. (2021) is a tractable quantity to assess the design sequence quality. For number K of experiments, DR (1)1 E.), [y), and I. contrastive variables, SPCE is defined as
K II PY Ex) k=1
SPCE (ER II pl p( 0o) fi p(2z)
log
(56)
K
II
pl Y 0r. Ex.
k=
21

---

Published as a conference paper at ICLR 2025
SPCE is a lower bound of the total EIG which is the expected information gained from the entire sequence of design parameters E1..... Ex and it becomes tight when L tends to oc. In addition, SPCE has the advantage to use only samples from the prior p(0) and not from the successive posterior distributions. It makes it fair criterion to compare methods on design sequences only. Considering true parameter value denoted by 0", given a sequence of design values == observations {yx by>J-=a k=1:K are simulated using p(y .&.) respectively. Therefore, for given DE. the corresponding SPCE is estimated numerically by sampling 0,,,,0 On from the prior,
I1 1( EE))
N Σ
1 SPC = N
log
& Σ
II plyx O2 Ex)
I+I
Similarly, an upper bound on the total EIG has also been introduced by Foster et al. (2021) and named the Sequential nested Monte Carlo (SNMC) criterion,
Il p Ou Ex)
LNMC E 11 pluc & 00)
log
R=
fi p(e,) 2
Σ N (Y+ Ex, ==1=1
As shown in Foster ct al. (2021) (Appendix A), SPCE increases with L to reach the total EIG when - oo at a rate O(L of convergence. It is also shown in Foster et al. (2021) that for a given L, SPCE is bounded by log( [. 1) while the upper bound SNMC below is potentially unbounded. As in Blau et al. (2022), if we use L = 10" to compute SPCE and SNMC, the bound is log(L
(L + 1) = 16.12 for SPCE. In practice this does not impact the numerical methods comparison as the intervals [SPCE, SNMCJ containing the total EIG remain clearly distinct.
F.2
IMPLEMENTATION DETAILS
F.2.1 SOURCE EXAMPLE
For VPCE (Foster et al., 2020) and RL-BOED (Blau et al., 2022), we used the code available at using the settings recommended therein reproduce the results in the respective papers. VPCE optimizes an EIG lower bound in a myopic manner estimating posterior distributions with variational approximations. RL-BOED is non-myopic approach which does not provide posterior distributions. From the obtained sequences of observations and design values, we computed SPCE and SNMC as explained above and retrieved the same results as in their respective papers. For PASOA and SMC procedures, we used the code available at cathuu.coo/iolloj/pssoooo PASOA is a myopic approach, optimizing an EIG lower bound using sequential Monte Carlo (SMC) samplers and tempering to also provide posterior estimations. The method refered to as SMC is variant without tempering.
For CoDiff, the vs's in the pooled posterior distribution were set to V; N The current prior and posterior distributions at experimental step & were initialized using respectively the prior and posterior samples at step k 1. Design optimization was performed using the Adam optimizer with an exponential learning rate decay schedule with initial learning rate 10 and decay rate 0.98. The Langevin step-size in the DiGS method Chen et al. (2024) was set to 10 2 The joint potimization-samplin
loop was rum for 5000 steps. Figure shows samples from the current prior De)), which gradually concentrate around the sources k increases. The additional Figure 6 shows, at some intermediate step k samples from the current prior p(0) Dx-1) and from the pooled posterior distribution in comparison, to illustrate its contrastive nature.
F.2.2 MNIST EXAMPLE
In this example, the likelihood easily derives from Y Ago + n, where Y, 0 and n are familiarly seen as arrays of pixels. The transformation Tg.0 8 is simply Tu.)(U) = Ago + U with = n Gaussian variable. However, to fit in our theoretical framework, images need to be treated as mappings over a continuous spatial domain, so that & can be seen as a continuous parameter and Agf be differentiable with respect to &. This is common in image processing and analysis, where
22

---

Published as a conference paper at ICLR 2025
Particles measurement
Particlss a: meeturement
mensurement
messurement
Figure 5: Source localization example. Experiments 0 (prior samples), 4, 8 and 12. As design locations are selected (orange crosses), samples concentrate to the true sources (red crosses), Samples with lower weights in blue, higher weights in yellow.
2D images are often regarded as mappings over a continuous 2D domain (see egg. Rein van den Boomgaard and Leo Dorst, 2021 Lecture Notes) with ø1 and 22 axes. That is, for (21, E R & = (E1, 22) E R², we consider that an array of pixels 0 is discrete sampled representation of a 2D function 0 (21, .2) and more generally define, using abusively the same notation for continuous and sampled representations,
Y (2) Ag(0(z), x2)) + (71,, 22)
where Ag(-) is masking operator depending some length h and defined by
Ag(0(z1,22)) = 0 ((11.22) if (x1, x2) E Sg.
otherwise
with SE, = ((I1, E - hd, S 2 E + h4, h&2 & D2 S E2 + h42, denoting by A and A2 the sampling distances along the 21 and 3/2
To be able to derive with respect to f, we then need to consider a smooth version He.. of Ag. This is classically done by convoluting with 2D Gaussian kernel G., e.g. the product of two ID Gaussian kernels with positive scales 8 (8), 82). In practice, this consists in smoothing the sharp borders of the mask, noting that fg.a Ag when - 0.
Pg. (21,
(Ac(O) GZ1 22)
(Ao(O) Ga)(11 E1, E2)
= M0.a E1, T2 - -22)
where the second equality uses that (0 x2)) = Ao(0(zi - E1, D2 - E2)). It follows that, no matter what (0) is, such convolution with G2 makes H4,e into function that is continuous and
23

---

Published as a conference paper at ICLR 2025
Current Proo pipPt.1
Pooled Pasteriar qg p/
Serplo Grannd b Lesign
Entrotirc tresst
Current
Ponled Posterior 9p0710p..)
Sumples Ground trah
Castnative
Dengr
Outreittto pats
Cument Prior p1)
Surglce
Co Desiget
2
Figure 6: Several source localisation examples. Prior (left) and pooled posterior (right) samples at experiment k. Final & value (orange cross) the end of the optimization sequence fo Er (blue crosses). This optimization "contrasts" the two distributions by making the pooled posterior "as different as possible" from the prior.
24

---

Published as a conference paper at ICLR 2025
infinitely differentiable in its arguments. Then, it comes, for i 1,2,
0ug. (x1, de,
a
(x1
E1, E2)
de,
0po.
(21 -11,02 E2)
dz,
(Ao(0) x aG&
(x1 -1, 22 - E2)
dz,
These latter derivatives are all that is needed to define the gradients in our developments. If Ao is Heaviside step function, its convolution with a distribution leads to the cumulative density function (cdf) of this distribution. The same developments are valid by replacing the Gaussian kernel by a bivariate logistic distribution L with mean 0 and scales 8 (#1, $2). The multivariate logistic distribution (Gumbel, 1961; Malik and Abraham, 1973) generalizes the univariate logistic distribution. exp((= -212///22)
Its pdf (2(11.02:))
Its edf closed-form and is T+expl - toxp( mm//s2) In practice, we simply consider product of two independent univariate
$1 (1 (1 /11) +exp( ==))==)))
logistic distributions with mean and scale 8,, = 1,2, L (x,;;,),) = (1(1 (atroppp((x_s)))) -=//=))>" The edf of
exp(-2./s)
ID logistic distribution is the sigmoid function S(x,; s:) T+exp(-a./s.) The convolution of Heaviside step function with such a logistic distribution is thus a smooth sigmoid. In the MNIST example, the mask length is set to h = 7, with sampling distances A A2 = 1 and we use 21) product logistic kernel with 81 = 52 = 0.1. Using that the 2D mask Ag can be written as the following product (= < +h) (6:
1) [H (x2 - &2 h)) + [f2 hh-+hhha- x2) 1), where H is the ID Heaviside step function, it follows that the smooth jgg.. is
Me. (21, = (S(2 &1 th - 1) (S(22- +1; 82) S(E2 4- -
For the numerical example of Section 6.3, we used the MNIST dataset (LeCun et al., 1998), the
time varying SDE (34) with noise schedule ((t) = bmin + (bmin bmar> (t to) to) (with bmaz = 5. bmin = 0.2, lo 0, T 2). The training of the usual score matching was done for 3000 epochs with a batch size of 256 and using Adam optimizer Kingma and Ba (2015). We used gradient clipping and the training was done
single A100 GPU.
Update equations for the sampling operators were derived from SDE (19) for the contrastive samples of the pooled posterior ge.v (0') and (45) for samples from the current prior p(O) DB-1), where D can be added in the conditioning part without difficulty. Those updates are equivalent to (55) and (53) respectively. The resampling weights were computed as in Section E.
Figures and 8 show additional image reconstruction processes. The digit to be recovered is shown in the first column. The successively selected masks shown (red line squares) in the second column with the resulting gradually discovered part of the image. The reconstruction per se can be estimated from the posterior samples shown in the last 16 columns. At each experiment, the upper sub-row shows the 16 most-likely reconstructed images, while the lower sub-row shows the 16 less-probable ones. As the number of experiments increases the posterior samples gradually concentrate on the right digit.
Figure then shows that design optimization is effective by showing better outcomes when masks locations are optimized (second column) than when masks are selected at random centers (third column). The highest posterior weight samples in the last 14 columns also clearly show more resemblance with the true digit in the optimized case. The superior performance of design optimization confirmed quantitatively in Table 1, which reports the reconstruction quality as measured by the structural similarity index measure (SSIM) (Wang et al., 2004), for both CoDiff and random design. 20 ground truth digit images are randomly selected and the SSIM is computed for the CoDiff and random reconstructions. after each successive experiment out of 6. Table 1 reports the median SSIM over the 20 selected digits. The SSIM a decimal value between -1 and 1, where indicates perfect similarity, O indicates similarity, and -1 indicates perfect anti-correlation.
25

---

Published as a conference paper at ICLR 2025
\
8
7
3
7
8
4
8
9
8
7
7
8
3
3
\
7
3
9
8
3
8
9
9
\
3
9
7
7
7
9
9
7
9
9
7
7
9
9
2
2
7
4
6
3
Q
2
6
0
6
6
2
S 8
6
6
8
7
4
4
7
9
4
9
7
9
4
7
7
9
9
7
9
NAY DOPER TOM
0
4
8
8
5
2
4
8
J
3
7
6
J
y
4
4
9
4
4
9
4
4
7
4
4
4
4
7
9
4
7
4
8
3
8
7
7
7
4
3
3
4
V
4
C
3
8
4
4
4
4
4
4
4
4
4
4
4
4
4/
4
4
4
9
7
7
7
9
3
2
8
7
8
J
9
4/
8
6
4
4
/
4
4
4
4
4
4
9
9
4
4
4
4
9
4
HS. mm www. w www Indeo So AMA
4
2
9
2
3
2
3
0
3
8
7
//
2
2
3
4
4
4
4
4
4
4
4
9
4
4
4
H
4
4
4
4
7
3
/
9
Q
3
C
8
2
7
3
9
7
7
8
Figure 7: Image reconstruction. First experiments (rows): image ground truth, measurement at experiment k, samples from current prior p(0 D with best (resp. worst) weights in upper (resp. lower) sub-row. The samples incorporate past measurement information as the procedure advances.
F.3 HARDWARE DETAILS
The source example 6.2 can be run locally. It was tested on an Apple M1 Pro 16Gb chip but faster running times can be achieved on GPU. The MNIST example 6.3 was run on single A100 80Gb GPU.
F.4 SOFTWARE DETAILS
Our code is implemented in Jax Bradbury et al. (2020) and uses Flax as Neural Network library and Optax as optimization one Babuschkin et al. (2020). The code is available at http://gghalcooomoooooooooootttt
26

---

Published as a conference paper at ICLR 2025
4
9
4
9
4
6
4
9
4
4
4
4
9
4
G
4
6
9
6
4
9
4
9
O
4
4
4
9
4
9
4
K
4
S<
2
4
4
2
W
+
4
2
4
Hit 1 T00 and fo Nole CORCR da
6
4
0
4
4
5
4
4
6
9
2
4
4
6
4
4
4
4
4
4
4
4
4
4
9
9
S
4
4
6
6
9
&
S
6
6
6
4
6
9
9
6
4
9
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
6
6
9
6
&
4
4
4
6
6
E
6
4
6
8
4
4
7<
<<
7
7
4
U
7<
7<
7
K
4
weliet Se
4
4
4
4
4
4
4
4
4
Z
4
4
//
4
7
7/
7
4
K
4
4
K
K
7
7
7
4
4
X
4
4
2
&
7
4
4
4
4
&
&
2
4
4
4
4
4
4
U
4
4
4
4
4
4
4
4
4
4
4
4
4
Figure 8: Image reconstruction. First 7 experiments (rows): image ground truth, measurement at experiment k, samples from current prior p(0 D with best (resp. worst) weights upper (resp. lower) sub-row. The samples incorporate past measurement information as the procedure advances
4
4
4
4
4
4
4
9
4 4 4
4
4
5
7
4
7
4
4
9
4
4
9
7
S
5
5
5
5
S
5
5
5
5
5
5
5
5
5
S
MCP NANIN Fre fed FOOAN TO Frst WEA RITTIIC POMEO SoAAnd
J
5
V
S
5
S
S
S
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
5
6
1
6
6
6
6
6
6
6
5
5
6
5
5
5
5
5
5
5
5
5
5
5
7
5
5
5
5
5
6
6
6
3
5
3
Figure 9: Image (Ist column) reconstruction from 7 sub-images = Ag +n selected sequentially 7 central pixel &. Optimized vs. random designs: measured outcome y (2nd vs. 3rd column) and parameter 0 estimates (reconstruction) with highest weights (upper vs. lower sub-row).
27


Reviews:
[{"review_id": "yyWUCSGrA2", "reviewer": "Reviewer_7k3H", "summary": "This paper presents a new expression for the gradient of the expected information gain (EIG)---a central quantity in Bayesian optimal experimental design (BOED). Their method, named CoDiff, combines the two steps of maximizing the EIG and then sampling from the posterior in BOED into a single joint step by leveraging the sampling-as-optimization setting of e.g. Marion et al. (2024). Doing so results in a more efficient BOED method, as demonstrated by numerical experiments.", "questions": "See the \"Weaknesses\" section.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper advances the field of BOED by making it more computationally efficient, whilst broadening the scope of problems that can be tackled by BOED by extending it to diffusion-based generative models. The idea of performing joint sampling and optimization using bi-level optimization seems novel to, although I am not very familiar with this literature. The experimental results are significant and demonstrate the efficacy of CoDiff compared to other baselines. The method is presented clearly along with discussions w.r.t. to the related works. Overall, this is a well-written and technically solid paper, and I am happy to recommend its acceptance.", "weaknesses": "Some minor points:\n* It would be nice to see some uncertainty bars around the results in Figure 3.\n* Parts of Section 4 and 5 were a bit difficult for me to follow, but that could just be because I am not familiar with the relevant literature. However, if possible, it might be a good exercise for the authors to think about making these parts a bit more accessible to the BOED community who are not familiar with the sampling-as-optimization literature.\n* It would have been nice to empirically test the performance of CoDiff as e.g. the dimensions of $\\theta$ or the hyperparameters (if any) vary.\n* Please include some limitations of CoDiff in the conclusion section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hCA8Vo9Phh", "reviewer": "Reviewer_j7MS", "summary": "This paper is a little outside of my expertise and I apologize in advance for any misunderstanding I may have and look forward to any correction from the authors. I have set my confidence to low accordingly.\n\nThe authors consider the problem of experimental design, given a statistical model of the form $P(y | \\xi, \\theta)$, we can choose values of $\\xi$ and observe a corresponding $y$, and the goal is to infer $\\theta$. In the first example in the paper, $\\xi$ are locations on a 2D plane, $y$ are the measurements at the $\\xi$ locations, and $\\theta$ represents the source locations of contamination in the plane.\n\nAssuming we have chosen the $\\xi$ values, we can use the $p(y|\\xi)$ to simulate $y_i$ values, one hallucinated rollout into the future, then infer the hallucinated posterior $p(\\theta'|y_i, \\xi)$ and measure the posterior distribution's entropy, $H(p(\\theta'|y_i,\\xi))$, typically requiring Monte-Carlo over the $\\theta'$. We may repeat this process many many times for $i=1,...,N$, and compute the average posterior entropy, wrapping the inner Monte-Carlo over $\\theta'$ in another outer Monte-Carlo over $y$. This average posterior entropy is the Mutual Information (MI) between $y$ and $\\theta$ up to an additive constant (that is the prior entropy $H(p(\\theta))$). This MI quantifies the benefit of sampling $y$ at $\\xi$, we can therefore change $\\xi$ in a way to maximise MI and this is a standard method for experimental design known as Expected Information Gain, or EIG.\n\nHowever as evaluating EIG requires performing nested Monte-Carlo, and optimizing $\\xi$ requires the gradient of EIG, much work has focused on various methods to construct efficient EIG gradient estimators. In past works and this works, the second inner integration over $\\theta'$ is performed with MCMC or by Monte Carlo with importance sampling using a proposal $q(\\theta'|y_i)$ that must be updated for every outer MC iteration $i$.\n\nThis work proposes to use a single proposal distribution across all outer MC iterations $i$ called the \"expected posterior\" which is a geometric mixture of all the hallucinated posteriors\n$$\nq(\\theta') \\propto \\prod_i p(\\theta|y_i, \\xi)^{\\nu_i}\n$$\n\nThe paper proposes to Langevin dynamics to generate samples for the outer Monte Carlo $\\\\{y_i, \\theta_i\\\\}$ from the joint distribution. As these samples are updated in a iterative manner, and the goal is to iteratively optimize $\\xi$, these nested iterative procedures can be merged and updated together, if $\\xi$ only incrementally moves, it is wasteful ti discard all the old samples $\\\\{y_i, \\theta_i\\\\}$ but keep incrementing them so they're up-to-date.\n\nFinally the proposed method is applied to two sequential design applications, the first is the source location example in which the $p(\\theta)$, $p(y|\\theta, \\xi)$ have known analytical forms and the proposed method outperforms baselines. The second is where $p(\\theta)$ cannot be evaluated but can be sampled, specifically a diffusion model trained on MNIST images, $\\theta$ is a full image,  $\\xi$ are pixel locations, and $y$ are the 7 X 7 patch of pixel values around $\\xi$. The goal is to infer the full image by cherry-picking image patches.", "questions": "- Can the authors provide an intuition as to why the proposed expected posterior is preferable to the true expected posterior = prior? Is this purely for the sequential setting?\n- how are the $\\nu_i$ values determined in the expected posterior?\n- what stop other methods from being used on the MNIST example? Is it just the computational cost from multiple nested loops and the sequential use case?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- rigorously cites and discusses related ideas\n- merging multiple nested loops into a single loop is intuitive idea and has been applied in multiple separate fields (training a VAE is effectively the EM algorithm was simplified into a single step, in Bayesian optimzation few-shot Knowledge Gradient)\n- the writing is clear and each idea is easy enough to follow.", "weaknesses": "### Technical Comments\n- __expected posterior__ the paper states that they aim to find one proposal distribution that covers all possible future posteriors and proposes a distribution that is a geometric mixing of hallucinated posteriors $q(\\theta) \\propto \\prod_i p(\\theta|y_i, \\xi)^{\\nu_i}$. The prior is the true expected posterior $p(\\theta) = \\mathbb{E}_Y[p(\\theta|y)]$, exactly the average density of all possible posterior densities.  Why not use the prior $p(\\theta)$ as the one global proposal distribution for integration over all the hallucinated posteriors? In the sequential case (introduced at the end of the paper)\n\n### Writing Comments\n- __Writing Density__ while the writing is thankfully clear, the paper cites and quickly introduces many ideas from prior works in quick succession making the paper very hard work to read for me, a lot of context switching making the paper very tiring. I would suggest reducing the details of referenced works, and providing higher level descriptions where possible and elaborating more on the novel parts of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7vMIfkeYcZ", "reviewer": "Reviewer_VEZf", "summary": "This paper makes an interesting connection between Bayesian optimal experimental design (BOED) and optimization by sampling using a diffusion model. While an interesting connection that may prove fruitful with some refinement, the authors seem to misunderstand the basis of BOED and maximization of a lower bound of mutual information using contrastive learning, as well as miss prior work that might better-inform this method. This disconnect, as well as confusing notation regrading how the diffusion model is used in multi-round BOED and lackluster empirical evaluation, makes me think this work is too early and could benefit from refinement in writing and additional experimental evaluations of the method.", "questions": "My questions are contained in the Weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors make a great connection to optimization by sampling via diffusion models. \n- Thorough gradient analysis.\n- The connection to inverse problems e.g. inpainting in MNIST images as a BOED problem is an interesting new task in the BOED setting and could see BOED objectives (really mutual information objectives) gain further adoption in generative modeling.", "weaknesses": "While the authors make a great connection to optimization by sampling, theoretical analysis of the convergence of the desired goal distribution is missing. This would e.g. especially be helpful in the multi-step BOED setting to show how to improve bounds of the mutual information. I list out my critiques and questions below:\n\n- The authors miss that they are also optimizing a lower bound of the mutual information by the InfoNCE bound, thus their comments that their approximation is exact is incorrect.\n- The definition of information in Equation 2 does not make sense. The KL divergence is written in a non-standard way - do you mean $E_{p(y|\\xi)}KL(p(\\theta|y, \\xi) || p(\\theta))$?\n- The authors miss the connection to Fotster et al. 2019b of LF-ACE that uses a posterior estimate in the InfoNCE bound to refine samples. e.g. in equation 5. \n- The notation for the expected posterior in equation 9, it's not clear what is the proposal distribution $q(\\theta)$ in the paper. Are you using likelihood evaluations? Also, how are you choosing the weighting of $\\nu_i$?\n- I would prefer to see the notation for iterative sampling operators similar to Marion et al. 2024, where the Y, and $\\Theta$ are variables of the operator instead of in the superscript. I think that would make the optimization problem more clear. \n- After equation 15, are you using two generative models to define the observed distribution $y$ and the parameter distribution $\\theta$?\n- in terms of the location finding experiment, usually we would expect to see the designs and final posterior distribution. the way it's portrayed in the paper is confusing. I'm not sure which samples are shown, which design round the designs come from, and why the seeming optimal design isn't on the ground truth values. It looks like the chosen design is quite far from the true theta_0.\n- While the EIG seems to improve over others in Figure 3, the upper bound is _much_ worse. Why is that? \n- The MNIST experiment is interesting but lacks quantitative evaluation. Maybe including a classifier and using the number of successfully classified images would help? Also, there seems to be opportunity to change the design size to see how that influences outcomes.\n\nIn summary, I think the notation is a little bloated. The authors tend to repeat themselves in the text and could be more concise. Concision would help make space for exposition of more salient points that would improve this paper. Additionally, it would help the authors to explicitly lay out what are the likelihood, posterior, and priors they use in experiments. Also, it would be nice to see calibration of the posterior that results of the final design round using e.g. the l-C2ST method or other simulation based calibration methods. Most applications of amortized posteriors using BOED will be interested in this, but it can be computationally expensive, which is why I think just on the final round is sufficient. This is an interesting method that I think needs cleaner exposition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)