NOVAS: Non-convex Optimization via Adaptive Stochastic Search for End-to-end Learning and Control
Authors: Ioannis Exarchos
Marcus Aloysius Pereira
Ziyi Wang
Evangelos Theodorou
Venue: ICLR 2021 Poster (2021)
Keywords: deep neural networks
nested optimization
stochastic control
deep FBSDEs

Abstract: In this work we propose the use of adaptive stochastic search as a building block for general, non-convex optimization operations within deep neural network architectures. Specifically, for an objective function located at some layer in the network and parameterized by some network parameters, we employ adaptive stochastic search to perform optimization over its output. This operation is differentiable and does not obstruct the passing of gradients during backpropagation, thus enabling us to incorporate it as a component in end-to-end learning. We study the proposed optimization module's properties and benchmark it against two existing alternatives on a synthetic energy-based structured prediction task, and further showcase its use in stochastic optimal control applications.

Published as a conference paper at ICLR 2021
NOVAS: NON-CONVEX OPTIMIZATION VIA ADAPTIVE STOCHASTIC SEARCH FOR END-TON TO-END LEARNING AND CONTROL
Ioannis Exarchos'
Marcus A. Pereira; Ziyi Wang Evangelos A. Theodorou
Department of Computer Science Stanford University
School of Aerospace Engineering
Georgia Institute of Technology Atlanta, GA 30332
Stanford, CA 94305
exa Searchos@stanforree
(marcus. perei ræ, ziyiwang, evangelos. theodoroo/@@aatehhhcee
ABSTRACT
In this work we propose the use of adaptive stochastic search building block for general, non-convex optimization operations within deep neural network ar- chitectures. Specifically, for an objective function located at some layer in the network and parameterized by some network parameters, we employ adaptive stochastic search to perform optimization over its output. This operation is dif- ferentiable and does not obstruct the passing of gradients during backpropagation, thus enabling us to incorporate it as a component in end-to-enn learning. We study the proposed optimization module's properties and benchmark it against two exist- ing alternatives on a synthetic energy-based structured prediction task, and further showease its use in stochastic optimal control applications.
INTRODUCTION
Deep learning has experienced drastic increase in the diversity of neural network architectures, both in terms of proposed structure, as well as in the repertoire of operations that define the inter- dependencies of its elements. With respect the latter, significant amount of attention has been devoted to incorporating optimization blocks or modules operating at some part of the network. This has been motivated by large number of applications, including meta-learning (Finn al., 2017 Rusu et al. 2018: Bartunov et al. 2019), differentiable physics simulators (de Avila Belbute- Peres et al. 2018 classification (Amos et al. 2019). GANs (Metz et al. (016). reinforcement learning with constraints, latent spaces, or safety (Amos & Kolter 2017; Srinivas et al. [2088] Amos & Yarats) 2019 Cheng et al. 2019 Pereira Pereira et 2020), model predictive control (Amos et al. [018]] Pereira et al. 2018). as well as tasks relying on the use of energy networks (Belanger et al., [ [unov et al. 2019), among many others. Aocaaoootimmzaatoo modules lead to nested optimization operations, as they interact with the global, end-to-end training of the network that contains them. Consider some component within the neural network architecture, e.g. a single layer, whose input and output I Rn and 2;41 E RT respectively. Within that layer, the input and output linked via the solution of the following optimization problem:
1++1 = arg minF x4, 2,,0),
that is, the output 24+1 #4+1 is defined the solution remains temporarily fixed, i.c., acts as a parameter. Here, F(z;z,, 0) HT x R" x θ + R is function possibly further parameterized by some subset of the neural network parameters Note that here is an independent variable which is free to vary. The result of this optimization could potentially also be subject to a set of (input-dependent) constraints, though in this paper we
optimization problem for which the input d.;
O.
Equua contribution.
'To distinguish between the optimization of the entire network as opposed to that of the optimization mod- ule, we frequently refer to the former as global or outer-loop optimization and to the latter as local or inner-loop optimization.

---

Published as a conference paper at ICLR 2021
will consider only unconstrained optimization. It is also important to note that, depending on the problem, F can be a given function, or it can itself be represented by a multi-layer neural network (trained by the outer loop). in which case the aforementioned optimization layer consists of multiple sub-layers and is more accurately described as a module rather than a single layer. Examples of this type of optimization are structured prediction energy networks (e.g. [Belanger al. (2017)); another such example is Amos & Kolter (2117) which treats the case of convex F(: 0).
In order to facilitate end-to-end learning over the entire network, the gradient of its loss function L with respect to 0 will require during backpropagation passing the gradient of the module's output x²+1 with respect to parameters and z,. Depending on the nature of the optimization problem under consideration, several procedures have been suggested; among them, particularly appealing is the case of convex optimization (Gould et al. Johnson et al., 2016) Amoss et al. 2017] Kolter, 2017), in which the aforementioned gradients can be computed efficiently through an appli- cation of the implicit function theorem to set of optimality conditions, such as the KKT conditions. In the case of non-convex functions however, obtaining such gradients is not as straightfforaad:; solutions involve either forming and solving a locally convex approximation of the problem. or un- rolling gradient descent (Domke 2012 Metz al. 2016 Belanger et al. 2017; Finn et al., 2017, Srinivas et al. 2018 Rusu et al. 2018 Foerster et al. [2018] [2018] Amos et al. 2018). Unrolling gradient descent approximates the arg min operator with fixed number of gradient descent iterations during the forward pass and interprets these as an unrolled compute graph that can be differentiated through during the backward pass. One drawback in using this unrolled gradient descent operation however is the fact that doing so can lead to over-fitting to the selected gradient descent hyper- parameters, such as learning rate and number of iterations. Recently, Amms & Yarats (019)) demonstrated promising results in alleviating this phenomenon by replacing these iterations of gradient descent by iterations of sampling- based optimization, in particular a differentiable approximation of the cross-entropy method. While still unrolling the graph created by the fixed number of iterations, they showed empirically that no over-fitting to the
2017|Amos |mmos &
occurred by performing inference on the trained network with altered inner-loop optimization hyper-pane parameters. Another significant bottleneck in all methods involving graph unrolling is the number of iterations, which has to be kept low to prevent prohibitively large graph during backprop. avoid issues in training.
Note that in eq. I the variable of optimization is free to vary independently of the network. This is in contrast to many applications involving nested optimization, mainly in the field of meta- beta-learning, in which the inner loop, rather than optimizing a free variable, performs adaptation to an initial value which is supplied to the inner loop by the outer part of the network. For example, MAML (Finn) et al. 2017) performs the inner-loop adaptation 0 - 0', in which the starting point 0 is not arbitrary (as x is in eq. (T) but is supplied by the network. Thus, in the context of adaptation, unrolling the inner-loop graph during back-prop is generally necessary to trace the adaptation back to the particular network-supplec
initial value. Two notable exceptions are first-order MAML (Finn et al.) 2017 Nichol et al.) 2018), which ignores second derivative terms, and implicit MAML (Rajeswaran ct al. 2019). which relies on local curvature estimation.
In this paper we propose Non-convex Optimization Via Adaptive Stochastic Search (NOVAS), a module for differentiable, non-convex optimization. The backbone of this module is adaptive stochastic search (Zhou & Huj 2014), a compling-based method within the field of stochastic opti- mization. The contributions of our work are as follows: (A). We demonstrate that the NOVAS mod- ule does not over-fit to optimization hyper-paramedt
and offers improved speed and convergence rate over its alternative (Amos & Yarats 2019). (B). If the inner-loop variable of optimization is free to vary (i.e., the problem fits the definition given by eq. (I), we show that there is no need unroll of gradients. The latter advantage is critical, as it drastically
the graph during the back-propagatie reduces the size of the overall end-to-end computation graph, thus facilitating improved ability to learn with higher convergence rates, improved speed, and reduced memory requirements. Further- more, allows us to use higher number of lnner-loop iterations. (C). If the inner-loop represents an adaptation to a network-suppe network-supplied value as it is the case in meta-learning applications, NOVAS may still be used in lieu of the gradient descent rule (though unrolling the graph may be necessary here). Testing NOVAS in such a setting is left for future work. (D). We combine the NOVAS module with the framework of deep FBSDEs, a neural network-based approach to solving nonlinear partial differential equations (PDEs). This combination allows to solve PDEs of the most general form, i.e., those in which the min operator does not have a closed-form solution, class of problems that previously impossible to address due to the non-convexity of
(HJB)
2

---

Published as a conference paper at ICLR 2021
the corresponding Hamiltonian. We validate the algorithm cart-pole task and demonstrate its scalability on a 101-dimensional continuous-time portfolio selection problem. The code is available at attttts: /ttp//gtttb/btt wwt.e/occccceeeeee www./eexarchos///O/C....... gitt
2 FURTHER BACKGROUND AND RELATED WORK
Relation to Differentiable Cross-Entropy: Particular importance should be given to Amos & Yarats 2019). since, to the best of our knowledge, it is the first to suggest sampling-basel opti- mization instead of gradient descent, and features some similarities with our approach. The authors therein propose a differentiable approximation of the cross-entropy method (CEM) (Rubinstein) 2001 De Boer et al. 2005), called differentiable cross-entropy (DCEM). To obtain this approxi- mation, they need to approximate CEM's s eliteness threshold operation, which is non-differentiable This is done by solving an additional, convex optimization problem separately for each inner loop step (and separately for each sample of z; in the batch, resulting in a total of N x M x K additional convex optimization problems, with N: batch size, M: number of inner loop iterations, K: number of outer loop iterations, i.e. training epochs). After CEM has been locally approximated by DCEM, they replace the usual inner-loop gradient descent steps with DCEM steps, and the entire inner-loop optimization graph is unrolled during the backward pass. Our method differs from this approach in the following ways: 1. we employ the already differentiable adaptive stochastic search algorithm, thus not having to solve any additional optimization problem obtain a differentiablee approxi- mation (speed improvement), while also showing some convergence rate improvements, and most importantly 2. In the case of inner-loop optimization over an independent variable (e.g., such as the problem defined by eq. (I)) we do not unroll the optimization graph, but instead pass the gradients only through the last inner-loop iteration. This drastically reduces its size during backpropagation. increasing speed, reducing memory requirements, and facilitating casier learning.
Sampling-based based Optimization: Adaptive stochastic search (Zhou & Hu 2014) is a sampling-bau -ased method within stochastic optimization that transforms the original optimization problem via a prob- abilistic approximation. The core concept behind this algorithm is approximating the gradient of the objective function by evaluating random perturbations around some nominal value of the indepen- dent variable, a concept that also appears under the name Stochastic Variational Optimization and shares many similarities with natural evolution strategies (Bird et al.. 2018). Another comparable approach is CEM (Rubinstein) 200T: De Boer et al. (205). In contrast to adaptive stochastic search, CEM is non-differentiaale (due to the eliteness threshold) and the parameters are typically updated de novo in each iteration, rather than as a gradient descent update to the parameter values of the previous iteration. In the case of Gaussian distributions, the difference between CEM and adaptive stochastic search boils down to the following: in adaptive stochastic search, the mean gets updated by calculating the average of all sampled variable values weighted by a typically exponential map- ping of their corresponding objective function values, whereas in CEM only the top-k performing values are used, and are weighted equally. Furthermore, this difference can be made even smaller one replaces the exponential mapping in the former method with differentiable (sigmoid) function that approximates the eliteness operation. More details are available in the Appendix.
Deep Learning Approaches for PDEs and FBSDEs: There has been recent surge in research and literature in applying deep learning to approximate solutions of high-dimensional PDEs. The transition from PDE formulation to trainable neural network is done via the concept of a system of Forward-Backwer wadd Stochastic Differential Equations (FBSDEs). Specifically, certain PDE solutions are linked to solutions of FBSDEs. Systems of FBSDEs can be interpreted to a two-point boundary value problem, and can be solved using a suitably defined deep neural network architecture. This is known in the literature as the deep FBSDE approach (Han al. 2018 Raissi 2018). While applied in high-dimensional PDEs, the aforrmentioned results have seen very limited applicability in the field of optimal control. Indeed, the HJB PDE in control theory has a much more complicated structure, and in its general form involves min operator applied on its Hamiltonian term over the control input. Exploiting certain structures of system dynamics and cost functions that allowed for closed-form expression for this operator, Exarchos & Theodorou (2018); Exarchos et al. (2018) 2019) developed a framework for control using FBSDEs, which was then translated to deep neural network setting in Pereira et al. (2019b); Wang et al. 2211). In this work, we incorporate the NOVAS module inside deep FBSDE neural network architectures to account for PDEs lacking a closed-form expression for their min and/or max operators. Thus, we are able to
stochastic equivalent
3

---

Published as a conference paper at ICLR 2021
address the most general description of a HJB PDE in which the corresponding Hamiltonian is non- convex. More information concerning the deep FBSDE framework can be found in the Appendix.
3 NON-CONVEX OPTIMIZATION VIA ADAPTIVE STOCHASTIC SEARCH
The cornerstone of our approach is method within stochastic optimization called adaptive stochas- tic search (Zhou & Hu 2014 Adaptive stochastic search addresses the general maximization problem
x² E arg maxF(x)) SEX
X XCR;
with X' being non-empty and compact, and F X + discontinuous and non-differemtiable function. Instead of dealing with this function that lacks desir- able properties such smoothness, adaptive stochastic search proposes the solution of a stochastic approximation of this problem in which x is drawn from a selected probability distribution fa: p) of the exponential family with parameters p and solve
IR real-valued, non-convex, potentially
p = arg max / F(x)f((;)) = E [F()).
P
This new objective function, due to its probabilistic nature, exhibits desirable properties for opti- mization. Algorithmically, this can be facilitated by introducing natural log and shape function S()) : R + R+ which continuous, non-decreasing, and with a non-negative lower bound (an ex- ample of such function would be the exponential). Due their properties, passing F(x)) through S(-) and the log does not affect the optimal solution. The final optimization problem is then
p arg max f
En [S(((z)))
(3)
To address this optimization problem, one can sample candidate solutions x from [xx; p) the solu- tion space X' and then use gradient ascent method on eq. (3) to update the parameter p. Depending on the chosen probability distribution for sampling x, a closed-form solution for the gradient of the above objective function with respect to p is available. Thus, while still being sampling-base based method at its core, adaptive stochastic search employs gradient ascent on a probabilistic mapping of the initial objective function. While any probability density function of the exponential family will work, in our work we sample from a Gaussian distribution. The resulting update scheme is shown in Alg. I More details concerning its derivation are included in the Appendix.
As described in the introduction, the standard approach employed in the literature for non-convex inner-loop optimization is apply optimization procedure (either gradient descent DCEM) for a fixed number of iterations during the forward pass and interpret these as an unrolled compute graph that can be differentiated through during the backward pass. In this work we argue that this needs to be done only in cases of adaptation (e.g.. in meta-learning): the variable to be adapted supplied to the inner-loop by the outer part of the network, and is adapted using a rule such as gradient descent with respect to the inner-loop objective function for a fixed number of steps. Crucially, the process is not initialized with an arbitrary initial value for the adapted variable but with the one that is supplied by the network; the backward pass which needs to pass through the optimization module also needs to flow through the input of the variable of optimization in the module. Thus, the variable's values pre- and post-adaptation need to be linked via the unrolled computational graph of the adaptation. In contrast, in the case in which the variable of optimization is not an input the layer, but rather is allowed vary freely and independently of the outer part of the network (e.g. as described by problem ([). such a process is not only unnecessary, but further leads to complications such as reduced network trainability and learning speed, as well as increased memory usage and computation time. Given that in this case the variable of optimization initialized by what is typically no more than random guess, there no need to trace the gradients during backpropagation all the way back to that random guess.
After fixing a number of iterations for optimization, say, n, a simple way to implement NOVAS in a unn-unrolled fashion is to take n 1 of the iterations off the graph and perform only the n-th on
While presented for maximization, we deploy it for minimization by switching the sign of the objective function.
4

---

Published as a conference paper at ICLR 2021
Algorithm 1: Non-convex Optimization Via Adaptive Stochastic Search (NOVAS)
Given:
Neural network architecture containing NOVAS module, mini-batch dataset {Xj, Y3
NOVAS objective function F(::2;; 6).
Parameters:
Neural network parameters: number of layers L, layer transformations fi, network parameters 0.
NOVAS parameters: initial mean and standard deviation fu. do. learning rate a, shape function
S, number of samples M, number of iterations N, small positive number 10-3 (sampling variance lower bound).
Set x1 {X,)
for i = to do
if' ft is NOVAS layer then
Set μ f [40, < do
for n 1 to N (off-graph operations) do
(H. σ) t NOVAS Layer (μi J, S, M, N. c, F)
end for
(1, σ) <- Layer_Layer (j, a, a S. M, N, F)
F(+1 = P
else
2i41 f((x;;0) end if
(Y3)
#L+1
Compute Loss: L = MSE(Y,, Y2) Update Parameters: < Adam( (L, 0)
end for
(11, < NOVAS Layer (1,, o, a, S, M, c,
Generate M samples of xm (M, o²).
m 1,..,; M;
for m = 1 to M (vectorized operation) do
Evaluate Fm F(x for maximization Fm p===-- =-F(x" for minimization;
Normalize Fm FmM minm (Fm))/ (max,. (00xmm((F) (Fm) min (F"));
Apply shape function 5Th = =(F")); Normalize S Sm Σ M 1==1= STD;
end for
Update == + EM 5th (."" p),
sqrt (0m= 5t (x
p1) + E):
graph. This amounts essentially to getting a good initial guess solution and performing single step optimization. The latter is enough supply gradient information during back-prop as it relates to the relevant, optimized value x", rather than the intermediate steps forming the trajectory from its initial guess to its final optimal value. From coding perspective, most automatic differentiation packages allow for localized deactivation of gradient information flow; in PyTorch (Paszke et al.) 2019), this simple adding "with torch. no_grad () :" for the to the proposed Alg. [] we would like to mention following: A. The forward pass of the given neural network containing the NOVAS module nested operation of layer transformations fa given by ...))), wherein fa can be either a NOVAS Jlayer or any standard
first iterations. With regards
Y, = S neural network layer. Also, note that the above transformation applies recurrent layers with the forward pass unrolled wherein each fi corresponds to a time step. B. For clarity of presentation, we outline the forward and backward (i.e. backprop) passes for single sampled mini-batch from the training dataset. The user is free to choose a training strategy that best suits the problem and apply the proposed algorithm to every sampled mini-batch with any variant of stochastic gradient descent for the outer loop. C. The "Normalize 1" operation in NOVAS layer is optional, but may lead to some numerical improvements. Further algorithm implementation details are given in the Appendix. Finally, an important remark is necessary concerning the differentiability of the module. Nonconvex optimization problems do not necessarily have a unique optimum, and as a result, the arg min operator is set-valued map and not differentiable in the classical sense; even when the
5

---

Published as a conference paper at ICLR 2021
optimum is unique almost everywhere in the parameter space it is possible for the optimal value to be discontinuous in the parameter. For non-global optimizers, an initialization that avoids wrong local optima is crucial. With respect to the latter, there is some evidence (see, e.g.. Bhhraddwaaa et al. (0200) that sampling-based optimization methods are more robust compared to gradient descent as they evaluate the objective function over an extended area of the input space and, given enough sampling variability for exploration, have thus more chances of escaping narrow local optimum.
4 APPLICATIONS
In this section we explore the properties of NOVAS and test its applicability in few problems involving end-to-end learning. The first task is a Structured Prediction Energy Network (SPEN) leaming task, which we adopted directly from Amos & Yarats (2019). We found it to be ideal environment to test NOVAS against unrolled DCEM and unrolled gradient descent because it is simple, allows for fast training. and, being two-dimensional, one can visualize the results. We would like to stress that this is merely an example for illustrating various algorithm differences and behavior rather than claim on state-of-the-ard results in the domain of SPENs. We then address two optimal control problems by incorporating the NOVAS module in deep FBSDE neural network architecture as shown in Fig. 3] the first problem is the cart-pole swing-up up task, a diveennnniioal problem that has been successfully addressed with already existing deep FBSDE approaches (Pereira al. 2019b) that exploit the structure of dynamics and cost (dynamies are affine in control and the cost quadratic in the control) in order to perform minimization of the Hamiltonian explicitly. This problem merely serves as a means to validate the NOVAS-FBSDE algorithm. The second problem demonstrates the establishment of new -fthe-oftthe-art in solving high-dimensional HJB PDEs using the deep FBSDE method; specifically, we address a continuous-ting problem that leads to general (i.e., without an explicit solution for the min operator) HJB PDE in 101 dimensions. This HJB PDE form could not be addressed by deep FBSDE methods previously.
portfolio optimization
4.1 STRUCTURED PREDICTION ENERGY NETWORKS
The goal in energy-based learning is to estimate a conditional probability PFyla) of an output y E y given input E X' using a parameterized energy function E(z, y; 0) : X x ) x θ + R, wherein E θ are the energy function's 's trainable parameters. The conditional probability is approximated as F(yl=r) xx exp(-E(y;; 0)). Predictions can be made by minimizing the trained energy function with respect to y:
① = arg minE 0).
Initially studied in the context of linear energy functions Taskar et al. 2005 LeCun et al. 2006). the
field recently adopted deep neural networks called structured prediction energy networks (SPENs) (Belanger & McCallum 2016) to increase the complexity of learned energy functions. In particular, Belanger et al. al.((2))) (2017) suggested training SPENs in a supervised manner; unrolled gradient descent used to obtain y which then compared to the ground-truth y² by suitable loss function. Mimicking
the unrolled gradient descent suggested by Belanger et al. (2017), Amos & Yarats (2019) replaced the gradient descent operations with differentiable cross entropy iterations, also using an unrolled computation graph during backpropagation. Here we adopt the same example as in Amos & Yarats (2019) to benchmark NOVAS against unrolled gradient descent and unrolled DCEM, compare their properties, and further show that unrolling the inner-loop graph is not necessary. We consider the simple regression task where ground-truth data are generated from f(x) x (x)), E [0, 2r). and we use a neural network of 4 hidden layers to approximate E. This problem belongs to the class of problems described by eq. in the introduction: the variable of optimization, y. is not an input to the optimization module from the exterior part of the network. The energy function represented by the multi-layer neural network defines the objective function within the module (that is, it corresponds to F(;;4;, 0) of eq. ①I. While in this case the entire network is within the module, the input could instead be features extracted from an anterior part of the network, e.g. through convolution. The results are shown in Figs. and 2 As can be seen from Fig. a(a), unrolled gradient descent converges to a very low loss value (which implies good regression performance), but the trained energy function does not reflect the ground truth relationship. Fig. [a). This implies a "faulty" inner-loop optimization, which the energy network itself learns to compensate for. The result resembles more ordinary regression than energy- structured prediction, since no useful

---

altered inner-loop Published as a conference paper at ICLR 2021
-245
7A0
(a)
(c)
Figure Learned energy functions: plots depict the log of the normalized energy function trained us-
ing 10 inner loop iterations of (a) unrolled gradient descent (reproduced from Amos & Yarats (2019)) (b)
NOVAS with the entire inner-loop graph unrolled, and (c) NOVAS without unrolling (only the last iteration on on-graph). Black/white curve denotes the ground truth solution sin(x).
Convergence
Computation Time
Inference
10
100
10a
U-DCEM
U-NOVAS
Comp. time (ms/iter)
80
10"
10-1
Average Loss
10
60
U-GD
Logg
10-2
U-DCEM
U-NOVAS
NOVAS
10->
10-3
10
10-4
10-5
10-5
100
200
300
400
500
U-GD
U-DCEM U-NOVAS NOVAS Methad
5
10 15 20
Iteration x100
Number of inner loop iterations
(a)
(b)
(c)
Figure 2: Convergence analysis and computation time: (a). convergence (test set loss) and (b), computation
time for unrolled gradient descent (U-GD), unrolled DCEM (U-DCEM), NOVAS in unrolled mode (U-NOVAS) without unrolling (NOVAS). terms speed, NOVAS offers approximate speedup 5x compared U-DCEM, and unrolled on-par with U-GD. (c). Inference on trained models
parameters: All models were trained using 10 inner loop iterations. Unrolled GD causes the energy network over-fit fit to that number (as noted in Amos & Yarats ) Unrolled DCEM does not suffer from this phenomenon, but optimization loss as NOVAS), NOVAS, both regular and unrolled form does not over-fit, and the most efficient.
to be less efficient (more inner-loop iterations required for the
structure is learned; furthermore, changing the inner-loop optimization parameters during inference (after training) leads to an operation that the energy network has not learned to compensate for, as seen in Fig. c). The sampling sameling-bbsed based methods of unrolled DCEM and NOVAS both alleviate this phenomenon by learning the correct energy landscape. Furthermore, as seen in Fig. [bb and (c), unrolling the graph unnecessary, and avoiding leads to a significant speed-up. 5x with respect to unrolled DCEM, Fig. (b). Interestingly, an additional benefit is that NOVAS seems to offer greater inner-loop convergence rate than DCEM (Fig ()). Due to the simplicity of this example, there is no learning inhibition when using an unrolled graph, as seen from the comparison between NOVAS and unrolled NOVAS. However, this can be the case in more complex tasks and network
architectures, as we shall see in Section 4.2 Further details are given in the Appendix.
4.2
CONTROL USING FBSDES
4.2.1
CART-POLE SWING-UP TASK
We first validate the NOVAS-FBSDE algorithm (neural network architecture seen in Fig. 3) by solving a task whose special structure (dynamics affine in control and cost function quadratic in control) allows for a closed-form solution for the min operator of the Hamiltonian. Because of this special structure, this task can be solved by already existing deep FBSDE approaches (e.g., Pereira et al. (001b))) Here, replace the minimization explicit solution with NOVAS. The results shown in Fig. and are accordance with results obtained using explicit minimization (see Fig. in Pereira et al. 2019b)). Equations and implementation details are given in the Appendix.

---

Published as a conference paper at ICLR 2021
Awo
AW
AWN1
X-2
N. Vig. Vam Fax
VeL.N Viun)
x,0
$2,0
A/
LSTM
NOVAS
NOVAS
NOWA Lener
View
Viy
Lave
AJ.
LSIM Lager
LSIM Layer
LSIM Lager
5
Awo
AW N-
AWN4
Figure 3: NOVAS-FBSDE neural network architecture: the NOVAS module is incorporated in the neural network minimize the Hamiltonian over the control at each time step.
0.0
cart position (m)
-0.5
cart pelocity (mL5)
Aelecii (reeley
-1.0
2.0
ongle
1.5
nad veloo
-1.5
-2.0
0.0
0.5
1.0 Time (s) (a)
1.5
0.0
0.5
1.0
1.5
-5
0.0
0.5
1.0
1.5
0.0
0.5
1.0 Time (s) (d)
1.5
Time (s)
Time (s) (c)
(b)
Figure 4: Cartpole: Validation of the NOVAS-FBSDE algorithm on task that allows explicit minimization of the Hamiltonian with respect the control. Replacing the explicit minimization with NOVAS leads to the same solution (compare with Fig. 6 Pereira al. (20T9h)). (a) cart position, (b) pole angle, (c) velocity (d) pole angular velocity. Blue line denotes mean trajectory, shaded regions show controllcr's response to injected noise. Target states are indicated with red. Plots above show statistics over 128 test trials.
4.2.2 HIGH-DIMENSIONAL, CONTINUOUS PORTFOLIO SELECTION
We now demonstrate that augmenting the deep FBSDE method with NOVAS allows us to solve gen- cral, high-dimensional HJB PDEs by employing NOVAS-FBSDD on a stock portfolio optimization problem, defined as follows: select subset M = 20 of those for trading. There also risk-less asset with (relatively low) return rate. We may invest initial wealth capital W among these 20 assets, and the goal is to control the percentage allocation among these assets over time such that the wealth outperforms the market index probability. This optimal control formulation leads to HJB PDE state space of 100 + 1 dimensions (100 stocks of the market plus the wealth process, which incorporates the risk-less asset dynamics. Volatility clearly dominates in such short-term horizons, so successful trading strategy would be one that increases the odds of beating the market average compared to random selection. Index-tracking and vealth-maxiiizzttion have long been the subject study from controls perspective (Primbs) 2007]] Bae et al.) 2020), though with limited results due to the diffi- culty in dealing with such high uncertainties. Primbs (2007) investigates
consider market index / that consists of N 100 stocks, and
wiwddmensional variant only be
of this problem (5 stocks, 3 traded) but does not enforce the constraint that allocation positive (thus leading to negative investments, i.e. borrowing money from stocks), and fracks the in- dex (thus paying penalty also when the portfolio outperforms the index) due to the approach being restricted to consider only quadratic cost functions. We avoid these and enforce positive invest-
ments only by applying softmax the control input, and use (softplus(/ - w))² as cost function to incentivize outperforming rather than tracking the index. We consider two alternative investment strategies baselines: constant and equal allocation among all traded assets. well as random allocations. The results are shown in Fig. [ As indicated by the violin plots, the FBSDE investment strategy outperforms in probability the alternative strategies, and is the only who
8

---

Published as a conference paper at ICLR 2021
performance
Terminal Wealth
0.075
0.050
Total
0.023
0.000
CONSTANT
RAMDOM
FICDS
CONSTANT
ROLDON
Strategies
Strategies
(a)
(b)
Best Case Trajectories
Warst Case Trajectories
1.25
FRIDe
CONSTANT
reSDe
CONSTANT
1.23
Memetted portolte
2.25
3.20
1.30
Time (sears)
[peand)
(c)
(d)
Figure 5: Portfolio Optimization: (a) Violin plots for the comparison of trading strategies in terms of the defined loss. The FBSDE controller outperforms in probability both strategies of equal and random allocation among traded stocks. (b) Wealth market index difference (W I) the end of year for the three
investment strategies. (c) Trajectory corresponding to the lowest cost (best case). (d) Trajectory corresponding to the highest cost (worst case).
outperforms the market index (by almost 5% average) the end of the planning horizon of one year. These statistics are obtained by running 128 market realizations with test sel volatility profiles (noise profiles not during training validation). All equations and implementation details provided in the Appendix. We note that solving this problem in an unrolled graph setting was nor possible: either because was impossible facilitate learning when a high number of inner-loop iterations was applied (presumably due to the excessive total graph depth), or because of memory issues. Thus, eliminating the unrolled graph is absolutely critical in this case.
CONCLUSIONS
In this paper we presented NOVAS, an optimization module for differentiable, non-convex inner- loop optimization that can be incorporated in end-to-end learning architectures involving nested optimizations.
demonstrated its advantages over alternative algorithms such as unrolled gradi- SPEN benchmark. We also showed that NOVAS allows us to expand
ent decent and DCEM the class of PDEs that can be addressed by the deep FBSDE method while resisting the curse of dimensionality, demonstrated by the solution of a 101-dimensional HJB PDE associated with portfolio optimization problem. NOVAS a general purpose differentiable non-convex optimiza- tion approach and thus, owing its broad description and its generality, could be useful in plethora of other applications involving nested optimization operations. We hope that the results of this work will inspire continued investigation.
ACKNOWLEDGMENNTS
This research was supported by AWS Machine Learning Research Awards and NASA Langley. We would also like to thank Brandon Amos for sharing his code for the example of Section 4.1 with us, which allowed to reproduce his results and benchmark all algorithms.

---

Published as a conference paper at ICLR 2021
REFERENCES
Brandon Amos and Zico Kolter. Opinet: Differentiable optimization as a layer in neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 136- 145. JMLR. org, 2017.
Brandon Amos and Denis Yarats. The differentiable cross-entropy method. arXiv preprint arXiv: 2019.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning- Learning-Volume 70, pp. 146-155. JMLR. org, 2017.
Brandon Amos, Ivan Jimenez, Jacob Sacks, Byron Boots, and Zico Kolter. Differentiable MPC for end-to-end planning and control. In Advances in Neural Information Processing Systems, pp. 2018
Brandon Amos, Vladlen Koltun, and Zico Kolter. The limited multi-label projection layer. arXiv preprint arXiv:/906.08877, 2019.
Se Yung Bae, Junkee Jeon, and Hyeng Keun Koo. Continuous-time portfolio selection: A cursory survey. Frontiers in Applied Mathematics and Statistics, 6:4, 2020. ISSN 2297-4687. doi: 10. 2889ffams.2220.0000 Fams 2020. 00004
URL [RL[https: wwww wwontiersin org/ wwt./oorrggeeeeeee http///rtceeeeee...
Sergey Bartunov, Jack W Rae, Simon Osindero, and Timothy P Lillicrap. Meta-learning deep energy-based memory models. arXiv preprint arXiv: 19/0.0220, 2019.
David Belanger and Andrew McCallum. Structured prediction energy networks. In International Conference on Machine Learning, pp. 983-992, 2016.
David Belanger, Bishan Yang, and Andrew McCallum. End-to-end learning for structured prediction energy networks. Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 429-439. JMLR. org. 2017.
Homanga Bharadhwaj. Kevin Xie, and Florian Shkurti. Model-predictive control via cross-entropy and gradient- gradient-based optimization, arXiv preprint arXiv:2004.08763 2020.
Thomas Bird, Julius Kunze, and David Barber. Stochastic variational optimization. arXiv preprint 1809. 2018.
Richard Cheng, Gábor Orosz, Richard M Murray, and Joel W Burdick. End-to-end safe reinforce- ment learning through barrier functions for safety-critical continuous control tasks. Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3387-3395, 2019.
Filipe Avila Belbute-Peres, Kevin Smith, Kelsey Allen, Josh Tenenbaum, and J. Zico Kolter. End-to-end differentiable physics for learning and control. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cianchi, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems 31. pp. 7178- 7189. Curran Associates, Inc., 2018. URL httt:://ppppeeee . nips cc wctpppcppccppee
pdf
Pieter- -Teek De Boer, Dirk P Kroese, Shie Mannor, and Reuven Y Rubinstein. A tutorial on the cross-entropy method. Annals of operations research, 134(1):19-67, 2005.
Justin Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, pp. 318-326, 2012.
Ioannis Exarchos and Evangelos A. Theodorou. Stochastic optimal control via forward and back- ward stochastic differential equations and importance sampling Automatica, 87:159-165. 2018.
Joannis Exarchos, Evangelos A Theodorou, and Panagiotis Tsiotras. Stochastic L oppimal control via forward and backward sampling. Systems & Control Letters, 11::101-108 2018.
10

---

Published as a conference paper at ICLR 2021
Ioannis Exarchos, Evangelos Theodorou, and Panagiotis Tsiotras. Stochastic differential games: A sampling approach via FBSDEs. Dynamic Games and Applications. 9(2):486-500 Jun 2019. ISSN 2153-0793. doi: 10. 10.10077/133355-18--22222
S 15235-018-022
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1126-1135. JMLR. org. 2017.
Jakob Foerster, Richard Y Chen, Maruan Al-Shedivat, Shimon Whiteson, Pieter Abbeel, and Igor Mordatch. Learning with opponent-leaming awareness. In Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, pp. 122-130. International Founda- tion for Autonomous Agents and Multiagent Systems, 2018.
Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. CoRR, abs/ heabs//167.00m .5447, 2016.
Jiequn Han, Arnulf Jentzen, and E Weinan, Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, (15(34):8505-8510) 2018.
Matthew Johnson, David K Duvenaud, Alex Wiltschko, Ryan Adams, and Sandeep Datta. Composing graphical models with neural networks for structured representations and fast infer- ence. In Advances in neural information processing systems, pp. 2946-2954, 2016.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Huang. tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Luke Metz, Ben Poole, David Pfau, and Jascha Sohl Sohl-Dickstein. Unrolled generative adversarial networks. arXiv preprint a 1.00016,, 2016.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint arXiv:1803.02999) 2018
Etienne Pardoux and Shige Peng. Adapted solution of backward stochastic differential equation. Systems & Control Letters, 14(1):55-61, 1990.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary De Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024 8035.
Curran Associates, Inc., 2019,
Marcus Pereira, David D Fan, Gabriel Nakajima An, and Evangelos Theodorou. MPC-inspired neural network policies for sequential decision making. arXiv preprint 1rXiv:1802.058003 2018.
Mareus Pereira, Ziyi Wang, Tianrong Chen, Emily Reed, and Evangelos A Theodorou. Deep 2FBSDEs for systems with control multiplicative noise. arXiv,, pp. arXiv-1906, 2019a.
Marcus A Pereira, Ziyi Wang, Ioannis Exarchos, and Evangelos A Theodorou. Learning deep stochastic optimal control policies using backward SDEs. In Robotics: Science and Systems, 2019b.
Marcus A Pereira, Ziyi Wang, Ioannis Exarchos, and Evangelos A Theodorou. Safe optimal control using stochastic barrier functions and deep forward-backwarr SDEs. arXiv preprint 2.XX::2000.01196, 2020.
J. A. Primbs. Portfolio optimization applications of stochastic receding horizon control. In 2007 American Control Conference, pp. 1811-1816, July 2007. doi: 10. (0.1109/ACC..20072222551
Maziar Raissi.
Forward- backward stochastic neural networks: Deep learning of high nigh-dimensional partial differential equations. arXiv preprint 18440:1804.07010, 2018.
11

---

Published as a conference paper at ICLR 2021
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with im- plicit gradients. In Advances in Neural Information Processing Systems, pp. 113-124, 2019.
Reuven Y Rubinstein. Combinatorial optimization, cross-entropy, ants and rare events. In Stochastic optimization: algorithms and applications, pp. 303-363. Springer, 2001.
Andrei Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osin- dero, and Raia Hadsell. Meta-learning with latent embedding optimization. arXiv preprint arXiv: 00008800770006 2018.
Aravind Srinivas, Allan Jabri, Picter Abbeel, Sergey Levine, and Chelsea Finn. Universal planning networks. arXiv preprint arXiv: 1804.0065, 2018.
Ben Taskar, Vassil Chatalbashev, Daphne Koller, and Carlos Guestrin. Learning structured predic- tion models: A large margin approach. In Proceedings of the 22nd international conference on Machine learning, pp. 896-903, 2005.
Ziyi Wang, Keuntaek Lee, Marcus Pereira, loannis Exarchos, and Evangelos Theodorou. Deep forward-backwarte SDEs for min-max control. In 58th IEEE Conference Decision and Control (CDC), Nice, France, December 11-13, pp. 6807-6814, 2019.
Enlu Zhou and Jiaqiao Hu. Gradient-based adaptive stochastic search for non-differentiable opti- mization. IEEE Transactions Automatic Control, 59(7):1818 1832-18322200 2014.
12

---

Published as a conference paper at ICLR 2021
A APPENDIX
A.1 NOVAS: DERIVATION AND IMPLEMENTATION DETAILS
Given the optimization problem defined in (3) we calculate the gradient with respect to the sampling distribution parameter
[((F())) pf: p)dr
Voln in S(F()))((E
(5)
[ x: xf (a; p)dæ
J ((F(()) V In f(x; (f)
(6)
S(F [=)) p)dr
(f)(F(z)) In f(:;p) [[S(F(=))
where the second equality is obtained using the log trick. The exponential family distribution is characterized by the probability distribution function
f(::)) = ia(a) exp() A(p),,
(8)
where (z)) is the base measure whose formula depends on the particular choice distribution, T(x) is the vector of sufficient statistics and A(p) Inf h(x) (pr()T(())da is the log partition. The gradient with respect to log distribution can then be calculated as
In = T(x) vpA(p).
(9)
A Gaussian (M, 2) is fully defined by mean and covariance, but one can choose to optimize only over the mean μl, and sample using fixed covariance. This results in the following parameters of the distribution:
3
T(x) E-1/2z,
PpA(p) E-
(10)
and for the gradient can be calculated as
V In = E(S(F()))(- m)
(11)
B[S(F(())
One reason for choosing to optimize over the mean only is the simplicity of the resulting update expression, as well as numerical reasons, since applying gradient descent on the covariance matrix can lead to non-positive definiteness if the initial values or the learning rate are not chosen carefully. intermediate solution between using a fixed covariance matrix and its gradient descent-based update law is assuming diagonal covariance matrix and calculating each clement of the diagonal via a simple weighted empirical estimator. The resulting update scheme given in Alg. ! Note that we also investigated using the full gradient descent- based update rule for the covariance, well as the use of the Hessian of the objective function 33) and other techniques like momentum, line search, trainable optimization Syper-paramettr values etc. to speed up convergence, but the results were inconclusive as to their additional benefit. We tested two different shape functions: S(y; k) exp(ky), well shape function suggested by Zhou & Hu (2014), namely S 7, (y /min)/ //1/////11 exp(-x(y - -)))). where Ynii is the minimum of the sampled values and is the n-th biggest value of the assorted y values. For the first choice is numerically advantageous to include the normalization step in the function definition and replace the exponential with the softmax function. The latter choice is a differentiable function approximating the level/indicator function used in CEM. Though both shape functions exhibited similar performance, former was slightly faster, and the latter lead to slightly more accurate results in the SPEN example (but not in the FBSDE example, where they were equivalent). All parameter values such as can be made trainable parameters of the network, though we did not notice improvement in doing so. In fact, the algorithm seems be quite insensitive to parameter values including the learning rate a, with the sole exception of a which does indeed affect its output significantly.
noticed that the
A.2 STOCHASTIC OPTIMAL CONTROL USING FBSDES
In this section we introduce the deep FBSDE framework for solving PDEs and show that com- bining NOVAS with the deep FBSDE allows extend the capabilities of the latter framework
13

---

Published as a conference paper at ICLR 2021
PPereira et al. 2019b Wang et al. 2019) in addressing Stochastic Optimal Control (SOC) prob- lems. The mathematical formulation of a SOC problem leads to a nonlinear PDE, the Hamilton- Jacobi-Bellman PDE. This motivates algorithmic development for stochastic control that combine elements of PDE theory with deep learning. Recent encouraging results (Han et al.] 2018: Raissi 2018) in solving nonlinear PDEs within the deep learning community illustrate the scalability and numerical efficiency of neural networks. The transition from a PDE formulation to a trainable neural network is done via the concept of a system of (FBSDEs). Specifically, certain PDE solutions are linked to solutions of FBSDEs, which are the stochastic equivalent of two-point boundary value problem and can be solved using a suitably de- fined neural network architecture. This is known in the literature as the deep FBSDE approach. In what follows, we will first define the SOC problem and present the corresponding HJB PDE, as well as its associated system of FBSDEs. The FBSDEs are then discretized over time and solved on neural network graph.
Stochastic Differential Equations
Consider SOC problem with the goal of minimizing expected cost functional subject to dynam- ics:
inf [u) = inf E ((((() + I." (a((), (()) di Netr|o,T] s.t. dz(t)= = f(((), u(t)) dt + I(x(t), u(t)) dw((), x(0)= =. E.
(12)
N
(13)
where E R" and u RT are the state and control vectors respectively, f F. x RT R" is a non-linear vector-valued drift function, D : I" x I + gaxe is the diffusion matrix, w E IR" is vector of mutually independent Brownian motions, U is the set of all admissible controls and 1 R" x Rm + R and Ø : R - R are the running and terminal cost functions respectively. Equation
13 is controlled Itō drift-diffusion stochastic process.
Through the value function definition V V(z, t) = inf J (u))aa=z,to and using Bellman's principle of optimality, one derive the Hamilton Jacobi Bellman PDE, given by
V/ +
inf [.
(i + V22 f(x - (xx, ui) 0, V(z,T) = ox)),
(14)
where we drop explicit time dependencies for brevity, and use subscripts to indicate partial deriva- tives with respect to time and the state vector. The term inside the infimum operation is called the Hamiltonian:
1
H(,, u, V2, VaL EE') & (u(V=.EE") V2' f u) + [(x, u).
(15)
Given that a solution u² to the minimization of H exists, the unique solution of (14) corresponds by virtue of the non-linear Feynman-Kac lemma (see for example Pardoux & Peng (990)) to the following system of FBSDEs:
x(1) + L f(((), di L. X(0(0) (()) dwe.
(FSDE)
(16)
S. (x)().)(()) dt S: VT (TT)) u (t), dw, (BSDE)
= v(x(t),t)=
(17)
u²(1)= = arg min H(x(1), u, V=(I(1))), (), Van (x(1) (z(()) ( u))))))))))
arg))argmm))
(18)
Here, V(x((),() l) denotes an evaluation of V(a, 1) along a path of a(l), thus ((x((), 1) is a stochas- tic process (and similarly for V2 (=(x)) t) and Vzx (a(t), t)). Note that (t) evolves forward in time (due to its initial condition (0) = Q). whereas V(x(t), t) evolves backwards in time, due to its terminal condition (xx(T)), thus leading to a system that is similar to two-point boundary value problem. While we can easily simulate forward process by sampling noise and then performing Euler integration, a simple backward integration of V (x((),, E) would result in it depending explicitly on future values of noise, which is not desirable for non-anticipating process, i.e., process that does not exploit knowledge on future noise values. Two remedies exist to mitigate this problem: either mack-propagate the conditional expectation of (x)(t),t) (e.g., as in Exarchos & Theodorou (2018)), ) or proppaape V (2x(t), t) starting from an initial condition guess, compare its termi- nal value V (xTT), T' to the terminal condition, and adjust the initial condition accordingly so that
14

---

Published as a conference paper at ICLR 2021
the terminal condition is satisfied approximately. For this forward evolution of the BSDE, the above system is discretized in time as follows:
Ik+1 1 I& f At + (x Awk,
20 E
(FSDE)
(19)
Vx+1 = V4 - ( u At + V 22 (aRk ut) Awk.
Vo = v,
(BSDE)
(20)
Ux = arg 7 22, Vi,k, VZA Ex u E()k w)²).
(21)
Here, Awn is drawn from (0, At) and H. is given by eq. 15 Note that for every sampled tra- jectory {ar k==1 there is corresponding trajectory {V&] k=1- Under the deep FBSDE controller framework, Vo = v and V2,0 are set to be trainable parameters of a deep neural network that ap- proximates V2 (zr(4), () every time step under forward-propagatiinn
using an LSTM 3 The terminal value of the propagated (=x(t), t). namely (==(T), T), is then compared to o (x)(T) to compute loss function to train the network. Note that since the Hamiltonian can have any arbitrary non-linear dependence on the control, the resulting minimization problem (2I) is generally non-covex and does not have a closed-form solution. Furthermore, it must solved for each time step, and for utiliza- tion within the deep FBSDE controller framework, the non-convex optimizer must be differentiable to facilitate end-to-end learning. This makes NOVAS a good fit. The neural network architecture is shown in Fig. 3 Since the non-convex Hamiltonian minimization procedure is performed at every time step leading to a repeated use of NOVAS the architecture, the ability to avoid unrolling the inner-loop computation graph is crucial.
A.3 FBSDE STOCHASTIC OPTIMAL CONTROL FOR AFHINE-QADDRATIC SYSTEMS
We now show how the previous state-of-the-art (Exarchos & Theodorou 2018 Pereira et al. 2019b)
deals with the problem of the Hamiltonian min operator by assuming a special structure of the problem. Specifically, they restrict the dynamics of eg. [33 to be affine in control, i.c., of the form f(z, u) = F(x) + G(z)u, and the cost in eq. (22) to be quadratic in control, i.c., ((x, u) = +u")+u+ Ru. In this case, and if (x f) is not function of one can perform explicit minimization of the Hamiltonian with respect to u in eq. 21 to find the optimal control:
u R V.-
(22)
This is done by simply setting 0H/Ouu and solving for u. Substituted back into the HJB PDE, this yields a simplified expression without a min operator:
Ve 2
- (V VT + F VT V'V
V(x, = ((x).
Thus, for this restricted class of systems, the deep FBSDE neural neural network architecture does not require a numerical minimization operation over u at every time step, as in eq. (21). The cart-pole swing-up task of the next section is example of system that satisfies these restrictions. similar closed-form solution exists for some cases of L Loptiilll control (Exarchos et al. 2018). as well as some differential games (Exarchos et al., 2019). While simplifying the problem significantly, this approach comes with an important caveat; several dynamical systems do not have control- structure, and penalizing control energy (u" Ru) not always meaningful in every setting.
A.3.1 CART-POLE SWING-UP PROBLEM
We define the state vector to be X [x. 0, i. 0)- where represents the cart-position, 0 repre- sents the pendulum angular-position, å represents the cart-velocity, and 0 represents the pendulum angular-velocity. Let I be the control force applied to the cart. The deterministic equations of motion for the cart-pole system are,
21 + mn sin (1 + cos 0)
m. + m, sin
Ø u cos 0 - m1 cos 0 sin
1(m, + m, sin 0)
In this work, additionally use the L.STM to predict column of the Hessian V.a (x(().
15

---

Published as a conference paper at ICLR 2021
For our experiments, consider the case where noise enters the velocity channels of the state. The stochastic dynamics therefore take the following form,
T
0
Mp sin 0(10 + g cos 0) me + mp sin 0
0
0
dwi
dX = d
dt +
u dt +
m. + mp sin A cos
O
dw2
0
my.0 cos sin 0
0
1(m, + mp sin
1(m mp sin 0)_
The task is to perform swing-up i.e. starting from an initial state of X = [0, 0, 0J' at time to 0. reach the target state of X [0, x, 0, 0] by the end of the time horizon t T. We consider T 1.5s with a time discretization step of At = 0.02s. Notice that the dynamics are affine in control, and selecting the running cost to be 1 u Ru, minimization of the Hamiltonian with respect to assumes a closed-form solution, namely that of eq. (22). This fact allows us to replace the min operator favor of this solution (Pereira al. 9b). Here, we test NOVAS by avoiding this replacement. We consider running and terminal cost matrix of (Qa the control cost matrix of 0.1. The cart-pole parameters considered are mp 0.01 kg. m. 1.0 kg. I = 0.5 m, which are the mass of the pendulum, mass of cart, and length of the pendulum, respectively. For the noise standard deviation, ē 0.5 was used. far as the hyper-parameters for learning the deep FBSDE controller are concerned, we used two-layer L.TTM network as shown in Fig. ([c) with hidden dimension of 16 in each layer, a batch size of 128, and trained the network using the Adam optimizer for 3500 iterations with a learning rate of 5e -3 For the NOVAS layer at every time step, we used 5 inner-loop iterations and 100 samples for both training and inference. A shape function of S cxp(-), initial μı 0, and initial
[0.0, 10.0, 3.0, 0.5) and
10 were used.
With reference to parameters in Alg. [ for this experiment we used 75 time steps, which means that 75 layered fced-forward network when unrolled. Additionally,
the LSTM graph can be viewed at each time step we use a NOVAS Layer to compute the optimal control. Thus, the total number of network layers is L 75 74 149 with 's being NOVAS Layer for i = 4, 6,....
A.3.2 PORTFOLIO OPTIMIZATION PROBLEM
We now consider a problem for which an explicit solution of the Hamiltonian min operator does exist. Let V the total number of stocks that make up an index such that / Σ S. where S; is the stock price process of the i-th stock. Let M be the number of fixed selection of traded stocks taken from those N stocks such that M < N. Furthermore, let u R TRM++ be the control vector. The (N + 1) dimensional state vector is comprised N stock prices and wealth process W. The dynamics of each stock price and wealth process are given by
ak [softmax(w]]
(k: 1,2, .M 1)
(23)
Σ Em=l 64mm
ds,(t) St(t) #4 dt Sc(t) dn: where, i= 1,2, N and dr,, Σ osy dw,
(24)
M+1
M+1
dw (t) Ti rird + Σ Tm pm dt + m=
Σ m==
Tm dmmm
(25)
where T is the fraction of wealth invested in the k-th traded stock, r is rate of return per period of the risk-free asset, 14 is the rate of return of the 7th stock. Here, Øzj denotes the standard deviation of noise terms entering the i-the stock process wherein i j indicates the contribution of the process' own noise as opposed to i A j. which indicates the interaction of noises between stocks (correlation). All wy's are mutually independent standard Brownian motions. To obtain the o's, we used randomly generated synthetic covariance matrices which mimic real stock-market data. Note that the M traded stocks were randomly picked and were not constrained to be any specific sub- selection of the N stocks. Separate noise realizations were used during training, validation, and testing to ensure that the network does not over-fit fit to particular noise profile.
16

---

Published as a conference paper at ICLR 2021
For our experiments N 100 stocks that make up the index I and M 20 traded stocks. We used a scaled squared-softlle function as terminal cost, given by
((x(T)) = ( log (1 ( 2
with 33 10, q 500, and no running cost, focusing investment outperformance at the end of the planning horizon one year. To simulate the stock dynamics we used a time discretization of dt 1/52, which amounts to controls (and thus amounts invested) being applied for a total time of year. The deep FBSDE-NOVAS hyperparameters were as follows: 16 neurons each in two-layer L.TTM network to predict the gradient of the value function each time step, a batch size of 32, an initial learning rate set to le and reduced by factor of 0.1 after 4000 and 4500 training iterations. Training was done using the Adam optimizer for total of 5000 iterations. For NOVAS, we used 100 samples with 5 inner-loop iterations for training and 200 samples with 50 inner-loop iterations for inference. The shape function used was S(x) exp (z).
weekly basis,
With reference to parameters in Alg. | for this experiment we used 52 time steps, which means that the LSTT graph can be viewed as a 52 layered feed-forward network when unrolled. Additionally, at each time step (except for the last time step) use NOVAS_Laye to compute the optimal control. Thus, the total number of network layers I 52 51 103 with f 's being NOVAS_Layer for i 2, 4, 6,....
A.3.3 LOss FUNCTION FOR TRAINING DEEP FBSDE CONTROLLERS
The loss function used in our experiments to train the deep FBSDE controller with the NOVAS layer as follows:
L l H& (V T) V* T)) 12 Ha (Va T) V (FT, T))
+ la HA (VaV T) - (xT, + 14 (VT,,, T))² + ls (V2(TT, la (V2 T))²
where
[a,,
for |a| < 8, otherwise.
Hs(a)
88(2/a) 8),
Here, 2r denotes xTT), V(ET. . T), Va (ZT, T). and (ar. T) are the predicted value function, its predicted gradient, and the predicted last column of the Hessian matrix, respectively, at the terminal time step. The corresponding targets are obtained through the given terminal cost function (<x(T)) so that (²(=T,7) computed by averaging across the batch samples. Additionally, directly minimize the targets. This is possible because gradients flow through the dynamics functions and therefore the weights of the LSTM can influence what the terminal state xTT) will be.
d(тT), V (=T,,T)) 7) = Φz (xT) and (=T T) = Φzz (IT). Each term is
may choose add terms that
For the cart-pole problem we used d = 50 and [1,, 12, l2, 4, 4s (6] = [1, 1, 1, 1, 0J. and for the portfolio optimization problem used 8 50 and [4, 12, 43 14, 15 16] [1, 1, 1,00,
A.3.4 HARDWARE CONFIGURATION AND RUN-TIMES
All experiments were The PyTorch (Paszke et al. 2019) implementation the 101-dimensional portfolio optimization problem had run-tine time of 2.5 hours.
NVIDIA GeForce RTX 2080Ti graphics card with 12GB memory.
17


Reviews:
[{"review_id": "grUp8I9Xtqt", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThis paper aims to present a method that allows efficient learning in neural networks architecture that present optimization blocks. These blocks have the form of x_{i+1} = \\arg \\min_x F(x, x_i, \\theta), and can be thought of as a neural network layer. The addition of this block results in a complex optimization problem, since it presents a multi-level problem. The approach presented in this paper relies on adaptive stochastic search as a differentiable optimization procedure. The authors evaluate the proposed algorithm in a variety of applications, including structured prediction networks and control.\n\n**Assessment**\nWhile the paper is well written, its structure could be significantly improved and the problem statement could be made clearer -- since the topic of the paper is not common, I believe further explanation should be made (it could be by the means of a graph or a figure).\n\n*Pros:* The algorithm presented here it’s simple and seems to lead to good results. It is benchmarked in a variety of fields (1) energy-based learning, (2) robotic control, and (3) portfolio management. In all the cases exceeding the presented baselines.\n\n\n*Cons:* First of all, the paper does not present or mention the practical implementation of their algorithm, which I believe is the only contribution of their work. As it is right now, there is no section of their technical contribution -- just mention of the existing adaptive sampling. Secondly, the baselines introduced, except section 4.1, are non-existent (section 4.2) or extremely naive (section 4.3). Finally, it lacks further experiments underpinning the claims made and ablations. \n\nHere I provide more detailed feedback of my “cons” points. Regarding my first point, since this paper presents a simple-idea/existing-idea-in-another-context, I would expect more details of the practical implementation of their algorithm and sensitivity to the different hyper-parameters. Both of those are lacking in the main paper. The experiment section, while it tackles three different domains, it does not provide enough evidence with respect to previous methods. Section 4.2, depicts a toy reinforcement learning/control problem and no baseline is provided. I would like to (1) incorporate the baseline in Pereira et al. 2019 and LQR with an analysis of the run-time for each algorithm and performance, and (2) test the algorithm to be tested on more complex domains and provide baselines that present the same assumptions, e.g., iLQR. One of the main reasons for the presented method is that it can tackle non-convex objectives; however, the environment presented is convex. Finally, the authors make a lot of emphasis on the difference between their optimization approach and the one that meta-learning does. It seems to me that at least U-NOVA could perform the same task as MAM, or even NOVA with some context of the task. I would like further explanations of this, and some experiments with (U-)NOVA on meta-learning. Finally, I would like further clarifications on why the graph decoupling to get a better initialization is justifiable with NOVA and not other procedures.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "gVM6LPu_Ol3", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary \n\nThe authors present the idea of adaptive stochastic search as a building block for neural networks, as an alternative to other \"inner loop\" optimization methods like gradient descent.\n\n### Reasons for score \n\nPresented works well against the baselines selected and the claims are supported by empirical evidence.\n\n### Pros \n\n1. Paper is clear and grounded in existing literature and context\n1. Claims are supported by empirical evidence.\n1. Proposed method is robust to changes in hyperparameters like the number of inner loop iterations during inference.\n1. Proposed method allows learning the intended functions (see Figure 1) leading to better generalization, the modularity is a bonus.\n\n### Cons \n\n1. There are some cases (although not in general) of meta-learning for adaptation, for example where the update can be described in an implicit fixed point way or when the method used in [FirstOrderMAML](https://lilianweng.github.io/lil-log/2018/11/30/meta-learning.html#first-order-maml) , which allows not needing to unroll (when backpropagating) the function to be optimized. The paper doesn't present this relevant information in page 4. (Edit: this is fixed now)\n1. In Figure 2 clarifying if the loss is train set or test set would be nice, plotting both would be even better. (Edit: this is fixed now)\n1. Not comparing to existing task-adaptation/meta-learning methods and benchmarks (for example First-Order MAML's gradient)\n\n### Suggestions\n\n1. In Figure 1 both (b) and (c) visibly look really similar, I would suggest plotting the difference between them, so a reader can understand better via visualizing the energy function the advantages/disadvantages of (not) unrolling.  (Edit: the authors considered it)", "extra_scores": {}, "extra_text": {}}, {"review_id": "duXZnhy8Kq3", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors suggest a way to backpropagate through neural network with an embedded optimization problem.\nThey propose to perform the estimation of the gradient of the optimum of the embedded problem with respect to its parameterization by the differentiation of one step of a stochastic search algorithm.\nTheir suggested NOVAS approach together with the deep FBSDE method allows for solving optimal control problems efficiently and with low memory due to not unrolling multiple optimizer steps during the backpropagation.\nThe authors demonstrate the feasibility of their approach via an example with a structured prediction energy network and two optimal control tasks (cart-pole swing-up and portfolio selection).\n\nI do not recommend to accept the paper.\n\nThe results are interesting but require a more compelling presentation and in depth analysis.\nAt the current state it is hard for me to understand how their approach for differentiation of an embedded optimization problem is principled and why it should be better than other approaches.\nIt is also hard to evaluate whether the presented approach really is better from the given computational experiments.\n\nFor a locally strictly convex optimization problem of the form\n\nx^star = argmin_x f(x, theta)\n\nthe computation of the Jacobian dx^star / dtheta analytically requires the inversion of the full Hessian matrix via the implicit function theorem:\n\n[partial x^star / partial theta] = -[partial (df/dx) / partial x]^{-1} [partial (df/dx) / partial theta]\n\nwhere [partial (df/dx) / partial x] is the Hessian of the optimization problem at the optimum.\n\nIf x \\in R^n then in the locally strictly convex case the Hessian has full rank n.\n\nBackpropagation through only one gradient step contains only information of rank 1.\nSo the backpropagation of one gradient step generally only contains a small part of the sensitivity of the optimum with respect to theta.\n\nWhen starting the optimization from different points and repeating throughout multiple outer iterations we potentially gather enough sensitivity information to properly represent the shape of the loss function around the optimum in expectation.\nBut the approach certainly adds variance compared to having an exact gradient or unrolling many optimization steps.\n\nAdditionally, the authors suggest a stochastic search method that due to using the log function trick has low sample efficiency compared to using actual gradient information.\nThe stochastic search method proposed by the authors gains information about the loss function shape by sampling locally in the input space.\nThese types of gradient estimators have much higher variance for truly high dimensional loss functions due to the curse of dimensionality.\n\nSee Ben Rechts criticism of the policy gradient / REINFORCE gradient estimator that has the same issue:\nhttps://www.argmin.net/2018/02/20/reinforce/\n\nSo backpropagation of one gradient step with a high variance gradient estimate compared to the true derivative which has rank n curvature information should be inefficient in theory for truly high-dimensional optimization problems.\n\nPerhaps the authors could add at least some analysis of the variance of their gradient estimator compared to other gradient estimators for embedded optimization problems for different examples and show how the variance behaves depending on the difficulty of the optimization task (dimensionality, curvature) taking the above perspective into account.\n\nSince the authors explicitly treat nonconvex optimization problems they could also make more explicit that a nonconvex optimization problem does not necessarily have a unique optimum.\nThe argmin is therefore a set-valued map and not differentiable in the classical sense (even when the optimum is unique almost everywhere in the parameter space it is possible that the optimal value is discontinuous in the parameter).\nFor non-global optimizers (such as gradient descent) starting in such a way that we do not end up in the wrong local optimum is crucial.\n\nIt seems problematic that the comparison against the backpropagation through unrolled gradient descent is with an example where the latter finds a different (wrong) local optimum then the suggested NOVAS method.\nI don't see how the comparison results with respect to convergence, computational time and inference can be considered conclusive without making the tuning effort for the baseline to perform the desired task properly.\n\nThe authors refer to computation time as \"inner-loop convergence rate\" (Page 6) but convergence rate cannot be inferred from computation time if e.g. sampling gradients using the likelihood ratio trick is much faster than computing a backpropagation gradient.\nReverse mode autodiff gradients are only faster than finite difference / forward mode / sampling gradients for a non-zero constant number threshold of input dimension (how many depends on the efficiency of the autodiff module but it can be e.g. 100-200 input parameters).\nThe computation time for truly high-dimensional problems can thus look very different than medium dimensional problems (n = 100).\n\nRegarding the violin plots for the portfolio optimization I am not sure how much they benefit the paper.\nThe goal of the paper should be to demonstrate the superiority of the suggested method over baselines.\nThe violin plots are relevant to people who are interested in the outcome of that specific optimal control task.\nSomeone not knowledgeable about the control problem will not benefit much from knowing the terminal wealth of the different strategies.\n\nMore analysis about the quality of the gradients obtained by the suggested method compared to other methods would improve the paper:\n- How many unrolling steps do what to the gradient variance?\n- What is the difference of using backprop through gradients of the loss vs stochastic search gradients (log likelihood trick) gradients for different dimensions of the optimization variable?\n\nAnother small suggestions:\nOn Page 4 \"desirable properties of optimization\" seems rather vague, perhaps it should be stated more specifically (most likely what is meant is smoothness of the resulting objective function).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "2J-HqSErnnN", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes using adaptive stochastic search as an optimization module within deep neural networks to perform general non-convex optimization.  This is used as a block within deep FBSDEs,which in general do not have a closed form optimization solution, and use the resulting network to show state of the art results in solving high dimensional PDEs on a 101 dimensional portfolio optimization problem.\n\nPros : \n1. Using adaptive stochastic search allows the inner optimization module to take multiple iterations without unrolling the computation graph (unlike meta-learning methods), since the initial value used by the module is arbitrary, and not provided by the network. This greatly reduces memory requirements, and also speeds up learning and convergence. The authors empirically show a 5x speed-up in using their approach compared to unrolled differentiable cross entropy and better qualitative performance than unrolled gradient descent for training an energy function for a simple regression task.\n\n2. Authors confirm that their approach of using the optimization approach within a deep FBSDE works as expected by getting the optimal solution for cartpole, then use their method to beat random and constant strategies on a high dimensional portfolio optimization problem, for which it is not possible to use methods that unroll the computation graph (due to memory issues).\n\n\nCons:\n1. More experiments for Structured Energy Prediction Networks with more challenging functions would give a better indication of the limits of the proposed approach in comparison to prior work. It is also unclear if the proposed approach would show worse performance for non-toy datasets where a few iterations of unrolling gradient descent or differentiable cross entropy are sufficient. This is especially since the only non-toy experiment was performed in a domain where other methods (differentiable cross entropy and gradient descent) couldn't be run.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose an intriguing alternative to IFT or unrolled GD as a method for optimizing through arg min layers in a neural net, by using a differentiable sampling-based optimization approach. I found the general idea in the paper to be intriguing and thought-provoking. The reviewers generally seem to have also appreciated the method, and many of the reviewers' concerns were addressed by the authors during the rebuttal. Although the paper does have a number of flaws -- in particular, the evaluation is a bit hard to appreciate, since improvement over prior work is either unclear, or no meaningful comparison is offered, -- I think in this case the benefits outweigh the downsides. The work is far from perfect, but the ideas that are presented are interested and valuable to the community, and I think that ICLR attendees will appreciate learning about this work. I would encourage the authors however to improve the paper, and especially the empirical evaluation, as much as possible for the camera-ready, and to take reviewer comments into account insofar as feasible. I'm also not sure how much I buy the "overfitting to hyperparameters" argument for unrolled GD, and a less charitable interpretation is that the authors present this issue largely to make up for the comparative lack of other benefits. That's not necessarily a bad thing, but I think making such a big deal of it is a bit strange. It's probably fair to say at this stage that the actual benefits of this approach are a bit modest (though improvements in runtime are a good thing...), but the idea is interesting, and may spur future research.
Decision: Accept (Poster)