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

1 Introduction . Designing experiments to maximize the information gathered about an underlying process is a key challenge in science and engineering . Most such experiments are naturally adaptive—we can design later iterations on the basis of data already collected , refining our understanding of the process with each step [ 36 , 45 , 51 ] . For example , suppose that a chemical contaminant has accidentally been released and is rapidly spreading ; we need to quickly discover its unknown source . To this end , we measure the contaminant concentration level at locations ξ1 , . . . , ξT ( our experimental designs ) , obtaining observations y1 , . . . , yT . Provided we can perform the necessary computations sufficiently quickly , we can design each ξt using data from steps 1 , . . . , t− 1 to narrow in on the source . Bayesian optimal experimental design ( BOED ) [ 7 , 32 ] is a principled model-based framework for choosing designs optimally ; it has been successfully adopted in a diverse range of scientific fields [ 52 , 58 , 60 ] . In BOED , the unknown quantity of interest ( e.g . contaminant location ) is encapsulated by a parameter θ , and our initial information about it by a prior p ( θ ) . A simulator , or likelihood , model y|θ , ξ describes the relationship between θ , our controllable design ξ , and the experimental outcome y . To select designs optimally , the guiding principle is information maximization—we select the design that maximizes the expected ( Shannon ) information gained about θ from the data y , or , equivalently , that maximizes the mutual information between θ and y . This naturally extends to adaptive settings by considering the conditional expected information gain given previously collected data . The traditional approach , depicted in Figure 1a , is to fit a posterior p ( θ|ξ1 : t−1 , y1 : t−1 ) after each iteration , and then select ξt in a myopic fashion using the one-step mutual information ( see , e.g. , [ 51 ] for a review ) . Unfortunately , this approach necessitates significant computation at each t and does not lend itself to selecting optimal designs quickly and adaptively . Recently , Foster et al . [ 17 ] proposed an exciting alternative approach , called Deep Adaptive Design ( DAD ) , that is based on learning design policies . DAD provides a way to avoid significant computation 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Design Observe T iterations Live experiment Infer ( a ) Traditional BOED : costly computations ( design optimisation and parameter inference ) are required at each iteration . Design Observe T iterations L ive experim ent O ffl in e O ffl in e Infer Policy Critic data deploy ( b ) Policy-based BOED using iDAD : a design policy and critic are learnt before the live experiment . The policy enables quick and adaptive experiments , the critic assists likelihood-free inference . Figure 1 : Overview of adaptive BOED approaches applicable to implicit models . at deployment-time by , prior to the experiment itself , learning a design policy network that takes past design-outcome pairs and near–instantaneously returns the design for the next stage of the experiment . The required training is done using simulated experimental histories , without the need to estimate any posterior or marginal distributions . DAD further only needs a single policy network to be trained for multiple experiments , further allowing for amortization of the adaptive design process . Unfortunately , DAD requires conditionally independent experiments and only works for the restricted class of models that have an explicit likelihood model we can simulate from , evaluate the density of , and calculate derivatives for , substantially reducing its applicability . To address this shortfall , we instead consider a far more general class of models where we require only the ability to simulate y|θ , ξ and compute the derivative ∂y/∂ξ , e.g . via automatic differentiation [ 5 ] . Such models are ubiquitous in scientific modelling and include differentiable implicit models [ 19 ] , for which the likelihood density p ( y|θ , ξ ) is intractable . Examples include mixed effects models [ 15 , 18 ] , various models from chemistry and epidemiology [ 1 ] , the Lotka Volterra model used in ecology [ 19 ] , and models specified via stochastic differential equations ( such as the SIR model [ 10 ] ) . To perform rapid adaptive experimentation with this large class of models , we introduce implicit Deep Adaptive Design ( iDAD ) , a method for learning adaptive design policy networks using only simulated outcomes ( see Figure 1b ) . To achieve this , we introduce likelihood-free lower bounds on the total information gained from a sequence of experiments , which iDAD utilizes to learn a deep policy network . This policy network amortizes the cost of experimental design for implicit models and can be run in milliseconds at deployment-time . To train it , we show how the InfoNCE [ 57 ] and NWJ [ 37 ] bounds , popularized in representation learning , can be applied to the policy-based experimental design setting . The optimization of both of these bounds involves simultaneously learning an auxiliary critic network , bringing an important added benefit : it can be used to perform likelihood-free posterior inference of the parameters given the data acquired from the experiment . We also relax DAD ’ s requirement for experiments to be conditionally independent , allowing its application in complex settings like time series data , and , through innovative architecture adaptations , also provide improvements in the conditionally independent setting as well . This further expands the model space for policy-based BOED , and leads to additional performance improvements . Critically , iDAD forms the first method in the literature that can practically perform real-time adaptive BOED with implicit models : previous approaches are either not fast enough to run in real-time for non-trivial models , or require explicit likelihood models . We illustrate the applicability of iDAD on a range of experimental design problems , highlighting its benefits over existing baselines , even finding that it often outperforms costly non-amortized approaches . Code for iDAD is publicly available at https : //github.com/desi-ivanova/idad . 2 Background . The BOED framework [ 32 ] begins by specifying a Bayesian model of the experimental process , consisting of a prior on the unknown parameters p ( θ ) , a set of controllable designs ξ , and a data generating process that depends on them y|θ , ξ ; as usual in BOED , we assume that p ( θ ) does not depend on ξ . In this paper , we consider the situation where y|θ , ξ is specified implicitly . This means that it is defined by a deterministic transformation , f ( ε ; θ , ξ ) , of a base ( or noise ) random variable , ε , that is independent of the parameters and the design ; e.g. , ε ∼ N ( ε ; 0 , I ) . The function f is itself often not known explicitly in closed form , but is implemented as a stochastic computer program ( i.e . simulator ) with input ( θ , ξ ) and ε corresponding to the draws from the underlying random number generator ( or equivalently the random seed ) . Regardless , the resulting induced likelihood density p ( y|θ , ξ ) is still generally intractable , but sampling y|θ , ξ is possible . Having acquired a design-outcome pair ( ξ , y ) , we can quantify the amount of information we have gained about θ by calculating the reduction in entropy from the prior to the posterior . We can further assess the quality of a design ξ before acquiring y , by computing the expected reduction in entropy with respect to the marginal distribution of the outcome , p ( y|ξ ) = Ep ( θ ) [ p ( y|θ , ξ ) ] . The resulting quantity , called the expected information gain ( EIG ) , is of central interest in BOED and is defined as I ( ξ ) : = Ep ( θ ) p ( y|θ , ξ ) [ log p ( θ|ξ , y ) p ( θ ) ] = Ep ( θ ) p ( y|θ , ξ ) [ log p ( y|θ , ξ ) p ( y|ξ ) ] . ( 1 ) Note that I ( ξ ) is equivalent to the mutual information ( MI ) between the parameters θ and data y when performing experiment ξ . The optimal ξ is then the one that maximises the EIG , i.e . ξ∗ = argmaxξ I ( ξ ) . Performing this optimization is a major computational challenge since the information objective is doubly intractable [ 46 ] . For implicit models , the cost becomes even greater as the likelihood is also not available in closed form , so estimating it , along with the marginal likelihood p ( y|ξ ) , is already itself a major computational problem [ 11 , 20 , 33 , 54 ] . Jointly optimizing the design variables for all undertaken experiments at the same time using ( 1 ) is called static experimental design . In practice , however , we are often more interested in performing multiple experiments adaptively in a sequence ξ1 , . . . , ξT , so that the choice of each ξt can be guided by past experiments , namely the corresponding history ht−1 : = { ( ξi , yi ) } i=1 : t−1 . The typical approach in such settings is to sequentially perform ( approximate ) posterior inference for θ|ht−1 , followed by a one-step look ahead ( myopic ) BOED optimization that conditions on the observed history . In other words , to determine the designs ξ1 , . . . , ξT , we sequentially optimize the objectives Iht−1 ( ξt ) : = Ep ( θ|ht−1 ) p ( yt|θ , ξ , ht−1 ) [ log p ( yt|θ , ξ , ht−1 ) p ( yt|ξ , ht−1 ) ] , t = 1 , . . . , T. ( 2 ) However , such approaches incur significant computational cost during the experiment itself , particularly for implicit models [ 16 , 21 , 30 ] . This has critical consequences : in most cases they can not be run in real-time , undermining one ’ s ability to use them in practice . 2.1 Policy-based adaptive design with likelihoods . For tractable likelihood models , Foster et al . [ 17 ] proposed a new framework , called Deep Adaptive Design ( DAD ) , for adaptive experimental design that avoids expensive computations during the experiment . To achieve this , they introduce a parameterized deterministic design function , or policy , πϕ that takes the history ht−1 as input and returns the design ξt = πϕ ( ht−1 ) to be used for the next experiment as output . This set-up allows them to consider the objective IT ( πϕ ) = Ep ( θ ) p ( hT |θ , πϕ ) [ T∑ t=1 Iht−1 ( ξt ) ] , ξt = πϕ ( ht−1 ) , ( 3 ) which crucially depends on the policy π rather than the individual design ξt . Learning a policy up-front , rather than designs , is what allows adaptive experiments to be performed in real-time . Under the assumption that yt is independent of ht−1 conditional on the parameters θ and the design ξt , i.e . p ( yt|θ , ξt , ht−1 ) = p ( yt|θ , ξt ) , Foster et al . [ 17 ] showed that the objective can be simplified to IT ( πϕ ) = Ep ( θ ) p ( hT |θ , πϕ ) [ log p ( hT |θ , πϕ ) p ( hT |πϕ ) ] , p ( hT |θ , πϕ ) = T∏ t=1 p ( yt|θ , ξt ) . ( 4 ) To deal with the marginal p ( hT |πϕ ) in the denominator , they then derived several optimizable lower bounds on IT ( πϕ ) , such as the sequential Prior Contrastive Estimation ( sPCE ) bound LsPCET ( πϕ , L ) = Ep ( θ0 ) p ( hT |θ , πϕ ) p ( θ1 : L ) [ log p ( hT |θ0 , πϕ ) 1 L+1 ∑L ℓ=0 p ( hT |θℓ , πϕ ) ] ≤ IT ( πϕ ) ∀L ≥ 1 . ( 5 ) The parameters of the policy ϕ , which takes the form of a deep neural network , are now learned prior to the experiment ( s ) using stochastic gradient ascent on this bound with simulated experimental histories . Design decisions can then be made using a single forward pass of πϕ during deployment . Unfortunately , training the DAD network by optimizing ( 5 ) requires the likelihood density p ( hT |θ , π ) to be analytically available—an assumption that is too restrictive in many practical situations . The architecture for DAD also assumes conditionally independent designs , which is unsuitable in some settings like time-series data . Our method lifts both of these restrictions .

Summary: This paper proposes a method for Bayesian Optimal Experiment Design (BOED) which uses a policy network to decide which experiment to perform next based on the history of past experiments and their outcomes. The goal is to enable experiments to be chosen quickly for real-time applications. The method trains a policy network in simulation, and does not require assumptions necessary for a prior method (DAD), which also trains a policy network in simulation but requires additional assumptions such as the availability of a likelihood model of the history given the parameters and policy. The present approach removes this requirement by using variational approximations together with a reparameterization which enables optimization with gradient descent. Experiments are done in a 2D location finding environment, a pharmacokinetics model and an epidemiology simulation. Results show that the proposed method almost matches the performance of an exact method which has access to exact likelihood on the first task, and does better than baselines on all the tasks while being fast to evaluate.