The Differentiable Cross-Entropy Method

1 INTRODUCTION . Recent work in the machine learning community has shown how optimization procedures can create new building-blocks for the end-to-end machine learning pipeline ( Gould et al. , 2016 ; Johnson et al. , 2016 ; Amos et al. , 2017 ; Amos & Kolter , 2017 ; Domke , 2012 ; Metz et al. , 2016 ; Finn et al. , 2017 ; Belanger et al. , 2017 ; Rusu et al. , 2018 ; Srinivas et al. , 2018 ; Amos et al. , 2018 ) . In this paper we focus on the setting of optimizing an unconstrained , non-convex , and continuous objective function fθ ( x ) : Rn × Θ → R as x̂ = arg minx fθ ( x ) , where f is parameterized by θ ∈ Θ and has inputs x ∈ Rn . If it exists , some ( sub- ) derivative ∇θx̂ is useful in the machine learning setting to make the output of the optimization procedure end-to-end learnable . For example , θ could parameterize a predictive model that is generating potential outcomes conditional on x happening that you want to optimize over . End-to-end learning in these settings can be done by defining a loss function L on top of x̂ and taking gradient steps ∇θL . If fθ were convex this gradient is easy to analyze and compute when it exists and is unique ( Gould et al. , 2016 ; Johnson et al. , 2016 ; Amos et al. , 2017 ; Amos & Kolter , 2017 ) . Unfortunately analyzing and computing a “ derivative ” through the non-convex arg min here is not as easy and is challenging in theory and practice . No such derivative may exist in theory , it might not be unique , and even if it uniquely exists , the numerical solver being used to compute the solution may not find a global or even local optimum of f . One promising direction to sidestep these issues is to approximate the arg min operation with an explicit optimization procedure that is interpreted as just another compute graph and unrolled through . This is most commonly done with gradient descent as in Domke ( 2012 ) ; Metz et al . ( 2016 ) ; Finn et al . ( 2017 ) ; Belanger et al . ( 2017 ) ; Rusu et al . ( 2018 ) ; Srinivas et al . ( 2018 ) ; Foerster et al . ( 2018 ) . This approximation adds significant definition and structure to an otherwise extremely ill-defined desiderata at the cost of biasing the gradients and enabling the learning procedure to over-fit to the hyper-parameters of the optimization algorithm , such as the number of gradient steps or the learning rate . In this paper we show that the Cross-Entropy Method ( CEM ) ( De Boer et al. , 2005 ) is a reasonable alternative to unrolling gradient descent for approximating the derivative through an unconstrained , non-convex , and continuous arg min . CEM for optimization is a zeroth-order optimizer and works by generating a sequence of samples from the objective function . We show a simple and computationally negligible way of making CEM differentiable that we call DCEM by using the smooth top-k operation from Amos et al . ( 2019 ) . This also brings CEM into the end-to-end learning process in cases where there is otherwise a disconnection between the objective that is being learned and the objective that is induced by deploying CEM on top of those models . We first quickly study DCEM in a simple non-convex energy-based learning setting for regression . We contrast using unrolled gradient descent and DCEM for optimizing over a SPEN ( Belanger & McCallum , 2016 ) . We show that unrolling through gradient descent in this setting over-fits to the number of gradient steps taken and that DCEM generates a more reasonable energy surface . Our main application focus is on using DCEM in the context of non-convex continuous control . This setting is especially interesting as vanilla CEM is the state-of-the-art method for solving the control optimization problem with neural network transition dynamics as in Chua et al . ( 2018 ) ; Hafner et al . ( 2018 ) . We show that DCEM is useful for embedding action sequences into a lower-dimensional space to make solving the control optimization process significantly less computationally and memory expensive . This gives us a controller that induces a differentiable policy class parameterized by the model-based components . We then use PPO ( Schulman et al. , 2017 ) to fine-tune the modelbased components , demonstrating that it is possible to use standard policy learning for model-based RL in addition to just doing maximum-likelihood fitting to observed trajectories . 2 BACKGROUND AND RELATED WORK . 2.1 DIFFERENTIABLE OPTIMIZATION-BASED MODELING IN MACHINE LEARNING . Optimization-based modeling is a way of integrating specialized operations and domain knowledge into end-to-end machine learning pipelines , typically in the form of a parameterized arg min operation . Convex , constrained , and continuous optimization problems , e.g . as in Gould et al . ( 2016 ) ; Johnson et al . ( 2016 ) ; Amos et al . ( 2017 ) ; Amos & Kolter ( 2017 ) , capture many standard layers as special cases and can be differentiated through by applying the implicit function theorem to a set of optimality conditions from convex optimization theory , such as the KKT conditions . Non-convex and continuous optimization problems , e.g . as in Domke ( 2012 ) ; Belanger & McCallum ( 2016 ) ; Metz et al . ( 2016 ) ; Finn et al . ( 2017 ) ; Belanger et al . ( 2017 ) ; Rusu et al . ( 2018 ) ; Srinivas et al . ( 2018 ) ; Foerster et al . ( 2018 ) ; Amos et al . ( 2018 ) ; Pedregosa ( 2016 ) ; Jenni & Favaro ( 2018 ) ; Rajeswaran et al . ( 2019 ) , are more difficult to differentiate through . Differentiation is typically done by unrolling gradient descent or applying the implicit function theorem to some set of optimality conditions , sometimes forming a locally convex approximation to the larger non-convex problem . Unrolling gradient descent is the most common way and approximates the arg min operation with gradient descent for the forward pass and interprets the operations as just another compute graph for the backward pass that can all be differentiated through . In contrast to these works , we show how continuous and nonconvex arg min operations can also be approximated with the cross entropy method ( De Boer et al. , 2005 ) as an alternative to unrolling gradient descent . 2.2 EMBEDDING DOMAINS FOR OPTIMIZATION PROBLEMS . Oftentimes the solution space of high-dimensional optimization problems may have structural properties that an optimizer can exploit to find a better solution or to find the solution quicker than an otherwise naïve optimizer . This is done in the context of meta-learning in Rusu et al . ( 2018 ) where gradient-descent is unrolled over a latent space . In the context of Bayesian optimization this has been explored with random feature embeddings , hand-coded embeddings , and auto-encoder-learned embeddings ( Antonova et al. , 2019 ; Oh et al. , 2018 ; Calandra et al. , 2016 ; Wang et al. , 2016 ; Garnett et al. , 2013 ; Ben Salem et al. , 2019 ; Kirschner et al. , 2019 ) . We show that DCEM is another reasonable way of learning an embedded domain for exploiting the structure in and efficiently solving larger optimization problems , with the significant advantage of DCEM being that the latent space is directly learned to be optimized over as part of the end-to-end learning pipeline . 2.3 RL AND CONTROL . High-dimensional non-convex optimization problems that have a lot of structure in the solution space naturally arise in the control setting where the controller seeks to optimize the same objective in the same controller dynamical system from different starting states . This has been investigated in , e.g. , planning ( Ichter et al. , 2018 ; Ichter & Pavone , 2019 ; Mukadam et al. , 2018 ; Kurutach et al. , 2018 ; Srinivas et al. , 2018 ; Yu et al. , 2019 ; Lynch et al. , 2019 ) , and policy distillation ( Wang & Ba , 2019 ) . Chandak et al . ( 2019 ) shows how to learn an action space for model-free learning and Co-Reyes et al . ( 2018 ) ; Antonova et al . ( 2019 ) embed action sequences with a VAE . There has also been a lot of work on learning reasonable latent state space representations ( Tasfi & Capretz , 2018 ; Zhang et al. , 2018 ; Gelada et al. , 2019 ; Miladinović et al. , 2019 ) that may have structure imposed to make it more controllable ( Watter et al. , 2015 ; Banijamali et al. , 2017 ; Ghosh et al. , 2018 ; Anand et al. , 2019 ; Levine et al. , 2019 ; Singh et al. , 2019 ) . In contrast to these works , we learn how to encode action sequences directly with DCEM instead of auto-encoding the sequences . This has the advantages of 1 ) never requiring the expensive expert ’ s solution to the control optimization problem , 2 ) potentially being able to surpass the performance of an expert controller that uses the full action space , and 3 ) being end-to-end learnable through the controller for the purpose of finding a latent space of sequences that DCEM is good at searching over . Another direction the RL and control has been pursuing is on the combination of model-based and model-free methods ( Bansal et al. , 2017 ; Okada et al. , 2017 ; Jonschkowski et al. , 2018 ; Pereira et al. , 2018 ; Amos et al. , 2018 ; Okada & Taniguchi , 2019 ; Janner et al. , 2019 ; Pong et al. , 2018 ) . Amos et al . ( 2018 ) proposes differentiable MPC and only do imitation learning on the cartpole and pendulum tasks with known or lightly-parameterized dynamics — in contrast , we are able to 1 ) scale our differentiable controller up to the cheetah and walker tasks , 2 ) use neural network dynamics inside of our controller , and 3 ) backpropagate a policy loss through the output of our controller and into the internal components . 3 THE DIFFERENTIABLE CROSS-ENTROPY METHOD ( DCEM ) . We focus on uses of the Cross-Entropy Method ( CEM ) ( De Boer et al. , 2005 ) for optimization in this paper . In this setting , suppose we have a non-convex , deterministic , and continuous objective function fθ ( x ) parameterized by θ over a domain Rn and we want to solve the optimization problem x̂ = arg min x fθ ( x ) ( 1 ) The original form of CEM is an iterative and zeroth-order algorithm to approximate the solution of eq . ( 1 ) with a sequence of samples from a sequence of parametric sampling distributions gφ defined over the domain Rn , such as Gaussians . We refer the reader to De Boer et al . ( 2005 ) for more details and motivations for using CEM and briefly describe how it works here . Given a sampling distribution gφ , the hyper-parameters of CEM are the number of candidate points sampled in each iteration N , the number of elite candidates k to use to fit the new sampling distribution to , and the number of iterations T . The iterates of CEM are the parameters φ of the sampling distribution . CEM starts with an initial sampling distribution gφ1 ( X ) ∈ Rn , and in each iteration t generates N samples from the domain [ Xt , i ] N i=1 ∼ gφt ( · ) , evaluates the function at those points vt , i = fθ ( Xt , i ) , and re-fits the sampling distribution to the top-k samples by solving the maximum-likelihood problem1 φt+1 = arg max φ ∑ i 1 { vt , i ≤ π ( vt ) k } log gφ ( Xt , i ) , ( 2 ) where the indicator 1 { P } is 1 if P is true and 0 otherwise , gφ ( X ) is the likelihood of X under gθ and π ( x ) sorts x ∈ Rn in ascending order so that π ( x ) 1 ≤ π ( x ) 2 ≤ . . . ≤ π ( x ) n. We can then map from the final distribution gφT back to the domain Rn by taking the mean of it , i.e . x̂ = E [ gφT+1 ( · ) ] . Proposition 1 . For multivariate isotropic Gaussian sampling distributions we have that φ = { µ , σ2 } and eq . ( 2 ) has a closed-form solution given by the sample mean and variance of the top-k samples as µt+1 = 1/k ∑ i∈It Xt , i and σ 2 t+1 = 1/k ∑ i∈It ( Xt , i − µt+1 ) 2 , where the top-k indexing set is It = { i : vt , i ≤ π ( vt ) k } . This is well-known in statistics and is discussed in , e.g. , Friedman et al . ( 2001 ) . 1The Cross-Entropy Method ’ s name comes from eq . ( 2 ) more generally optimizing the cross-entropy measure between two distributions . Algorithm 1 DCEM ( fθ , gφ , φ1 ; τ , N , k , T ) DCEM minimizes a parameterized objective function fθ and is differentiable w.r.t . θ . Each DCEM iteration samples from the distribution gφ , starting with φ1 . for t = 1 to T do [ Xt , i ] N i=1 ∼ gφt ( · ) . Sample N points from the domain vt , i = fθ ( Xt , i ) . Evaluate the objective function at those points It = ΠLk ( vt/τ ) . Compute the soft top-k projection of the values with eq . ( 4 ) Update φt+1 by solving the maximum weighted likelihood problem in eq . ( 5 ) end for return E [ gφT+1 ( · ) ] Differentiating through CEM ’ s output with respect to the objective function ’ s parameters with ∇θx̂ is useful , e.g. , to bring CEM into the end-to-end learning process in cases where there is otherwise a disconnection between the objective that is being learned and the objective that is induced by deploying CEM on top of those models . Unfortunately in the vanilla form presented above the top-k operation in eq . ( 2 ) makes x̂ non-differentiable with respect to θ . The function samples can usually be differentiated through with some estimator ( Mohamed et al. , 2019 ) such as the reparameterization trick ( Kingma & Welling , 2013 ) , which we use in all of our experiments . The top-k operation can be made differentiable by replacing it with a soft version as done in Martins & Kreutzer ( 2017 ) ; Malaviya et al . ( 2018 ) ; Amos et al . ( 2019 ) , or by using a stochastic oracle as in ? . Here we use the Limited Multi-Label Projection ( LML ) layer ( Amos et al. , 2019 ) , which projects points from Rn onto the LML polytope defined by Ln , k = { p ∈ Rn | 0 ≤ p ≤ 1 and 1 > p = k } , ( 3 ) which is the set of points in the unit n-hypercube with coordinates that sum to k. Notationally , if n is implied by the context we will leave it out and write Lk . We propose a temperature-scaled LML variant to project onto the interior of the LML polytope with ΠLk ( x/τ ) = arg min 0 < y < 1 −x > y − τHb ( y ) s. t. 1 > y = k ( 4 ) where τ > 0 is the temperature parameter and Hb ( y ) = − ∑ i yi log yi + ( 1 − yi ) log ( 1 − yi ) is the binary entropy function . Equation ( 4 ) is a convex optimization layer and can be solved in a negligible amount of time with a GPU-amenable bracketing method on the univariate dual and quickly backpropagated through with implicit differentiation . We can use the LML layer to make a soft and differentiable version of eq . ( 2 ) as φt+1 = arg max φ ∑ i It , i log gφ ( Xt , i ) subject to It = ΠLk ( vt/τ ) . ( 5 ) This is now a maximum weighted likelihood estimation problem ( Markatou et al. , 1997 ; 1998 ; Wang , 2001 ; Hu & Zidek , 2002 ) , which still admits an analytic closed-form solution in many cases , e.g . for the natural exponential family ( De Boer et al. , 2005 ) . Thus using the soft top-k operation with the reparameterization trick , e.g. , on the samples from g results in a differentiable variant of CEM that we call DCEM and summarize in alg . 1 . We note that we usually also normalize the values in each iteration to help separate the scaling of the values from the temperature parameter . Proposition 2 . The temperature-scaled LML layer ΠLk ( x/τ ) approaches the hard top-k operation as τ → 0+ when all components of x are unique . We prove this in app . A by using the KKT conditions of eq . ( 4 ) . Corollary 1 . DCEM captures CEM as a special case as τ → 0+ . Proposition 3 . With an isotropic Gaussian sampling distribution , the update in eq . ( 5 ) becomes µt+1 = 1/k ∑ i It , iXt , i and σ2t+1 = 1/k ∑ i It , i ( Xt , i − µt+1 ) 2 , where the soft top-k indexing set is It = ΠLk ( vt/τ ) . This can be proved by differentiating eq . ( 5 ) , as discussed in , e.g. , ( De Boer et al. , 2005 ) . As a corollary , this captures prop . 1 as τ → 0+ .

Summary: Various optimization methods can be wrapped to form black-box differentiable deep learning modules. This allows end-to-end learning of energy functions that can be used, for example, in continuous control. There is a whole cottage industry of designing these modules. Advancements in this field are of broad interest to the ICLR community. This paper proposes to unroll the cross entropy method, which is very different than the standard practice of unrolling gradient descent. Experiments on continuous control benchmarks demonstrate that this can be used to learn a latent space in which test-time optimization is performed. By doing this optimization in latent space, it can be performed much faster than in the raw high-dimensional space.