Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: **Summary**
This 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.

**Assessment**
While 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).

*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.


*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. 

Here 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.

review_title: Official Blind Review #3

final_score: rating: 6: Marginally above acceptance threshold
review_content: ### Summary 

The 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.

### Reasons for score 

Presented works well against the baselines selected and the claims are supported by empirical evidence.

### Pros 

1. Paper is clear and grounded in existing literature and context
1. Claims are supported by empirical evidence.
1. Proposed method is robust to changes in hyperparameters like the number of inner loop iterations during inference.
1. Proposed method allows learning the intended functions (see Figure 1) leading to better generalization, the modularity is a bonus.

### Cons 

1. 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)
1. 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)
1. Not comparing to existing task-adaptation/meta-learning methods and benchmarks (for example First-Order MAML's gradient)

### Suggestions

1. 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)
review_title: Great work!

final_score: rating: 6: Marginally above acceptance threshold
review_content: 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.

Pros : 
1. 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.

2. 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).


Cons:
1. 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.
review_title: Efficient Inner Optimization module for networks

final_score: rating: 4: Ok but not good enough - rejection
review_content: The authors suggest a way to backpropagate through neural network with an embedded optimization problem.
They 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.
Their 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.
The 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).

I do not recommend to accept the paper.

The results are interesting but require a more compelling presentation and in depth analysis.
At 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.
It is also hard to evaluate whether the presented approach really is better from the given computational experiments.

For a locally strictly convex optimization problem of the form

x^star = argmin_x f(x, theta)

the computation of the Jacobian dx^star / dtheta analytically requires the inversion of the full Hessian matrix via the implicit function theorem:

[partial x^star / partial theta] = -[partial (df/dx) / partial x]^{-1} [partial (df/dx) / partial theta]

where [partial (df/dx) / partial x] is the Hessian of the optimization problem at the optimum.

If x \in R^n then in the locally strictly convex case the Hessian has full rank n.

Backpropagation through only one gradient step contains only information of rank 1.
So the backpropagation of one gradient step generally only contains a small part of the sensitivity of the optimum with respect to theta.

When 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.
But the approach certainly adds variance compared to having an exact gradient or unrolling many optimization steps.

Additionally, 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.
The stochastic search method proposed by the authors gains information about the loss function shape by sampling locally in the input space.
These types of gradient estimators have much higher variance for truly high dimensional loss functions due to the curse of dimensionality.

See Ben Rechts criticism of the policy gradient / REINFORCE gradient estimator that has the same issue:
https://www.argmin.net/2018/02/20/reinforce/

So 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.

Perhaps 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.

Since 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.
The 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).
For 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.

It 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.
I 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.

The 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.
Reverse 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).
The computation time for truly high-dimensional problems can thus look very different than medium dimensional problems (n = 100).

Regarding the violin plots for the portfolio optimization I am not sure how much they benefit the paper.
The goal of the paper should be to demonstrate the superiority of the suggested method over baselines.
The violin plots are relevant to people who are interested in the outcome of that specific optimal control task.
Someone not knowledgeable about the control problem will not benefit much from knowing the terminal wealth of the different strategies.

More analysis about the quality of the gradients obtained by the suggested method compared to other methods would improve the paper:
- How many unrolling steps do what to the gradient variance?
- 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?

Another small suggestions:
On 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).

review_title: Can be improved by a more compelling analysis and principled reasoning for the approach

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)