On Differentiating Parameterized Argmin and Argmax Problems with Application to Bi-level Optimization

Introduction

Bi-level optimization has a long history of study dating back to the early 1950s and investigation of the so-called Stackelberg model Bard 1998. In this model, two players—a market leader and a market follower—compete for profit determined by the price that the market is willing to pay as a function of total goods produced and each player’s production cost. Recently, bi-level optimization problems have found application in machine learning and computer vision where they have been applied to parameter and hyper-parameter learning Do et al. 2007; Domke 2012; Klatzer and Pock 2015, image denoising Samuel and Tappen 2009; Ochs et al. 2015, and most recently, video activity recognition Fernando and Gould 2016; Fernando et al. 2016.

A bi-level optimization problem consists of an upper problem and a lower problem. The former defines an objective over two sets of variables, say ${\boldsymbol{x}}$ and ${\boldsymbol{y}}$. The latter binds ${\boldsymbol{y}}$ as a function of ${\boldsymbol{x}}$, typically by solving a minimization problem. Formally, we can write the problem as

where $f^{U}$ and $f^{L}$ are the upper- and lower-level objectives, respectively. As can be seen from the structure of the problem, the lower-level (follower) optimizes its objective subject to the value of the upper-level variable ${\boldsymbol{x}}$. The goal of the upper-level (leader) is to choose ${\boldsymbol{x}}$ (according to its own objective) knowing that the lower-level will follow optimally.

In one sense the $\mathop{\textrm{argmin}}$ appearing in the lower-level problem is just a mathematical function and so bi-level optimization can be simply viewed as a special case of constrained optimization. In another sense it is useful to consider the structure of the problem and study bi-level optimization in its own right, especially when the $\mathop{\textrm{argmin}}$ cannot be computed in closed-form. As such, many techniques have been proposed for solving different kinds of bi-level optimization problems Bard 1998; Dempe and Franke 2015. In this technical report we focus on first-order gradient based techniques, which have become very important in machine learning and computer vision with the wide spread adoption of deep neural network models LeCun et al. 2015; Schmidhuber 2015; Krizhevsky et al. 2012.

Our main aim is to collect results on differentiating parameterized $\mathop{\textrm{argmin}}$ and $\mathop{\textrm{argmax}}$ problems. Some of these results have appeared in one form or another in earlier works, e.g., Faugeras 1993 considers the case of unconstrained and equality constrained $\mathop{\textrm{argmin}}$ problems when analysing uncertainty in recovering 3D geometry. However, with the growth in popularity of deep learning we feel it important to revisit the results (and present examples) in the context of first-order gradient procedures for solving bi-level optimization problems.

We begin with a brief overview of methods for solving bi-level optimization problems to motivate our results in Section 2. We then consider unconstrained variants of the lower-level problem, either $\mathop{\textrm{argmin}}$ or $\mathop{\textrm{argmax}}$ (in Section 3), and then extend the results to problems with equality and inequality constraints (in Section 4). We include motivating examples with gradient calculations and discussion throughout the paper leading to a small bi-level optimization example for learning a novel softmax classifier in Section 5. Examples are accompanied by supplementary Python code. Available for download from http://users.cecs.anu.edu.au/sgould/.

Background

The canonical form of a bi-level optimization problem is shown in Equation 1. There are three general approaches to solving such problems that have been proposed over the years. In the first approach an analytic solution is found for the lower-level problem, that is, an explicit function ${\boldsymbol{y}}^{\star}({\boldsymbol{x}})$ that when evaluated returns an element of $\mathop{\textrm{argmin}}_{{\boldsymbol{y}}}f^{L}({\boldsymbol{x}},{\boldsymbol{y}})$. If such a function can be found then we are in luck because we now simply solve the single-level problem

Of course this problem, may itself be difficult to solve. Moreover, it is not always the case that an analytic solution for the lower-level problem will exist.

The second general approach to solving bi-level optimization problems is to replace the lower-level problem with a set of sufficient conditions for optimiality (e.g., the KKT conditions for a convex lower-level problem). If we think of these conditions being encapsulated by the function $h^{L}({\boldsymbol{x}},{\boldsymbol{y}})=0$ then we can solve the following constrained problem instead of the original,

The main difficulty here is that the sufficient conditions may be hard to express and the resulting problem hard to solve. Indeed, even if the lower-level problem is convex, the resulting constrained problem may not be.

The third general approach to solving bi-level optimization problems is via gradient descent on the upper-level objective. The key idea is to compute the gradient of the solution to the lower-level problem with respect to the variables in the upper-level problem and perform updates of the form

In this respect the approach appears similar to the first approach. However, now the function ${\boldsymbol{y}}^{\star}({\boldsymbol{x}})$ does not need to be found explicitly. All that we require is that the lower-level problem be efficiently solveable and that a method exists for finding the gradient at the current solution. Note that we have made no assumption about the uniqueness of $y^{\star}$ nor the convexity of $f^{L}$ or $f^{U}$. When multiple minima of $f^{L}$ exist care needs to be taken during iterative gradient updates to select consistent solutions or when jumping between modes. However, these considerations are application dependent and do not invalidate any of the results included in this report.

This last approach is important in the context large-scale and end-to-end machine learning applications where first-order (stochastic) gradient methods are often the preferred method. This then motivates the results included in this technical report, i.e., computing the gradients of parameterized $\mathop{\textrm{argmin}}$ and $\mathop{\textrm{argmax}}$ optimization problems where the parameters are to be optimized for some external objective or are themselves the output of some other parameterized function to be learned.

Unconstrained Optimization Problems

where $f_{XY}\doteq\frac{\partial^{2}f}{\partial x\partial y}$ and $f_{YY}\doteq\frac{\partial^{2}f}{\partial y^{2}}$.

Equating to zero and rearranging gives the desired result

We now extend the above result to the case of optimizing over vector-valued arguments.

Note that the main computational challenge of the inverting the $n\times n$ matrix $f_{YY}$ (or decomposing it to facilitate solving each system of $n$ linear equations) only needs to be done once and can then be reused for the derivative with respect to each parameter. Thus the overhead of computing gradients for multiple parameters is small compared to the cost of computing the gradient for just one parameter. Of course if ${\boldsymbol{x}}$ (or ${\boldsymbol{g}}({\boldsymbol{x}})$) is changed (e.g., during bi-level optimization) then $f_{YY}$ will be different and its inverse recalculated.

So far we have only considered minimization problems. However, studying the proofs above we see that they do not require that ${\boldsymbol{g}}(x)$ be a local-minimum point; any stationary point will suffice. Thus, the result extends to the case of $\mathop{\textrm{argmax}}$ problems as well.

In this section we consider the simple example of finding the point whose sum-of-squared distance to all points in the set $\{h_{i}(x)\}_{i=1}^{m}$ is minimized. Writing out the problem as a mathematical optimization problem we have $g(x)=\mathop{\textrm{argmin}}_{y}f(x,y)$ where $f(x,y)=\sum_{i=1}^{m}(h_{i}(x)-y)^{2}$. Here a well-known analytic solution exists, namely the mean $g(x)=\frac{1}{m}\sum_{i=1}^{m}h_{i}(x)$.

which agrees with the analytical solution (assuming the derivatives $h^{\prime}_{i}(x)$ exist).

2 Example: Scalar Function with Three Local Minima

The results above do not require that the function $f$ being optimized have a single global optimal point. Indeed, as discussed above, the results hold for any stationary point (local optima or inflection point). In this example we present a function with up to three stationary points and show that we can calculate the gradient with respect to $x$ at each stationary point. Technically the function that we present only has three stationary points when $x<-2\left(\frac{4}{3}\right)^{\frac{1}{3}}$ or $x>0$. It has two stationary points at $x=-2\left(\frac{4}{3}\right)^{\frac{1}{3}}$ and a unique stationary point elsewhere. Consider the function

with the following partial first- and second-order derivatives

Then for any stationary point $g(x)$ of $f(x,y)$, where stationarity is with respect to $y$ for fixed $x$, we have

Here the gradient describes how each stationary point $g(x)$ moves locally with an infintisimal change in $x$. If such a problem, with multiple local minima, were to be used within a bi-level optimization learning problem then care should be taken during each iteration to select corresponding solution points at each iteration of the algorithm.

The (three) stationary points occur when $\frac{\partial f(x,y)}{\partial y}=0$, which (in this example) we can compute analytically as

which we use when generating the plots below.

Figure 1 shows a surface plot of $f(x,y)$ and a slice through the surface at $x=1$. Clearly visible in the plot are the three stationary points. The figure also shows the three values for $g(x)$ and their gradients $g^{\prime}(x)$ for a range of $x$. Note that one of the solutions (i.e., $y=0$) is independent of $x$.

3 Example: Maximum Likelihood of Soft-max Classifier

Now consider the more elaborate example of exploring how the maximum likelihood feature vector of a soft-max classifier changes as a function of the classifier’s parameters. Assume $m$ classes and let classifier be parameterized by $\Theta=\{(\boldsymbol{a}_{i},b_{i})\}_{i=1}^{m}$. Then we can define the likelihood of feature vector ${\boldsymbol{x}}$ for the $i$-th class of a soft-max distribution as

where $Z({\boldsymbol{x}};\Theta)=\sum_{j=1}^{m}\exp\left(\boldsymbol{a}_{j}^{T}{\boldsymbol{x}}+b_{j}\right)$ is the partition function. Note that we use a different notation here to be consistent with the standard notation in machine learning. In particular, the role of ${\boldsymbol{x}}$ is differs from its appearance elsewhere in this article.

The maximum (log-)likelihood feature vector for class $i$ can be found as

whose objective is concave and so has a unique global maximum. However, the problem, in general, has no closed-form solution. Yet we can still compute the derivative of the maximum likelihood feature vector ${\boldsymbol{g}}_{i}(\Theta)$ with respect to any of the model parameters as follows.

where $\boldsymbol{e}_{k}$ is the $k$-th canonical vector ($k$-th element one and the rest zero) and $\makebox[1.29167pt][l]{[}\makebox{[}\cdot\makebox[1.29167pt][l]{]}\makebox{]}$ is the indicator function (or iverson bracket), which takes value 1 when its argument is true and 0 otherwise.

This example will be developed further below by adding equality and inequality constraints and then applying it in the context of bi-level optimization.

4 Invariance Under Monotonic Transformations

These examples motivate the following lemma, which formalizes the fact that composing a function with a monotonically increasing or monotonically decreasing function does not change its stationary points.

where $f_{XY}\doteq\frac{\partial^{2}f}{\partial x\partial y}$ and $f_{YY}\doteq\frac{\partial^{2}f}{\partial y^{2}}$.

Follows from Lemma 3.1 observing that $\frac{\partial h(f(x,y))}{\partial y}=h^{\prime}(f(x,y))\frac{\partial f(x,y)}{\partial y}$ by the chain rule and that $h^{\prime}(f(x,y))$ is always non-zero for monotonically increasing or decreasing $h$. ∎

Constrained Optimization Problems

In this section we extend the results of the $\mathop{\textrm{argmin}}$ and $\mathop{\textrm{argmax}}$ derivatives to problems with linear equality and arbitrary inequality constraints.

Let us introduce linear equality constraints $A{\boldsymbol{y}}=\boldsymbol{b}$ into the vector version of our minimization problem. We now have ${\boldsymbol{g}}(x)=\mathop{\textrm{argmin}}_{{\boldsymbol{y}}:A{\boldsymbol{y}}=\boldsymbol{b}}f(x,{\boldsymbol{y}})$ and wish to find ${\boldsymbol{g}}^{\prime}(x)$.

Alternatively, we can construct the ${\boldsymbol{g}}^{\prime}(x)$ directly as the following lemma shows.

Consider the Lagrangian ${\cal L}$ for the constrained optimization problem,

From the first row of the block matrix equation above, we have

where $H=f_{YY}(x,{\boldsymbol{y}}^{\star})$.

Note that for $A=0$ (and $b=0$) this reduces to the result from Lemma 3.2. Furthermore, ${\boldsymbol{g}}^{\prime}(x)$ is in the null-space of $A$, which we require if the linear equality constraint is to remain satisfied.

2 Inequality Constraints

Consider a family of optimization problems, indexed by $x$, with inequality constraints. In standard form we have

where $f_{0}(x,{\boldsymbol{y}})$ is the objective function and $f_{i}(x,{\boldsymbol{y}})$ are the inequality constraint functions.

where $t>0$ is a scaling factor that controls the approximation (or duality gap when the problem is convex). We now have an unconstrained problem and can invoke Lemma 3.2.

For completeness, recall that the gradient and Hessian of the log-barrier function, $\phi({\boldsymbol{z}})=\sum_{i=1}^{m}\log(-f_{i}({\boldsymbol{z}}))$, are given by

Thus the gradient of an inequality constrained $\mathop{\textrm{argmin}}$ function can be approximated as

In many cases the constraint functions $f_{i}$ will not depend on $x$ and the above expression can be simplified by setting $\phi_{XY}(x,{\boldsymbol{y}})$ to zero.

3 Example: Positivity Constraints

We consider an inequality constrained version of our scalar mean example from above,

where we have added a positivity constraint on $y$. This problem has closed-form solution

Following Section 4.2 we can construct the approximation

where $f(x,y)=\sum_{i=1}^{m}\left(h_{i}(x)-y\right)^{2}$. Applying Lemma 3.1 gives

Note that here we can solve the quadratic equation induced by $\frac{\partial}{\partial y}\left(tf(x,y)-\phi(y)\right)=0$ to obtain a closed-form solution of $g_{t}(x)$ and hence $g^{\prime}_{t}(x)$. We leave this as an exercise.

Observe from Equation 76 that as $t\rightarrow\infty$,

We demonstrate this example on a special case with $m=1$ and $h_{1}(x)=x$. The true and approximate function and their gradients are plotted in Figure 3.

4 Example: Maximum Likelihood of Constrained Soft-max Classifier

Let us now continue our soft-max example from Section 3.3 by adding constraints on the solution. We begin by adding the linear equality constraint $\textbf{1}^{T}{\boldsymbol{x}}=1$ to give

Following Lemma 4.2 and letting $H^{\dagger}=\left(H^{-1}-\frac{1}{\textbf{1}^{T}H^{-1}\textbf{1}}H^{-1}\textbf{1}\textbf{1}^{T}H^{-1}\right)$ we have the following gradients with respect to each parameter,

where $H$ and $\bar{\boldsymbol{a}}$ are as defined in Section 3.3.

Figure 4 shows the constrained solutions before and after taking a gradient step for the same maximum-likelihood surface as described in Section 3.3. Notice that the solution lies along the line $x_{1}+x_{2}=1$. Moreover, the sum of gradients for $x_{1}$ and $x_{2}$ are zero, i.e., the gradient is in the null-space of $\textbf{1}^{T}$. To show this it is sufficient to prove that $H^{\dagger}\textbf{1}=0$, which is straightforward.

Next we consider adding the inequality constraint $\|{\boldsymbol{x}}\|^{2}\leq 1$ to the soft-max maximum-likelihood problem. That is, we constrain the solution to lie within a unit ball centered at the origin. The new function is given by

Following Section 4.2 we can compute approximate gradients as

where $H$ and $\bar{\boldsymbol{a}}$ are as defined in Section 3.3, and the second derivative for the log-barrier $\phi({\boldsymbol{x}})=\log\left(1-\|{\boldsymbol{x}}\|^{2}\right)$ is given by

Similar to previous examples, Figure 5 shows the maximum-likelihood surfaces and corresponding solutions before and after taking a gradient step (on all parameters) in the negative $x_{1}$ direction.

Bi-level Optimization Example

We now build on previous examples to demonstrate the application of the above results to the problem of bi-level optimization. We consider the constrained maximum-likelihood problem from Section 4.4 with the goal of optimizing parameter to achieve a desired location for the maximum-likelihood feature vectors. We consider a three class problem over two dimensional space. Here we use the constrained version of the problem since for a three class soft-max classifier over two-dimensional space there will always be a direction in which the likelihood tends to one as the magnitude of the maximum-likelihood feature vector ${\boldsymbol{x}}$ tends to infinity.

Let the target location for the $i$-th maximum-likelihood feature vector be denoted by ${\boldsymbol{t}}_{i}$ and given. Optimizing the parameters of the classifier to achieve these locations (with the maximum-likelihood feature vectors constrained to the unit ball centered at the origin) can be formalised as

Solving by gradient descent gives updates

for any $\theta\in\Theta=\{(a_{ij},b_{i})\}_{i=1}^{m}$. Here $\eta$ is the step size.

Example likelihood surfaces for the three classes in our example are shown in Figure 6 for both initial parameters and final (optimized) parameters where we have set the target locations to be evenly spaced around the unit circle. Notice that this is achieved by the final parameter settings. Also shown in Figure 6(c) is the learning curve (in log-scale). Here we see a rapid decrease in the objective in the first 20 iterations and final convergence (to within $10^{-9}$ of the optimal value) in under 100 iterations.

Discussion

We have presented results for differentiating parameterized $\mathop{\textrm{argmin}}$ and $\mathop{\textrm{argmax}}$ optimization problems with respect to their parameters. This is useful for solving bi-level optimization problems by gradient descent Bard 1998. The results give exact gradients but (i) require that function being optimized (within the $\mathop{\textrm{argmin}}$ or $\mathop{\textrm{argmax}}$) be smooth and (ii) involve computing a Hessian matrix inverse, which could be expensive for large-scale problems. However, in practice the methods can be applied even on non-smooth functions by approximating the function or perturbing the current solution to a nearby differentiable point. Moreover, for large-scale problems the Hessian matrix can be approximated by a diagonal matrix and still give a descent direction as was recently shown in the context of convolutional neural network (CNN) parameter learning for video recognition via stochastic gradient descent Fernando and Gould 2016.

The problem of solving non-smooth large-scale bi-level optimization problems, such as in CNN parameter learning for video recognition, present some interesting directions for future research. First, given that the parameters are likely to be changing slowly for any first-order gradient update it would be worth investigating whether warm-start techniques would be effectivey for speeding up gradient calculations. Second, since large-scale problems often employ stochastic gradient procedures, it may only be necessary to find a descent direction rather than the direction of steepest descent. Such an approach may be more computationally efficient, however it is currently unclear how such a direction could be found (without first computing the true gradient). Last, the results reported herein are based on the optimal solution to the lower-level problem. It would be interesting to explore whether non-exact solutions could still lead to descent directions, which would greatly improve efficiency for large-scale problems, especially during the early iterations where the parameters are likely to be far from their optimal values.

Models that can be trained end-to-end using gradient-based techniques have rapidly become the leading choice for applications in computer vision, natural language understanding, and other areas of artificial intelligence. We hope that the collection of results and examples included in this technical report will help to develop more expressive models—specifically, ones that include optimization sub-problems—that can still be trained in an end-to-end fashion. And that these models will lead to even greater advances in AI applications into the future.

References