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: In this work, the authors extend the DAD method by using implicit models to overcome some weaknesses of DAD. The authors further derive several bounds to design re-parametrization gradients. The authors show that the proposed method outperforms existing methods in several Bayesian experimental design tasks.