We present learn2convexify, a data-driven method to solve similar non-convex optimization problems in real-time. Non-convex optimization scales poorly with the problem size, and thus solution algorithms remain prohibitively expensive. Most of these algorithms are general-purpose and tackle each optimization instance individually. In practice, however, we often solve very similar problems with varying data. We propose a novel neural network architecture to solve non-convex problems that are similar in nature, , come from the same family. Our method leverages common structure across the family to learn a shared latent variable representation, which is then used as input to a convex optimization problem whose solution constitutes an approximate solution to the original non-convex optimization problem. To do this, we marry the strength of neural networks in uncovering structure from data with recent advances in differentiable convex optimization which allows end-to-end training of the latent variables and approximate optimization problems to minimize suboptimality. After training, our method consists of a feedforward network followed by a convex optimization problem which enables fast real-time decision making. Additionally, we provide generalization bounds that guarantee good out-of-sample performance.

§ INTRODUCTION

Consider the following parametric optimization problem.

\begin{equation}\label{eq:original-noncvx}
\begin{array}{ll}
\text{minimize} & f(x, \theta)\\
\text{subject to} & x \in \mathcal{X}(\theta)
\end{array}
\end{equation}

$x$ is the decision variable and $\theta$ is the parameter. $\theta$ is a random sample drawn from a distribution $\mathcal{D}$ with a bounded support set $\Theta$. The function $f$ is convex and $\alpha$-smooth. The constraint set is the intersection of two other sets, $\mathcal{X}(\theta) = \mathcal{N}(\theta) \cup \mathcal{C}(\theta)$. The set $\mathcal{N}(\theta)$ is potentially non-convex, but prox-regular at local minima. That is, it has a single valued Euclidean projection around local minima  <cit.>. Sparse and low-rank constraints are popular examples of prox-regular constraints. The set $\mathcal{C}$ is convex and compact.

Recent advances and limitations.
In recent years however, it has become increasingly clear that the computation time and memory needs to solve non-convex problems explode as the size of the problem grows.

Our contributions.
In this paper, we introduce Learn2Convexify, a novel data-driven method that can quickly solve non-convex problems that come from the same family. We create a neural network architecture, represented in Figure cite figure, which consists of two modules. The first module predicts an initial point using a neural network learned across the family of non-convex problems. These predictions are in turn used as an input to initialize several iterations of the inner nexos algorithm (<ref>), the second module.

We can treat each nexos iteration as a layer and backpropagate through the network to update all of its parameters in search of a better initial point as predicted by the first module.

* We effectively reduce the non-convex problem to a locally convex one.

* We then further increase the speed of our forward pass by only running several iterations of (<ref>) to compute the final solution. This is faster than solving the locally convex problem to convergence because we only require a few iterations of (<ref>).

* We use automatic differentiation to compute the derivative of the solution of the inner nexos module with respect to its initial point. This allows us to conduct end-to-end training and link the prediction of a good starting point to the nexos layers.

* We unroll a fixed number of iterates in training, making our forward pass time significantly faster than even solving a convex optimization problem. This allows us to reliably solve problems on the scale of microseconds.

* We use a smart initialization step that helps to get better global solutions.

* We provide an analysis of the Rademacher complexity of our method.

* We benchmark our approach on examples with real-time requirements, achieving better scalability than commercial non-convex solvers such as Gurobi.

§ RELATED WORK

Differentiable optimization.
Differentiable optimization is a method to compute gradients of optimization problem solutions with respect to their input parameters, by differentiating their optimality conditions <cit.>. It has been applied to several different areas such as control policies design <cit.>, meta learning <cit.>, and game theory <cit.>. While differentiable optimization has been used to perform end-to-end prediction and optimization, its focus has not been to speed up optimization algorithms.

Optimization as a classification problem.
In real world applications, practitioners who use optimization algorithms are often unable to understand how the optimal solution is arrived at. Thus, optimization algorithms are perceived as black-boxes. Bertsimas and Stellato <cit.> redefine optimization problems as multi-class classification problems where parameters are mapped to strategies. These strategies are used to solve much simpler optimization problems leading to quicker run-times. Additionally, using classification trees allows for practitioners to interpret the solution of the optimization problem and to efficiently solve them. However, this method relies on the concept of active linear constraints and therefore does not generalize to SDPs.

Learning heuristics for combinatorial optimization.
Since combinatorial optimization problems scale poorly, many methods to solve them rely on heuristics. For example, branch and bound methods rely on several levels of heuristics about the branching decisions that are made.
Khalil et al. <cit.> propose to learn these branching decisions from an expensive strategy that creates search trees. However, branch and bound methods can still be slow even if the heuristics are optimized. The work <cit.> also describes a method for learning heuristics, but specializes to graph problems while using a reinforcement learning framework.

Our approach differs from this strategy method.  Additionally, we focus on a family of parametric optimization problems, while these methods mostly focus on solving a single problem. Lastly, these works do not use differentiable optimization to embed optimization blocks in the architecture and cannot exploit their optimality conditions.

Unrolling iterates

§ NEXOS FRAMEWORK

Here, we introduce the Non-convex exterior optimization (NExOS) algorithm <cit.>. NExOS is used to solve problems of the form,

\begin{array}{ll} \mbox{minimize} & f(x) \\
\mbox{subject to} & x \in \mathcal{X}, \\
\end{array}

where $f(x)$ is a convex function and $\mathcal{X}$ is prox-regular around local minima.

NExOS proposed to solve this problem by solving a series of smoothened versions of the problem, each of which is parametrized by a smoothening parameter $\mu$. First, we introduce our notation for the Moreau envelope. The Moreau envelope for a function, $h(x)$, is

$${}^{\mu} h(x) = \min_z h(z) + \frac{1}{2 \mu} \|x-z\|_2^2

The Moreau envelope is directly related to the prox of a function. The prox is the minimizer of the above optimization problem, i.e.

\textbf{prox}_{\gamma h}(x) = \argmin_z h(z) + \frac{1}{2 \mu} \|x-z\|_2^2

Now, we smoothen the constraint.

\begin{array}{ll} \mbox{minimize} & f(x) + {}^{\mu} \mathcal{I}(x) \\
\end{array}

The NExOS algorithm works by solving a series of problems $\mathcal{P}_{\mu_1}, \dots, \mathcal{P}_{\mu_k}$ with decrease $\mu$ values to local minima. As $\mu$ approaches $0$, we get closer to approximating the original non-smooth optimization problem. Each problem is initialized with the solution of the previous problem.

The nexos inner algorithm

$n = 0$
$\psi = \frac{\mu}{\gamma + \mu}$
$\|x^n-y^n\| \leq \epsilon$
$x^{n+1} = \textbf{prox}_{\gamma f}(z^n)$
$\tilde{y}^{n+1} = 2x^{n+1} - z^n$
$y^{n+1} = \psi \tilde{y}^{n+1} + (1 - \psi) \mathbf{\Pi}(\tilde{y}^{n+1})$
$z^{n+1} = z^n + y^{n+1} - x^{n+1}$
$n = n + 1$

Algorithm (<ref>) can be rewritten in a single line as

$$z^{n+1} = \mathcal{T}_{\mu} (z^n)$$

where $\mathcal{T}_{\mu}$ is an operator defined as

$$\mathcal{T}_{\mu} = \textbf{prox}_{\gamma {}^\mu \mathcal{I}} (2 \textbf{prox}_{\gamma f} - I) + I - \textbf{prox}_{\gamma f}$$

Let $\bar{x}$ be a local min to $\mathcal{P}$. Following Proposition 2 of <cit.>, for any small enough $\mu$, $\mathcal{T}_{\mu}$ is a firmly nonexpansive operator over a ball centered at $\bar{x}$. Furthermore, if $f$ is strongly convex, then $\mathcal{T}_{\mu}$ is contractive with factor $\kappa \in (0, 1)$.

§ OUR METHOD

Neural Network Architecture.
The first module of our neural network consists of $L$ relu layers that predicts a warm-start, $\hat{z}$. That is, $\hat{z} = s(\theta) = h_L(h_{L-1} \dots h_1(\theta))$ where $h_l(y_l) = \sigma(W_l y_l + b_l)$. Here, $\sigma(y) = \max(y, 0)$. We write this prediction compactly as a mapping from the input space in $\mathbf{R}^d$ to $\mathbf{R}^n$. We let $\mathcal{H}$ denote the set of such mappings.

The second module is an algorithm that uses thee initial dual point and outputs a final dual point. Our algorithm consists of $k$ ADMM iteratioions.

Let $\mathcal{A}$ be the algorithm that maps $\mathbf{R}^n$ to $\mathbf{R}^n$. So $\mathcal{A}(\hat{z}) = z^{k}$ where $z^{i+1} = T_{\mu}(z^i)$ and $z^1 = T_{\mu}(\hat{z})$.

Thus, our neural network architecture is a mapping that composes the two modules mentioned above. Let $\mathcal{G} = \{h \circ \mathcal{A} \mid h \in \mathcal{H} \}$. Then our neural network is $g \in \mathcal{G}$.

Learning task.
There are two loss functions that we can consider. One is a function that combines suboptimality and innfeasibility.

The loss of a single problem parametrized by $\theta$ is

$$L(\hat{z}, \theta) = f(\textbf{prox}_{\gamma f}(\hat{z}), \theta) + \frac{1}{2 \mu} d^2(\textbf{prox}_{\gamma f}(\hat{z}), \theta) - f(x^\star, \theta).$$

Here, $x^\star$ is the solution to the non-convex problem parametrized by $\theta$. The value of $\mu$ is chosen to be a small so that we closely approximate the indicator function. Having a nonzero $\mu$ will allow us to compute derivatives.

The parameters that we optimize over are the weights of the neural network in module one. The goal is to minimize the expected loss, i.e. the risk,

$$R_{L2C}(g) = \mathbf{E}_{\theta \sim \mathcal{D}}[L(g(\theta), \theta)].$$

Since we don't know the underlying distribution, $\mathcal{D}$, we instead attempt to minimize the empirical risk,

$$\hat{R}_{L2C}(g) = \frac{1}{N} \sum_{i=1}^N L(g(\theta_i), \theta_i)$$

Thus we want to solve the optimization problem

\begin{array}{ll} \mbox{minimize}_{h \in \mathcal{H}} & \hat{R}_{L2C}(h \circ \mathcal{A}) \\
\end{array}

Each of the functions involved the composition of our neural network is differentiable almost everywhere. Thus, we can employ a gradient-based method to minimize the training loss.

Online optimization.
For a new test point, $\theta$, we simply need to do a forward pass of the neural network, i.e. compute $g(\theta)$.

§ GENERALIZATION BOUNDS

§.§ Convex case

If the set $\mathcal{X}$ is a convex set, then our problem is convex. Thus, if we pick $k$ to be infinity, then our problem will always converge to the unique, global optimal solution regardless of the weights of the neural network. Therefore, the Rademacher complexity will be zero. Now, we discuss what happens for finite $k$.

The empirical Rademacher complexity of a function class $\mathcal{G}$ with respect to a sample $S = (\theta_1, \dots, \theta_N)$ is defined as

$$\hat{\mathcal{R}}_S(\mathcal{G}) = \mathbf{E}_{\sigma}[\sup_{g \in \mathcal{G}} \frac{1}{N} \sum_{i=1}^N \sigma_i g(\theta_i)]$$

The Rademacher complexity is

$$\mathcal{R}^N(\mathcal{G}) = \mathbf{E}_{S \sim \mathcal{D}^N}[\hat{\mathcal{R}}_S(\mathcal{G})]$$

Our loss function given after all of the ADMM iterations for a problem parametrized by $\theta$ is

$$L^{\mu}(x, \theta) = f(x, \theta) + \frac{1}{2 \mu} d^2(x, \theta) - f(x^\star, \theta)$$

where $x^\star$ is the optimal solution of the parametric problem

If $\mu=0$, then we use the convention of the indicator function to write the loss as

$$L^0(x, \theta) = f(x, \theta) + \mathcal{I}(x, \theta) - f(x^\star, \theta)$$

The goal is to minimize the loss over the distribution of $\theta$ values, i.e. the risk. The risk is

$$R^{\mu}(g) = \mathbf{E}_{\theta \sim D}[L^{\mu}(g(\theta_i), \theta_i)].$$

The empirical risk over a set of samples is

$$\hat{R}^{\mu}_N(g) = \frac{1}{N} \sum_{i=1}^N L^{\mu}(g(\theta_i), \theta_i).$$

It is better to project the point, $x^k$ onto $\mathcal{X}$ so that we get a feasible point. Thus we can remove the Lipschitz constant by taking $y^k = \Pi(2x^k - z^k)$ as our final point.

Assume that $\|z^0-z^\star\|_2 \leq B$, $f$ is strongly convex, and that $\mathcal{X}$ is a convex set. Then $L^0(\mathcal{A}^k(z^0))$ has a Lipschitz constant of $\kappa^k (\|\nabla f(x^\star)\|_2 + \alpha \kappa^k B)$.

First, let $y^k = \mathcal{A}^k(z^0)$ and $v^k = \mathcal{A}^k(w^0)$ where $\|z^0 - z^\star\|_2 \leq B$ and $\|w^0 - z^\star\|_2 \leq B$. Since $d^2(y^k, \theta)=0$ and $d^2(v^k, \theta)=0$, we have $L^0(y^k, \theta) = f(y^k, \theta) - f(x^\star, \theta)$ and $L^0(v^k, \theta) = f(v^k, \theta) - f(x^\star, \theta)$.

From the mean value theorem we get the inequality

$$|L^0(y^k,\theta) - L^0(v^k,\theta)| = |f(y^k,\theta) - f(v^k,\theta)| \leq \sup_{u \in B(x^\star, r)} \|\nabla f(u, \theta)\|_2 \|y^k - v^k\|_2.$$

Using the triangle inequality,

$$\|\nabla f(u, \theta)\|_2 \leq \|\nabla f(u, \theta) - \nabla f(x^\star, \theta)\|_2 + \|\nabla f(x^\star, \theta)\|_2.$$

Then, using the smoothness of $f$,

$$\|\nabla f(u, \theta)\|_2 \leq \alpha \|u - x^\star\|_2 + \|\nabla f(x^\star, \theta)\|_2.$$

First, we prove that the mapping of $z^0$ to $z^k$ which we denote by $\mathcal{T}_{\mu}^k$ is contractive. Since $\mathcal{T}_{\mu}^k$ is the composition of $k$ operators, all of which are contractive with factor $\kappa$, we get that $\mathcal{T}_{\mu}^k$ is Lipschitz with factor $\kappa^k$. The operator $C(z^K) = (2 \textbf{prox}_{\gamma f} - I)(z^k)$ is nonexpansive. Thus $\mathcal{A}^k$ is $\kappa^k$ Lipschitz.

Plugging in and using the initial bound on $\|z^0-z^\star\|_2$

$$\|\nabla f(y^k, \theta)\|_2 \leq \alpha \kappa^k B + \|\nabla f(x^\star, \theta)\|_2.$$

Therefore, we have

$$|L^0(y^k,\theta) - L^0(v^k,\theta)| \leq (\alpha \kappa^k B + \|\nabla f(x^\star, \theta)\|_2) \|y^k-v^k \|_2.$$

Since, $\mathcal{T}_{\mu}^k$ is contractive, we get

$$|L^0(\mathcal{A}^k(z^0,\theta)) - L^0(\mathcal{A}^k(w^0,\theta))| \leq \kappa^k B (\alpha \kappa^k B + \|\nabla f(x^\star, \theta)\|_2) \|z^0-w^0 \|_2.$$

Let $\mathcal{H}$ be the set of mappings from $\mathbf{R}^d$ to $\mathbf{R}^n$ with bounded norm to the optimal solution. That is, $\mathcal{H} = \{ h \mid \max_{\theta \in \Theta} \| h(\theta) - z^\star(\theta) \|_2 \leq B \}$. Then we achieve the following bounds using the Rademacher complexity and the empirical Rademacher complexity. With probability of at least $1 - \delta$ over the draw of i.i.d samples the following holds.

R_{L2C}(g) \leq \hat{R}_{L2C}(g) + 2 c_0 \mathcal{R}^N(\mathcal{H}) + c_1 \frac{\log(2/\delta)}{2N}

R_{L2C}(g) \leq \hat{R}_{L2C}(g) + 2 c_0 \mathcal{R}^S(\mathcal{H}) + 3 c_1 \frac{\log(1/\delta)}{2N}

$$c_0 =  \kappa^k (\alpha \kappa^k B + \max_{\theta \in \Theta} \|\nabla f(x^\star, \theta)\|_2)$$

$$c_1 = B c_0$$

With probability at least $1 - \delta$ over the draw of i.i.d samples the following holds for all $g \in \mathcal{G}$ <cit.>. Here, $c_1$ is the bound on the worst-case possible loss over $g \in \mathcal{G}$.

R^0(g) \leq \hat{R}_{L2C}(g) + 2 \mathcal{R}_N^0(\mathcal{G}) + c_1 \frac{\log(2/\delta)}{2N}

We can also provide a bound in terms of the empirical Rademacher complexity

R^0(g) \leq \hat{R}_{L2C}(g) + 2 \mathcal{R}_S^N(\mathcal{G}) + 3 c_1 \frac{\log(1/\delta)}{2N}

Our result is to use the Lipschitz nature of our forward pass and bound the Rademacher complexity $\mathcal{R}^N_{L2C}(\mathcal{H})$ and worst-case loss, $c_1$.

First, we use the fact that our second module composed with the loss function, (i.e. $L \circ \mathcal{A}^k$), is Lipschitz conditioned on the input into the second module, $z^0$ being within distance $B$ to $z^\star$. Thus, we know that

$$\mathcal{R}^N_{L2C}(\mathcal{G}) \leq \kappa^k B (\alpha \kappa^k B + \|\nabla f(x^\star, \theta)\|_2)  \mathcal{R}^N_{L2C}(\mathcal{H})$$

since $\kappa^k B (\alpha \kappa^k B + \|\nabla f(x^\star, \theta)\|_2) $ is the Lipschitz constant.

The second term is, $c_1$, the worst-case possible suboptimality.

We use the fact that $f$ is Lipschitz under the condition that $z^0$ is a bounded distance away from $z^\star$ and that we get $\|y^k-x^\star \|_2$ is bounded in terms of $\kappa$ and $k$.

Therefore, $|f(x) - f(x^\star)| \leq (\alpha \kappa^k B + \|\nabla \bar{f}(x^\star) \|_2) \kappa^k B = c_1$

Let's consider the case where $\mathcal{H}$ is the set of linear functions with bounded norm. Specifically, $\mathcal{H} = \{h \mid h(\theta) = W \theta\}$ where $\frac{1}{2}\|W\|_F^2 \leq w$. Then we have that $\mathcal{R}^N (\mathcal{H}) \leq \rho_2(\mathcal{\theta}) w \sqrt{\frac{2d}{N}}$. Here, $\rho_2(\mathcal{\theta}) = \sup_{\theta \in \Theta} \|\theta\|_2$

If we assume that $\max_{\theta \in \Theta} \| z^\star(\theta) \|_2 \leq r$, then $B \leq r + \sqrt{2 w} \rho_2(\mathcal{\theta})$

§.§ Equality Constraints

In order to make real-time decisions, we need our $\textbf{prox}_{\gamma f}$ and $\Pi$ operations to admit a closed-form solution. Unfortunately, this is not always the case. For instance, if $\mathcal{X}$ holds equality and inequality constraints. When this happens, we move the equality constraints to the $f$ part, and we take $x^k$ rather than $y^k$. Thus, the equality constraints are satisfied, but the inequality constraints are not. Now, we need to bound the level of inequality constraint violation, $d(x^k) = \|\Pi(x^k) - x^k\|$ as well as the suboptimality, $f(x^k)-f(x^\star)$.

Assume that $\|z^0-z^\star\|_2 \leq B$, $f(x) = \bar{f}(x) + \mathcal{I}_{ \{A x = b\} }$, $\bar{f}(x)$ is strongly convex, and that $\mathcal{X}$ is a convex set. Then $f(\mathcal{A}^k(z^0)) - f(x^\star)$ has a Lipschitz constant of $\kappa^k (\|\nabla f(x^\star)\|_2 + \alpha \kappa^k B)$. Furthermore, $d(\mathcal{A}^k(z^0))$ has a Lipschitz constant of $2 \kappa^k$.

For the part about suboptimality, we reuse the proof from the part without equality constraints in $f$.

For measuring infeasibility, we need a Lipschitz constant on $d(x^k) = \| \Pi(x^k) - x^k \|_2$ where $x^k=\mathcal{A}^k(z^0)$.

Since $T_{\mu}$ is contractive, we know that $\| \mathcal{A}^k(z^0) - \mathcal{A}^k(w^0) \|_2 \leq \kappa^k \|z^0 - w^0\|_2$. It remains to show that $d(x) = \| \Pi(x) - x\|_2$ is Lipschitz with factor $2$.

\begin{align*}
|d(x) - d(u)| &= \| \Pi(x) - x\|_2 - \|\Pi(u)-u\|_2| \\
&\leq \|\Pi(x)-\Pi(u) + u -x\|_2 \\
&\leq \|\Pi(x)-\Pi(u)\|_2 + \| u -x\|_2 \\
&\leq 2 \|u-x\|_2
\end{align*}

The first two inequalities both use the triangle inequality. The last inequality uses the non-expansiveness of $\Pi$.

§.§ Tradeoff between generalization and time

Suppose we are given a budget of time $T$ that we cannot exceed for our online forward pass of our neural networks. We measure the generalization as the ratio of the generalization bound after $k$ ADMM iterations and the generalization bound after $0$ iterations, i.e. $\text{gen} = \frac{\text{bound}(k)}{\text{bound}(0)}$

Now, we can plot the Pareto frontier, $\lambda \frac{\text{bound}(k)}{\text{bound}(0)} + \frac{\text{time}}{T}$.

§.§ Non-strongly convex case - New

If $f$ is not strongly convex, then the operator $T$ is not contractive. Instead, it is an averaged operator. An operator $F$ is $\nu$-averaged if $F = (1-\nu) I + \nu G$ where $G$ is nonexpansive for $\nu \in (0,1)$. Suppose that an operator, $F$ is averaged. Then the fixed point iteration converges to an optimal solution. The rate of convergence is $O(\frac{1}{k})$ for the updates on the size of the iterates,
$\| z^{k+1}-z^k \|$ which is typically the stopping criterion.  Thus, we change our loss to be $L(z) = \|T(z)-z\|_2$.

Let Z be a closed, compact set that contains all the fixed points of a $\nu$-averaged operator $T$. Let $\bar{z}$ be a maximizer to the problem $\max_{z \in Z} L(\mathcal{A}(z))$. Denote $\bar{z}^{\infty}$ as the point that the sequence of iterates, $\{z^0, z^1, \dots \}$ converges to. Then for any $z, w \in Z$ we get

|L(\mathcal{A}(z)) - L(\mathcal{A}(w))| \leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \|\bar{z} - \bar{z}^{\infty}\|_2

Consider the optimization problem

\max_{z,w \in Z}|L(\mathcal{A}(z)) - L(\mathcal{A}(w))|.

Since the loss is non-negative, a maximizer can be achieved by picking $w \in Z$ that minimizes $L(\mathcal{A}(w))$ and picking $z \in Z$ that maximizes $L(\mathcal{A}(z))$. Since $Z$ includes all of the fixed points of $T$ by assumption, we can pick any fixed point as our optimal point for $w$. In particular, we can pick $\bar{z}^{\infty}$ for $w$. Thus, we know $(\bar{z}, \bar{z}^{\infty})$ is a maximizer.

From <cit.>, we know that

\|z^{k+1}-z^k\| \leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \|z^0 - z^\star\|_2.

Now consider any $z,w \in Z$.

\begin{align*}
|L(\mathcal{A}(z)) - L(\mathcal{A}(w))| &\leq
|L(\mathcal{A}(\bar{z})) - L(\mathcal{A}(\bar{z}^{\infty}))| \\
&= L(\mathcal{A}(\bar{z}))\\
&= \|T(\mathcal{A}(\bar{z}))-\mathcal{A}(\bar{z})\| \\
&\leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \|\bar{z} - \bar{z}^{\infty}\|_2
\end{align*}

Since $L(\mathcal{A}(\bar{z}^{\infty})) = 0$,

First, we claim that $(\bar{z}, \bar{z}^{\infty})$ is a maximizer to the problem, $\max_{z,w \in Z} |L(\mathcal{A}(z)) - L(\mathcal{A}(w))|$.

\begin{align*}
\max_{z,w \in Z}|L(\mathcal{A}(z)) - L(\mathcal{A}(w))| &\leq \max_{z \in Z} L(\mathcal{A}(z))
\end{align*}

The first inequality comes from the non-negativity of the loss.

§.§ Non-strongly convex case

If $f$ is not strongly convex, then the operator $T_{\mu}$ is not contractive. Instead, it is an averaged operator. An operator $F$ is $\nu$-averaged if $F = (1-\nu) I + \nu G$ where $G$ is nonexpansive for $\nu \in (0,1)$. Suppose that an operator, $F$ is averaged. Then the fixed point iteration converges to an optimal solution. The rate of convergence is $O(\frac{1}{k})$ for the updates on the size of the iterates,
$\| z^{k+1}-z^k \|$ which is typically the stopping criterion.  Thus, we change our loss to be $L^k(z^0) = \|z^{k+1}-z^k\|_2$ where $z^{i+1} = T_{\mu}(z^i)$.

Consider the loss function, $L^k(z^0) = \|z^{k+1}-z^k\|_2$, where $T_{\mu}$ is an $\nu$-averaged operator. Let $Z$ be a closed and bounded set that contains all the fixed points of $T_{\mu}$. Denote $(Z^0, W^0)$ as the set of optimal solutions to $\max_{z^0 \in Z, w^0 \in Z} |L^k(z^0) - L^k(w^0)|$. By the Weierstrass theorem, this set is non-empty. There exists $(\bar{z}^0, \bar{w}^0) \in (Z^0, W^0)$ such that

$$|L^k(\bar{z}^0) - L^k(\bar{w}^0)| \leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \|\bar{z}^0 - \bar{w}^0\|_2.$$

Let $z^\star$ denote the point that the sequence of $\{z^i\}$ converges to starting from $z^0$. Let $w^\star$ be the same for the sequence that starts from $w^0$. We use the following inequality as a starting point for averaged operators <cit.>.

$$\|z^{k+1}-z^k\|_2 \leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \|z^0-z^\star\|_2$$

We know that

$$|L^k(z^0) - L^k(w^0)| \leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \max(\|z^0-z^\star\|_2, \|w^0-w^\star\|_2).$$

We want to provide a bound over the maximum of $|L^k(z^0) - L^k(w^0)|$ over $z^0$ and $w^0$. We denote $\bar{z}^0$ and $\bar{w}^0$ as the maximizers. Note that if the input to $L^k(w^0)$ is a fixed point then $L^k(w^0) = 0$. Either $\bar{w}^0$ or $\bar{z}^0$ will be a fixed point of $T_{\mu}$ since both $L^k(z^0) \geq 0$ and $L^k(w^0) \geq 0$. Without loss of generality, let $\bar{w}^0$ be a fixed point of $T_{\mu}$. We are free to pick any fixed point since they can all be maximizers. So we pick $\bar{w}^0 = z^\star$. Then $z^\star = w^\star$ and $w^0 = w^\star$. Finally, we get

$$ |L^k(\bar{z}^0) - L^k(\bar{w}^0)| \leq \frac{1}{\sqrt{(k+1) \nu (1 - \nu)}} \|\bar{z}^0-\bar{w}^0\|_2.$$

The result allows us to bound the Rademacher complexity in the averaged operator case. Note that the result is not the same as the Lipschitz result. Instead, it holds over the worst case $z^0$ and $w^0$. However, this suffices to prove a bound on the Rademacher complexity since we only need to consider the worst-case result.

Let $T_{\mu}$ be a $\nu$-averaged operator. Let the loss be $L^k(z^0) = \|z^{k+1}-z^k\|_2$ where $z^{i+1} = T_{\mu}(z^i)$. Let $\mathcal{H}$ be the set of mappings from $\mathbf{R}^d$ to $\mathbf{R}^n$ with bounded norm to the optimal solution. That is, $\mathcal{H} = \{ h \mid \max_{\theta \in \Theta} \| h(\theta) - z^\star(\theta) \|_2 \leq B \}$. Then we achieve the following bounds using the Rademacher complexity and the empirical Rademacher complexity. With probability of at least $1 - \delta$ over the draw of i.i.d samples the following holds.

R_{L2C}(g) \leq \hat{R}_{L2C}(g) + 2 c_0 \mathcal{R}^N(\mathcal{H}) + c_1 \frac{\log(2/\delta)}{2N}

R_{L2C}(g) \leq \hat{R}_{L2C}(g) + 2 c_0 \mathcal{R}^S(\mathcal{H}) + 3 c_1 \frac{\log(1/\delta)}{2N}

$$c_0 =   \frac{1}{\sqrt{(k+1) \nu (1 + \nu)}} $$

$$c_1 = \frac{1}{\sqrt{(k+1) \nu (1 + \nu)}} B$$

The proof is similar to the strongly convex case. We bound the worst-case loss with $c_1$. We use the weaker Lipschitz-like condition to write the Rademacher complexity of $\mathcal{G}$ in terms of $\mathcal{H}$. The proof follows exactly that of the proof for Lemma 7 in <cit.>. In line six, we only need the Lipschitz-like inequality to hold for any maximizer, $(s_1^\star, s_2^\star)$. We redo the key lines of the proof. $\epsilon$ is a Rademacher random variable. We assume that there exists $(\bar{s_1}, \bar{s_2})$ that maximizes $\psi(s_1) +f(s_1) - \psi(s_2) + f(s_2)$ and also the Lipschitz-like condition, $\psi(\bar{s}_1) - \psi(\bar{s}_2) + f(\bar{s}_1 + f(\bar{s}_2) \leq \| \phi(\bar{s}_1) - \phi(\bar{s}_2) \| + f(\bar{s}_1) + f(\bar{s}_2)$, holds.

\begin{align*}
2 \mathbf{E}[\sup_{s \in \mathcal{S}} \epsilon \psi (s) +f(s) - \delta] &= \sup_{s_1,s_2 \in \mathcal{S}} \psi(s_1) +f(s_1) - \psi(s_2) + f(s_2) -\delta \\
&\leq \psi(s_1^\star) - \psi(s_2^\star) + f(s_1^\star) + f(s_2^\star) \\
&= \psi(\bar{s}_1) - \psi(\bar{s}_2) + f(\bar{s}_1 + f(\bar{s}_2) \\
&\leq \| \phi(\bar{s}_1) - \phi(\bar{s}_2) \| + f(\bar{s}_1) + f(\bar{s}_2) \\
\end{align*}

To get from the first line to the second, we use a maximizer, $(s_1^\star, s_2^\star)$. To get the third line with equality, we use the fact that any maximizer works and therefore we are free to select that one that fits our needs, which we call $(\bar{s}_1, \bar{s}_2)$. To get the last line, we use the Lipschitz-like condition which may not hold for every maximizer, but will hold for $(\bar{s}_1, \bar{s}_2)$. The rest of the proof follows as is.

To apply to our work, we use the maximizer, $(\bar{s}_1, \bar{s}_2) = (\bar{z}^0, z^\star)$. The set $\mathcal{S}$ is any closed and bounded set that contains all of the fixed points of $T_{\mu}$. Remember, we found that $\bar{z}^0$ is a point that maximizes $L^k(z^0)$ and we define $z^\star$ is the limit point of the iterates generated starting at $\bar{z}^0$. For us, $\phi$ is the identity function and $\psi(z^0) = L^k(z^0)$.

Let's consider the case where $\mathcal{H}$ is the set of linear functions with bounded norm. We have that $\mathcal{R}^N (\mathcal{H}) \leq \rho_2(\mathcal{\theta}) w \sqrt{\frac{2d}{N}}$. Here, $\rho_2(\mathcal{\theta}) = \sup_{\theta \in \Theta} \|\theta\|_2$. If we assume that $\max_{\theta \in \Theta} \| z^\star(\theta) \|_2 \leq r$, then $B \leq r + \sqrt{2 w} \rho_2(\mathcal{\theta})$. We get the benefit of having the generalization bound with rates $O(\frac{1}{\sqrt{N}})$ and $O(\frac{1}{\sqrt{k}})$.

§.§ Non-convex, Discrete Case

Here, we consider the case where $\mathcal{X}$ limits the decision variables to discrete variables only. $\mathcal{X} = \{x \mid x \in \{0,1\}^n\}$.

In this case, we only need to run one iteration of ADMM.

$$\mathcal{A}^{\infty}(z^0) = \Pi(2 \textbf{prox}_{\gamma f}(z^0) - z^0)$$

We assume that the solutions to all of the problems are unique. We can consider this as a classification problem with the loss as

$$L^{0-1}(h, \theta) = \mathbf{1}_{ \{\mathcal{A}^\infty(h(\theta)) = x^\star(\theta) \}}.$$

The empirical risk, which gives the proportion of predictions that are in the convex ball, is

$$\hat{R}^{0-1}(h) = \frac{1}{N} \mathbf{1}_{ \{\mathcal{A}^\infty(h(\theta_i)) = x^\star(\theta_i) \}}.$$

The risk is

$$R^{0-1}(h) = \mathbf{E}_{\theta \sim \mathcal{D}}[L^{0-1}(h, \theta)].$$

Here, $\mathcal{S}$ is the set of all possible local minima over all of the problems. For many problems, this is combinatorial in nature.

$$\mathbf{E}[R^{0-1}(h)] \leq \hat{R}^{0-1}(h) + 2 \sqrt{\frac{d_N(\mathcal{A^{\infty}(\mathcal{H})}) \log(N |\mathcal{S}|^2)}{N}} + \frac{\log(1/\delta)}{2 N}$$

Here, $d_N$ is the Natarajan dimension of a hypothesis class.

The need to bound two things: the empirical Rademacher complexity and the worst-case loss. The worst-case loss for a single problem is just one. Thus we only need to bound the Rademacher complexity.

$$\mathcal{R}^N(\mathcal{G}) = \mathbf{E}_{\sigma} [\sup_{h \in \mathcal{H}} \frac{1}{N} \sum_{i=1}^N \sigma_i \mathbf{1}_{ \{ \mathcal{A}^{\infty}(h(\theta_i)) = x^\star_i \} }]$$

Let $\mathcal{F} = \{A^\infty(h(\theta_1)), \dots, A^\infty(h(\theta_N)) \mid h \in \mathcal{H}\}$

$$\mathcal{R}^N(\mathcal{G}) \leq \sqrt{\frac{2 \log |\mathcal{F}| }{N}}$$

Massart's lemma states the following. Let $\mathcal{A} \in \mathbf{R}^m$ be a finite set with $r = \max_{x \in \mathcal{A}} \|x\|_2$ then the following holds <cit.>.

$$\mathbf{E}_{\sigma}[\frac{1}{m} \sum_{i=1}^m \sigma_i x_i] \leq \frac{r \sqrt{2 \log (|\mathcal{A}|)}}{m}$$

where $\sigma_i$'s are independent uniform random variables taking values in $\{-1, 1\}$ and $x_1, \dots, x_m$ are the components of the vector $x$.

To bound, $| \mathcal{F} |$, we use the Natarajan lemma. We observe that $\mathcal{A}^{\infty}(\cdot)$ has at most $| \mathcal{S} |$ values it can take.

The Natarajan lemma for a hypothesis class $\Theta \rightarrow [k]$ states that $|\mathcal{H}| \leq (|\Theta | k^2)^{\text{Ndim}(\mathcal{H})}$

Thus we get the bound on the Rademacher complexity of

$$\mathcal{R}^N(\mathcal{G}) \leq \sqrt{\frac{2 d_N(\mathcal{A}^{\infty}(\mathcal{H})) \log (N |\mathcal{S}|^2)}{N}}$$

In the case of a linear prediction first module and where $\theta$ enters linearly in $f$, we can bound the Natarajan dimension by $p n + 1$ where $\theta \in \mathbf{R}^p$

Let $\psi: \Theta \times \mathcal{S} \rightarrow \mathbf{R}^{d,n}$ be a function that takes in a feature vector and label and outputs a matrix. Let
$\mathcal{H}_{\psi} = \{\theta \rightarrow \argmax_{x \in S} \Tr(B^T \psi(\theta, x)) \}$.

\begin{align*}
\mathcal{A}^{\infty}(B \theta) & = \Pi (2 \textbf{prox}_{\gamma f}(B \theta) - B \theta) \\
&= \argmin_{x \in \mathcal{X}} \frac{1}{2}\| x - (2 H - I)(B \theta) + 2h\|_2^2 \\
&= \argmax_{x \in \mathcal{X}} -\frac{1}{2} \|x\|_2^2 - 2h^T x - \Tr(B^T (2 H - I)\theta x^T ) \\
\end{align*}

Theorem 7 in Daniely and Shalev-Shwartz (2014) shows that $d_N(\mathcal{H}_{\psi}) \leq dn$. We have an additional constant term, so we can consider our prediction variable as having an additional dimension that is fixed to one to take care of this bias term. Our function is $\psi(\theta, x) = (2 H - I)\theta x^T$.

Since $\mathcal{A}^{\infty} \subset \mathcal{H}_{\psi}$, $d_N(\mathcal{A}^{\infty}(\mathcal{H}_{lin})) \leq dn + 1$.

§.§ Non-convex case

In general, we cannot use the above argument if $\mathcal{X}$ is not a convex set. If we use $\mu=0$ and $k=\infty$, we can think of our neural network as one that predicts a point and then the ADMM iterates map this point to the local minimum of this locally convex problem. Thus, we want the prediction to be inside the correct locally convex ball. We now consider this as a classification problem with the empirical 0-1 loss as

$$L^{0-1}(h, \theta) = \mathbf{1}_{ \{\mathcal{A}^\infty(h(\theta)) = x^\star(\theta) \}}.$$

The empirical risk, which gives the proportion of predictions that are in the convex ball, is

$$\hat{R}^{0-1}(h) = \frac{1}{N} \mathbf{1}_{ \{\mathcal{A}^\infty(h(\theta_i)) = x^\star(\theta_i) \}}.$$

The risk is

$$R^{0-1}(h) = \mathbf{E}_{\theta \sim \mathcal{D}}[L^{0-1}(h, \theta)].$$

In general, we cannot use the above argument if $\mathcal{X}$ is not a convex set. Our operator $\mathcal{T}_{\mu}$ is only locally Lipschitz, rather than globally Lipschitz. At points where the projection is not single-valued, $\mathcal{T}_{\mu}$ can be multi-valued and thus cannot be Lipschitz. We will call these points degenerate. Formally, the set of degenerate points is

$$\mathcal{X}^{\rm{deg}} = \{x | |\mathcal{T}_{\mu}(x)| > 1\}.$$

We denote the distance to this set as

\nu(x) =
\begin{array}{ll} \mbox{minimize} & \|x-z\|_2 \\
\mbox{subject to} & z \in \mathcal{X}^{\rm{deg}}
\end{array}

And set the argmin as

u(x) =
\begin{array}{ll} \mbox{argmin} & \|x-z\|_2 \\
\mbox{subject to} & z \in \mathcal{X}^{\rm{deg}}
\end{array}

Now we rewrite our operator $\mathcal{T}_{\mu}$ so that it is globally Lipschitz.

\mathcal{T}_{\mu}^{l}(z) =
\begin{cases}
\mathcal{T}_{\mu}(z) & \nu(z) \geq l \\
\frac{\nu(z)}{l}\mathcal{T}_{\mu}(z) + (1-\frac{\nu(z)}{l}) u(z)  & \nu(z) < l
\end{cases}

The operator $\mathcal{T}_{\mu}^{l}(z)$ is globally Lipschitz.

The main issue concerns the situation where we have two points that are near points of degeneracy. Then using the old operator, $\mathcal{T}_{\mu}$, it's easy to see the global Lipschitz property is not satisfied.
Now, consider two points, $z_1$ and $z_2$ and a degenerate point that is closest to both of them, $\bar{z}$, i.e. $\nu(z_1) \leq l$ and $\nu(z_2) \leq l$ and $u(z_1) = u(z_2)$.

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  = \| \frac{\nu(z_1)}{l}\mathcal{T}_{\mu}(z_1) + (1-\frac{\nu(z_1)}{l}) u(z_1)  -  \frac{\nu(z_2)}{l}\mathcal{T}_{\mu}(z_2) - (1-\frac{\nu(z_2)}{l}) u(z_2)\|_2$$

But since $u(z_1) = u(z_2)$,

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  = \| \frac{\nu(z_1)}{l}\mathcal{T}_{\mu}(z_1) + (\frac{\nu(z_2) - \nu(z_1)}{l})  u(z_1)  +  \frac{\nu(z_2)}{l}\mathcal{T}_{\mu}(z_2) \|_2$$

By the triangle inequality,

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  \leq \| \frac{\nu(z_1)}{l}\mathcal{T}_{\mu}(z_1)  -  \frac{\nu(z_2)}{l}\mathcal{T}_{\mu}(z_2) \|_2 + \|(\frac{\nu(z_2) - \nu(z_1)}{l})  u(z_1)\|_2$$

Since $\nu(z_1)$ and $\nu(z_2)$ both use the same degenerate point, we can use the fact that $|\nu(z_1) - \nu(z_2)| \leq \|z_1 - z_2\|$, also from the triangle inequality.

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  \leq \| \frac{\nu(z_1)}{l}\mathcal{T}_{\mu}(z_1)  -  \frac{\nu(z_2)}{l}\mathcal{T}_{\mu}(z_2) \|_2 + \|(\frac{\nu(z_2) - \nu(z_1)}{l})  u(z_1)\|_2$$

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  \leq \| \frac{\nu(z_1)}{l}\mathcal{T}_{\mu}(z_1)  -  \frac{\nu(z_2)}{l}\mathcal{T}_{\mu}(z_2) \|_2 + \frac{\| u(z_1)\|_2}{l} \|z_1 - z_2\|_2$$

Without loss of generality, assume that $\nu(z_1) \geq \nu(z_2)$. Then,

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  \leq \frac{\nu(z_1)}{l} \| \mathcal{T}_{\mu}(z_1) - \mathcal{T}_{\mu}(z_2) \|_2 + \frac{\| u(z_1)\|_2}{l} \|z_1 - z_2\|_2$$

$$\| \mathcal{T}_{\mu}(z_1) - \mathcal{T}_{\mu}(z_2)\|_2 \leq \| \mathcal{T}_{\mu}(z_1) - \mathcal{T}_{\mu}^1(\bar{z})\|_2 + \|\mathcal{T}_{\mu}(z_2) - \mathcal{T}_{\mu}^1(\bar{z})\|_2$$

where $\mathcal{T}_{\mu}^1(\bar{z})$ is the $\mathcal{T}_{\mu}(\bar{z})$ that is closest to $z_1$.

$$\| \mathcal{T}_{\mu}(z_1) - \mathcal{T}_{\mu}(z_2)\|_2 \leq \| \mathcal{T}_{\mu}(z_1) - \mathcal{T}_{\mu}^1(\bar{z})\|_2 + \|\mathcal{T}_{\mu}(z_2) - \mathcal{T}_{\mu}^2(\bar{z})\|_2 + \|\mathcal{T}_{\mu}^1(z_2) - \mathcal{T}_{\mu}^2(\bar{z})\|_2$$

where $\mathcal{T}_{\mu}^2(\bar{z})$ is the $\mathcal{T}_{\mu}(\bar{z})$ that is closest to $z_2$.

We denote $T_{\max}$ to be the largest jump the operator $\mathcal{T}_{\mu}$ can make, i.e.

$$T_{\max} = \max_z \text{dist} (\mathcal{T}_{\mu}(z))$$

By the locally contractive of $\mathcal{T}_{\mu}$ and using the definition of $T_{\max}$, we get

$$\| \mathcal{T}_{\mu}(z_1) - \mathcal{T}_{\mu}(z_2)\|_2 \leq \| T_{\max} + 2 \|z_1-z_2\|_2$$

$$\|\mathcal{T}_{\mu}^{l}(z_1) - \mathcal{T}_{\mu}^{l}(z_2)\|_2  \leq \frac{ \nu(z_1)}{l} (2\| z_1 - z_2 \|_2 + T_{\max}) + \frac{\| u(z_1)\|_2}{l} \|z_1 - z_2\|_2$$

§ VANISHING GRADIENT

The function $L(z) = f(x) + \frac{1}{2 \mu} d^2(x)$ where $x$ is defined by $x = \textbf{prox}_{\gamma f}(z)$ is smooth.

Recall that the function $f(x)$ is convex and need not be smooth. It can incorporate equality constraints. The inequality constraints will be moved to the projection step. Since $f$ is convex, its proximal operator is always single-valued. By looking at the KKT conditions for the prox, we can see that solving the linear system leads to $x$ being linear in $z$. We can write this as $x(z) = Hz + h$. Thus we can write the loss explicitly as a function of $z$ and remove $x$ altogether.

L(z) = f(Hz+h) + \frac{1}{2 \mu} d^2(Hz+h)

Let's write $f(x) = \bar{f}(x) + \mathcal{I}_C(x)$. Here, $\bar{f}$ is smooth and $C = \{x \mid Ax=b\}$ .

So by rewriting, we have

L(z) = \bar{f}(Hz+h) + \mathcal{I}_C(Hz+h) + \frac{1}{2 \mu} d^2(Hz+h)

But notice that by definition of solving the proximal optimization problem, $\mathcal{I}_C(Hz+h) = 0$, so wee can get rid of it.

L(z) = \bar{f}(Hz+h) + \frac{1}{2 \mu} d^2(Hz+h)

The composition of two smooth functions is smooth. The first terms is the composition of $\bar{f}$ which is smooth by assumption and a linear term which is always smooth. For the second term, since $\nabla d^2(x) = 2(\Pi(x)-x)$ and the projection is locally Lipschitz continuous, we see that $d^2$ is smooth and the composition with the linear piece is also smooth.

The vanishing gradient. The gradient $\nabla L(z^0)$ decays exponentially in the number of ADMM iterates.

By the chain rule,

$$\nabla L(z) = \frac{\partial z^1}{\partial z^0} \dots \frac{\partial z^k}{\partial z^{k-1}} \frac{\partial L}{\partial z^k}$$

We apply Cauchy-Schwarz to get

$$\nabla L(z) \leq \| \frac{\partial z^1}{\partial z^0}\| \dots \| \frac{\partial z^k}{\partial z^{k-1}} \| \| \frac{\partial L}{\partial z^k}\|$$

We can deal with the last piece. We use the fact that $L(z^k)$ is $\beta$ smooth. Therefore, $\|\frac{\partial L}{\partial z^k}\|_2 \leq \beta$.

For the other pieces, we use that the fact that $T_{\mu}$ is a contraction where $z^{i+1} = T_{\mu}(z^i)$. The object we want to bound is the norm of the Jacobian, $\| \frac{\partial z^{i+1}}{\partial z^i}\|$.

By definition of the Jacobian, for every vector $\epsilon$ and $p$ such that $\|p\|_2 = 1$, and for small enough $h$ we have the following inequality where $J(z)$ is the Jacobian of the functiono $T_{\mu}$ at the point $z$. We assume differentiability of $T_{\mu}$ which is true almost everywhere.

\frac{T_{\mu}(z + hp) - T_{\mu}(z)}{h} - J(z) p \leq o(h)

Then for every $\epsilon > 0$ we can make $h$ small enough such that

\|J(z) p\| - \frac{\|T_{\mu}(z + hp) - T_{\mu}(z)\|}{h} \leq \epsilon

Hence by $T_{\mu}$ being locally contractive with factor $\kappa \in (0, 1)$, we get

\|J(z) p\| \leq \epsilon + \kappa

By definition of the largest singular value of $J(z)$, $\|J(z)\|_2 = \max_{\{p \mid \|p\|_2=1\}}  \|J(z) p\|$. Since $\|J(z) p\| < \kappa$ for every unitary $p$, we get $\|J(z)\|_2 < \kappa$

Now we are ready to provide a bound on $\nabla L(z^0)$.

\|\nabla L(z^0)\| \leq \kappa^k \beta

As $k$ goes to infinity, the gradient vanishes.

§ IMPLEMENTATION

§.§ Smart initialization

In order to get a better starting point, we choose to do a pretraining step. In this step, we first run the full outer nexos algorithm on each of the training problems with possibly many random initializations. Then, we take the best solution for each problem and label as $\{\bar{x}_i\}_{i=1}^{N_{train}}$. The pretraining loss is defined as

$$L^{pretrain} = \sum_{i=1}^{N_{train}} \|h(\theta_i) - \bar{x}_i\|_2^2$$

where $h(\theta)$ maps the input to the output of the standard neural network layers.

§.§ Cache Factorizations

In many practical applications, the KKT matrix that appears in the linear system which needs to be solved in the proximal step of the function $f$ does not change from problem to problem. In this work we factorize it offline and cache it. By doing so, our online nexos iterates only require forward-backward substitutions that we can carry out very efficiently.

§.§ Closed-form steps

<cit.> solves (<ref>) by repeatedly evaluating the proximal operator of the function $f$. In some cases, this cannot be done in closed-form and thus, requires an iterative algorithm to solve. For example, if the function $f$ includes inequality constraints, then there is no closed-form solution. We improve upon this idea by moving inequality constraints into the projection step onto $\mathcal{X}$. Consider a closed, convex set $\mathcal{C}$ and a non-convex set $\mathcal{X}$. Denote the union of the two sets as $\mathcal{Y} = \mathcal{X} \cup \mathcal{C}$. If $\mathcal{X}$ is prox-regular at $x \in \mathcal{X}$ then $\mathcal{Y}$ is prox-regular at $x \in \mathcal{Y}$.

So if $f(x) = g(x) + i(x)$ where $g$ admits a closed-form solution and $i(x)$ is the indicator function of a convex set $\mathcal{X}$, then we can rewrite the problem as the following where $\mathcal{Y} = \mathcal{X} \cup \mathcal{C}$. Both the proximal and projection steps can be evaluated quickly.

\begin{equation}\label{eq:inner-nexos-prob}
\begin{array}{ll}
\text{minimize} & g(x) + \frac{\beta}{2}\|x\|_2^2 + \frac{1}{2 \mu} d_{\mathcal{Y}}^2(x)\\
\end{array}
\end{equation}

§.§ JAX automatic differentiation

Just-in-time (jit) compilation. jaxopt. TODO cite

§ COMPUTATIONAL BENCHMARKS

We benchmark the learn2convexify scheme on several parametric optimization problems. We compare the computational time to solve the problems to OSQP, SCS, and CVXOPT.

Infeasibility and suboptimality. We use the same criteria for calculating relative suboptimality and infeasibility as in <cit.>. We define the suboptimality of a solution, $\hat{x}$ as

$$d(\hat{x}) = \frac{f(\hat{x}, \theta) - f(x^\star, \theta)}{|f(x^\star, \theta)|}$$

§.§ Markowitz

\begin{equation}\label{eq:sparse-markowitz}
\begin{array}{ll}
\text{minimize} & \frac{1}{2}x^T \Sigma x - \rho \mu^T x\\
\text{subject to} & \ones^T x = 1\\
& x \geq 0
\end{array}
\end{equation}

Parameters. The problem parameters are $\theta = \mu$ which corresponds to the returns vector. It is typical for the covariance matrix to be updated once a month.

Problem Setup.
We use S$\&$P500 returns data from 2007 to 2017. $\Sigma$ is computed as the sample covariance matrix from the returns. The $\mu$ vector changes for each problem in and is drawn from an i.i.d. distribution in the following way. $\mu \sim \mathcal{N}(\mu_0, \sigma)$ where $\mu_0$ is the vector of average returns and $\sigma$ is the standard deviation for each asset. Although this problem may not need to be solved in real-time, it is important to solve quickly to perform many backtesting experiments. We use $\rho = 1$.

For this problem, we need to split the non-negativity constraints from the single equality constraint in ADMM. In this particular instance, we are able to take our solution as $\hat{x} = \frac{y^k}{\mathbf{1}^T y^k}$. $x^k$ always satisfies feasibility assuming that $y^k \neq 0$. We just need to check suboptimality. We are able to solve the Markowitz problem with only 5 ADMM iterations with a suboptimality level of about $0.02$ while satisfying feasibility.

Convex Markowitz training results for five unrolled iterates.

Convex Markowitz benchmarks on test data.

Method     Time (seconds)     Suboptimality     Infeasibility

[head to column names, late after line=

§.§ Path-planning

§ APPENDIX

Case where constraints are split.

Assume that $\|z^0-z^\star\| \leq B$ and that $\bar{f}$ is $\alpha$-smooth. Then $L(\mathcal{A}^k(z^0))$ is Lipschitz with factor $\kappa^k B ( (\alpha + \frac{1}{2 \mu}) \kappa^k B + \|\nabla \bar{f}(x^\star) \|_2 )$.

The composition of two Lipschitz functions with factors $\alpha_1$ and $\alpha_2$ is a Lipschitz function with factor $\alpha_1 \alpha_2$. Our function $L(\mathcal{A}^k(z^0))$ consists of the composition of many Lipschitz functions. $T_{\mu}$ is Lipschitz with factor $\kappa$. $\textbf{prox}_{\gamma f}$ is nonexpansive. The last part to deal with is the loss $L(x^k) = f(x^k) + \frac{1}{2 \mu}d^2(x^k) - f(x^\star)$.

Since $x^k = \textbf{prox}_{\gamma f}(z^k)$ it will satisfy the equality constraint. Therefore we only need to look at $\bar{f}$ and can write the loss as $L(x^k) = \bar{f}(x^k) + \frac{1}{2 \mu}d^2(x^k) - f(x^\star)$. By the contractive property of $T_{\mu}$, we have $\|z^k-z^\star \|_2 \leq \kappa^k \|z^0 - z_star\|_2 \leq \kappa^k B$. Since $\textbf{prox}_{\gamma f}$ is nonexpansive, $\|x^k-x^\star \|_2 \leq \kappa^k B$.

Then using the fact that $\bar{f}(x)$ is $\alpha$-smooth and that $\|x^k-x^\star\|_2 \leq \kappa^k B$ we can bound $\| \nabla \bar{f}(x^k) \|_2$.

By the smoothness of $\bar{f}(x)$,

$$\| \nabla \bar{f}(x^k) - \nabla \bar{f}(x^\star)\|_2 \leq \alpha \|x^k-x^\star\|_2 \leq \alpha \kappa^k B$$
. Using the triangle inequality,

$$\|\nabla \bar{f}(x^k)\|_2 \leq \|\nabla \bar{f}(x^k) - \nabla \bar{f}(x^\star)\|_2 + \|\nabla \bar{f}(x^\star)\|_2$$

Plugging in gets

$$\|\nabla \bar{f}(x^k)\|_2 \leq \alpha \kappa^k B + \|\nabla \bar{f}(x^\star)\|_2$$

Now to find the Lipschitz constant for the other term, $d^2(x)$.
$d^2(x^k) = \|\Pi(x^k)-x^k\|_2^2$ and $\nabla d^2(x^k) =2  \| \Pi(x^k)-x^k\|_2$

Since $\mathcal{X}$ is convex by assumption, $\Pi(x)$ is single-valued and nonexpansive.

$$\| \Pi(x^k)-x^k\|_2 \leq \|x^k - x^\star\|_2 \leq \kappa^k B$$

Thus we have

$$\frac{1}{2 \mu} \nabla d^2(x^k) \leq \frac{1}{2 \mu} \kappa^k B$$

Consider the following parametric optimization problem.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & f(x, \theta)\\
\text{subject to} & g_i(x, \theta)\le 0, \quad \forall i = 1, \dots, m\\
& h_j(x, \theta) = 0, \quad \forall j = 1, \dots, p
\end{array}
\end{equation*}

$x$ is the decision variable and $\theta$ is the parameter.

If $f, g_1, \dots, g_m$ are all convex functions in $x$ and $h_1, \dots, h_p$ are all affine in $x$ then this optimization problem is convex. If one of them is non-convex, then the optimization problem is non-convex. Solving non-convex optimization problems is considerably more difficult than convex ones. Typically, we cannot expect to find the global solution quickly.

In many real-time applications, we need to solve non-convex problems on the order of milliseconds. State-of-the-art or commercial-grade solvers are nowhere near sufficient for these applications. Most of these algorithms are general-purpose and tackle each optimization instance individually. In practice, however, we often solve a family of very similar problems with varying data. The central question to this paper is to reuse the effort of solving these problems offline to help solve a similar one online.

There are several ways to convexify a non-convex problem given the current iterate $x^k$. One way would be to add quadratic penalty terms to $f, g_1, \dots, g_m$ and take the first order Taylor expansion for $h_1, \dots, h_p$. That is, $\hat{f}(x) = f(x) + \frac{1}{2 \gamma}\|x-x^k\|_2^2$. We can convexify the constraint $\hat{g_i}(x) = g_i(x) + \frac{1}{2 \gamma}\|x-x^k\|_2^2$. We can find an affine approximation of $h_j(x)$ by taking $\hat{h_j}(x) = h(x^k) + \nabla h(x^k)^T (x - x^k)$. By doing these approximations, we have formed a convex optimization problem that behaves similarly to the non-convex one in the region around $x^k$.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \hat{f}(x)\\
\text{subject to} & \hat{g_i}(x)\le 0, \quad \forall i = 1, \dots, m\\
& \hat{h_i}(x) = 0, \quad \forall i = 1, \dots, p
\end{array}
\end{equation*}

Another common technique is to use a trust region around the current iterate. The reasoning is that the convexified problem is locally accurate but is decreasingly reliable the further you move away from the current iterate. This trust region problem can be written as follows.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \hat{f}(x)\\
\text{subject to} & \hat{g_i}(x)\le 0, \quad \forall i = 1, \dots, m\\
& \hat{h_i}(x) = 0, \quad \forall i = 1, \dots, p \\
& \|x - x^k\|_2 \leq \rho
\end{array}
\end{equation*}

Another variant is to relax the trust region constraint as follows.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \hat{f}(x) + \frac{1}{2 \gamma} \|x - x^k\|_2^2 \\
\text{subject to} & \hat{g_i}(x)\le 0, \quad \forall i = 1, \dots, m\\
& \hat{h_i}(x) = 0, \quad \forall i = 1, \dots, p
\end{array}
\end{equation*}

Sequential convex programming has no convergence guarantees at all. It's performance can heavily depend on the initial point. Additionally, it take take many iterates to get a good solution. In our work, we propose a method to use machine learning to solve similar non-convex optimization problems. In it, we use a neural network to map the differing parameters of each non-convex problem, $\theta$, to a vector $y$ who should serve as a good approximation of $x^k$. Then using $y$, we solve a convex problem to obtain the solution to the non-convex one. Since it is convex, we can solve this problem orders of magnitude more quickly than the original non-convex one. Our method is valuable in cases where we need to solve non-convex problems in real-time.

For the rest of the paper, we remove the affine constraint for simplicity of notation. To transform an affine constraint into inequality form, we simply require both greater or equal to and less than or equal to inequalities.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & f(x, \theta)\\
\text{subject to} & g_i(x, \theta)\le 0, \quad \forall i = 1, \dots, m
\end{array}
\end{equation*}

§ DIFFERENTIATING THROUGH OPTIMIZATION PROBLEMS

§.§ Differentiable Optimization

To differentiate through the QP where $Q \succ 0$, we use the work of Amos TODO:cite.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2} x^T Q x + c^T x \\
\text{subject to} & Ax = b \\
& Gx \leq h \\
\end{array}
\end{equation*}

To compute the vector Jacobin product (VJP), we need to solve the following linear system.

\begin{equation*}
\begin{bmatrix}
d_x \\
d_\lambda \\
\end{bmatrix} = -
\begin{bmatrix}
Q & G^T D(\hat{\lambda}) &  A^T\\
G & D(G\hat{x} - h) & 0 \\
A & 0 & 0
\end{bmatrix}^{-1}
\begin{bmatrix}
\frac{\partial L}{\partial x} \\
\frac{\partial L}{\partial \lambda} \\
\frac{\partial L}{\partial \nu}
\end{bmatrix}
\end{equation*}

In our case, we are mainly interested in computing $\frac{\partial L}{\partial q} = d_x$. If the solution is non-degenerate, then the matrix in this linear system is invertible. If the solution is degenerate, i.e. $\exists i$ such that $\hat{\lambda}_i = 0$ and $(G\hat{x} - h)_i = 0$, then there will be no solution. In this case, we solve a least squares problem to minimize the loss.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \|K d_z - y\|_2^2
\end{array}
\end{equation*}

K = -\begin{bmatrix}
Q & G^T D(\hat{\lambda}) &  A^T\\
G & D(G\hat{x} - h) & 0 \\
A & 0 & 0
\end{bmatrix},
d_z = \begin{bmatrix}
d_x \\
d_\lambda \\
\end{bmatrix},
y = \begin{bmatrix}
\frac{\partial L}{\partial x} \\
\frac{\partial L}{\partial \lambda} \\
\frac{\partial L}{\partial \nu}
\end{bmatrix}

§.§ Envelope Theorem

If we want to compute the derivative of the objective of an optimization problem with respect to its parameters, we can use the Envelope theorem. In this case, there is no need to compute the Jacobians $\frac{\partial \hat{x}}{\partial \theta}, \frac{\partial \hat{\lambda}}{\partial \theta}$. This can be advantageous because it does not require solving a linear system which can be expensive.

Consider the optimization problem

\begin{equation*}
\begin{array}{ll}
\text{minimize} & f(x, \theta)\\
\text{subject to} & g_i(x, \theta)\le 0, \quad \forall i = 1, \dots, m
\end{array}
\end{equation*}

Then define the value function as $V(\theta) = f(\hat{x}(\theta), \theta)$. The derivative can then be computed as follows,

$$\nabla_{\theta} V(\theta) = \nabla_{\theta} f(\hat{x}, \theta) + \sum_{i=1}^{m} \hat{\lambda}_i \nabla_\theta g_i(\hat{x}, \theta).$$

Here, we write $\nabla_{\theta} f(\hat{x}, \theta)$  where $\hat{x}$ is not a function of $\theta$. So we do not need to calculate the Jacobian $\frac{\partial \hat{x}}{\partial \theta}$. The same can be said for the other term with $\lambda$.

§ LEARNING STATEMENT

Given a family of $N$ training problems with parameters $(\theta_i, \dots, \theta_N)$, we aim to minimize the loss of a neural network. The first part of the neural network predicts a vector $y$ for each $\theta$. This first part can consist of standard neural network layers with activation functions such as relu. The second part of the neural network is the convexified optimization problem which takes $y$ as an input. The output is the solution $\hat{x}$. If the convexified problem has a strongly convex objective, then this mapping is a function. Furthermore, the convexified problem is always feasible. Therefore it has an unrestricted domain. The loss of the neural network can either be unsupervised or supervised. The distinction between the two is made by seeing if the loss uses the information of the optimal solution to the non-convex optimization problem. Then, the neural network training is aiming to do imitation learning on the optimal solutions. We denote the primal and dual optimal solutions of the ith convexified problem given $y_i$ as $(\hat{x}_i(y_i), \hat{\lambda}_i(y_i))$. The subscript under the functions $\hat{x}_i$ and $\lambda_i$ make it clear that each training problem parametrized by $y$ is different since the problem data will be different. The first part of the neural network maps $\theta$ to $y$. Since in the general case, $y$ should not have any domain restrictions we don't use any activation function to go from $z_i^{l-1}$ to $y_i$. The decision variables in this training optimization problem are the weight and bias terms of the neural network of the various layers, $\{(W^k, b^k)\}_{k=1}^{l}$. The optimal primal and dual non-convex solutions of the ith training problem are $(x^\star_i, \lambda^\star_i)$.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \sum_{i=1}^N \frac{1}{2} \|x^\star_i - \hat{x}_i(y_i)\|_2^2 +  \frac{1}{2} \|\lambda^\star_i - \hat{\lambda}_i(y_i)\|_2^2\\
\text{subject to} & y_i = W^l z^{l-1}_i + b^{l-1} \quad \forall i=1, \dots, N \\
& z^k_i = \phi^k(W^k z^{k-1}_i + b^k) \quad \forall k=1, \dots, l-1, \quad \forall i=1, \dots, N \\
& z^0_i = \theta_i \quad \forall i=1, \dots, N
\end{array}
\end{equation*}

If instead, the loss is unsupervised, then the loss does not use the non-convex optimal solutions. Rather, we directly minimize either the non-convex loss directly or a proxy for it. If we directly aim to minimize the objectives of the training problems, then this unsupervised problem is as follows.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \sum_{i=1}^N f_i(\hat{x}_i(y_i)) \\\
\text{subject to} & y_i = W^l z^{l-1}_i + b^{l-1} \quad \forall i=1, \dots, N \\
& z^k_i = \phi^k(W^k z^{k-1}_i + b^k) \quad \forall k=1, \dots, l-1, \quad \forall i=1, \dots, N \\
& z^0_i = \theta_i \quad \forall i=1, \dots, N
\end{array}
\end{equation*}

There is another unsupervised loss function we consider. In this loss, we penalize the envelope rather than the non-convex objective. The envelope serves as a proxy to the non-convex objective. We use this method to minimize over the weights of the neural network due to theoretical properties that this envelope gives. It allows us to use the Envelope theorem rather than the implicit function theorem which means we don't need to solve a linear system. Additionally, we find that the optimization landscape is friendlier over the loss with the envelope rather than either the non-convex objectives or the imitation learning losses.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \sum_{i=1}^N V_{i,\gamma} (y_i)  \\\
\text{subject to} & y_i = W^l z^{l-1}_i + b^{l-1} \quad \forall i=1, \dots, N \\
& z^k_i = \phi^k(W^k z^{k-1}_i + b^k) \quad \forall k=1, \dots, l-1, \quad \forall i=1, \dots, N \\
& z^0_i = \theta_i \quad \forall i=1, \dots, N
\end{array}
\end{equation*}

In the above problem, $V_{i,\gamma} (y_i)$ is equal to the objective of the convexified ith problem with parameter $y_i$. As mentioned, there are many ways to convexify problems. In the section called Envelopes, we discuss several of these functions: Moreau, Forward-backward, and SCP.

§ ENVELOPES

§.§ Moreau Envelope (Weakly convex only)

If the original problem is weakly convex, then we consider the Moreau envelope. The Moreau envelope is a function equal to the objective of the non-convex problem with an additional quadratic penalty in the objective, $\frac{1}{2 \gamma}\|x-y\|_2^2$. For small enough $\gamma$, the Moreau envelope is equal to the minimization of a convex problem.

We have that the Moreau envelope is equal to the minimal objective of the following optimization problem.

\begin{equation*}
M_\gamma (y) = \begin{array}{ll}
\text{minimize} & f(x) + \frac{1}{2 \gamma} \|x-y\|_2^2  \\\
\text{subject to} & g_i(x) \leq 0 \quad \forall i=1, \dots, m
\end{array}
\end{equation*}

Assuming $\gamma$ small enough, $M_\gamma(y)$ is convex. The minimization of $M_\gamma (y)$ corresponds to stationary points of the original non-convex problem.

Applying gradient descent to the Moreau envelope is equivalent to applying the proximal point method which is

$$x^{k+1} = \textbf{prox}_{\gamma(f + g)}(x^k)$$

where $g(x)$ is a non-smooth function defined as

$$g(x) = \sum_{i=1}^m \mathcal{I}(g_i(x) \leq 0).$$

The gradient of the Moreau envelope is

$$\nabla M_\gamma(y) = \frac{1}{\gamma} (y - \textbf{prox}_{\gamma(f + g)}(y)).$$

§.§ Forward Backward Envelope (Weakly convex only)

If we have the optimization problem

$$\min_x f(x) + g(x)$$

where $g$ is closed, convex, and proper and $f$ is smooth, one technique to solve this problem is the proximal point method as mentioned above. Another popular method to solve this problem is the proximal gradient descent algorithm. The iterates are

$$x^{k+1} = \text{prox}_{\gamma g}(x^k - \gamma \nabla f(x^k)).$$

The question of whether these iterates are equivalent to performing gradient descent on some differentiable, unconstrained optimization problem has been answered affirmatively. The forward backward envelope (FBE) is this loss function. In this case, we linearize the non-convex part and add a quadratic penalty term. The FBE is defined as

$$\phi_\gamma (y) = \min_x l(x,y) + \frac{1}{2 \gamma} \|y-x\|_2^2$$

where $l(x,y) = f(y) + \nabla f(y)^T (x - y) + g(x)$

Notice that the subproblems here are strongly convex for any value of $\gamma$. This is advantageous since many algorithms are better-suited for optimization problems  where the objective is more strongly convex. Unlike the Moreau envelope, the FBE is not convex even if $f$ and $g$ are convex. We can use both the Moreau envelope and the FBE to find stationary points of the original non-convex problem.

§.§ Sequential Convex Programming Envelope (Any non-convex)

We propose a new envelope called the Sequential Convex Programming Envelope (SCPE). This envelope is the only one that can handle non-convex constraints out of the three we mention. It is defined as

\begin{equation*}
\psi_\gamma (y) = \begin{array}{ll}
\text{minimize} & f(x) + \frac{1}{2 \gamma} \|x-y\|_2^2  \\\
\text{subject to} & g_i(y) + \nabla g(y)^T (x - y) + \frac{1}{2} (x - y)^T P_i^+ (x - y) \leq 0 \quad \forall i=1, \dots, m
\end{array}
\end{equation*}

Here, $P_i^+$ is the projection of the Hessian of $g_i(x)$ onto the positive semidefinite cone. We can also use the first order approximation of $g_i(x)$ instead.

This convex optimization problem has a strongly convex objective. Thus it always has a unique optimal solution. Furthermore, $\psi_\gamma (y)$ is differentiable  everywhere. If we ignore the second order term, then the gradient is

$$\nabla \psi_\gamma (y) = \hat{\lambda}(y)^T g(y)  + (\frac{1}{\gamma} I + \sum_{i=1}^m (\hat{\lambda}(y))_i \nabla^2 g_i(y)) (\hat{x}(y) - y).$$

One drawback of sequential convex programming is that the iterates may not be feasible.

§ UNSUPERVISED LOSS

Even in time-insensitive (offline) settings, it can be prohibitively expensive to solve non-convex problems. Thus, in these settings, it can be advantageous to use our method without the need to solve the non-convex problems in offline.

In the objective-based loss, the objective and loss are equivalent. That is, we no longer need to solve for the non-convex solutions offline and instead try to minimize the Moreau envelope directly. However, notice that the Moreau envelope is not equal to the loss of the non-convex problem since the former includes the quadratic term. We can use the Envelope theorem to differentiate the through the loss (i.e. we don't need to use differentiable optimization). The loss is now an unconstrained optimization problem over $y$.

L(y) = \frac{1}{2} \hat{x}^T Q \hat{x} + c^T \hat{x} + \frac{1}{2 \gamma} \|x-y\|_2^2
where $\hat{x}$ is the minimizer of the optimization problem.

Furthermore, $L(q)$ is convex, continuous, and piece-wise quadratic. However, the minimizer of $L(y)$ is only a stationary point of the original non-convex problem. Regardless, since $L(y)$ is convex, we can minimize over it using subgradient descent. We can also reformulate this is a QP. This is further explained in the next part.

§.§ Multiple training problems

Of course, in training, we want to predict $y$ from $\theta$ given many problem instances. In this section, we focus on the linear prediction $y = B \theta$ where $\theta \in \mathbf{R}^p$ (the number of features of each problem) and $B \in \mathbf{R}^{n,p}$. Let $N$ denote the number of training problems.

\min \frac{1}{N} \sum_{i=1}^N g_{\gamma}^i (B \theta_i)

where $g_{\gamma}^i (y)$ denotes the Moreau envelope of the i-th problem (as seen in the FBE section). This problem is still convex and differentiable almost everywhere. It can be solved for with subgradient descent. Alternatively, we can use a duality trick and solve this optimization problem with a single QP.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{N} \sum_{i=1}^N \frac{1}{2 \gamma} \|B \theta_i\|_2^2 \frac{1}{2}  v_i^T (Q + \frac{1}{\gamma} I)^{-1} v_i \\
\text{subject to} & \lambda_i \geq 0 \\
& v_i = c - \frac{1}{\gamma} B \theta_i + G^T \lambda_i \\
\end{array}
\end{equation*}

Here, I've omitted the $Ax = b$ constraints, and only entereed $Gx \leq h$ constraints.

Proof sketch: Write full proof (TODO)
1. Form Lagrangian of the $N$ problems
2. Write dual of each of them ($\lambda_i$ is the dual variable vector for the ith problem)
3. Introduce the new variable $v_i = c  - \frac{1}{\gamma} B \theta_i + G^T \lambda_i$
4. Now flip the sign of the max to get a min
5. Take advantage of the double min and put the decision variables $B$ and $(\lambda_1, \dots, \lambda_N, v_1, \dots, v_N)$ as the decision variables of this single problem.

§ TOY EXAMPLE TO SHOW IMPORTANCE OF DUAL PENALTY

Simply doing imitation learning on the primal variables alone will not guarantee convergence to the optimal solution. The root of this issue is inequality constraints. Assuming the solution to the parametric optimization problem is non-degenerate, then if the inequality constraint, $x_j \geq 0$, is tight, then
$\frac{\partial x_j}{\partial c}$ will be zero. But in this case, the corresponding dual variable $\lambda_j$ will be non-zero by strict complementarity. Since the complementarity constraint $x_j \lambda_j = 0$ is satisfied for any local perturbation of $\lambda_j$, the partial derivative  $\frac{\partial \lambda_j}{\partial c}$ is free to be non-zero.

Consider the following 1-D optimization problem where $c$ is a parameter.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & x^2 + cx \\
\text{subject to} & x \geq 0
\end{array}
\end{equation*}

The dual problem is

\begin{equation*}
\begin{array}{ll}
\text{maximize} & -\frac{1}{4} \lambda^2 + \frac{1}{2} c \lambda\ \\
\text{subject to} & \lambda \geq 0
\end{array}
\end{equation*}

The optimal solutions are

\[ \hat{x}(c) = \begin{cases}
-\frac{c}{2} & c \leq 0 \\
0 & c > 0
\end{cases}
\]

\[ \hat{\lambda}(c) = \begin{cases}
0 & c \leq 0 \\
c & c > 0
\end{cases}
\]

Say the optimal $x$ is $x^\star = 1$, then the optimal $c$ is $c^\star=-\frac{1}{2}$. The loss function with only the primal penalty, $L_p(c) = \frac{1}{2}\|\hat{x}(c) - x^\star\|_2^2$, is

\[ L_p(c) = \begin{cases}
\frac{1}{2} (1+\frac{c}{2})^2 & c \leq 0 \\
\frac{1}{2} & c > 0
\end{cases}
\]

If we initialize $c$ with any value greater than zero, the derivative $\frac{\partial L_p}{\partial c} = 0$. Hence, doing gradient descent will not work.

However, if we use the loss that we suggest which penalizes the primal and dual variables, then we can obtain convergence with gradient descent assuming small enough step size. The loss is $L(c) =  \frac{1}{2}\|\hat{x}(c) - x^\star\|_2^2 + \frac{1}{2}\|\hat{\lambda}(c) - \lambda^\star\|_2^2$ and can be simplified to

\[ L(c) = \begin{cases}
\frac{1}{2} (1+\frac{c}{2})^2 & c \leq 0 \\
\frac{1}{2}  + \frac{1}{2}c^2 & c > 0
\end{cases}
\]

The derivative is

\[ \frac{\partial L(c)}{\partial c} = \begin{cases}
\frac{1}{2} (1+\frac{c}{2}) & c \leq 0 \\
c & c > 0
\end{cases}
\]

So the derivative is not smooth at the degenerate point $c=0$. However, the loss, $L(c)$ is continuous and piecewise strongly convex. Furthermore, the only point where $\frac{\partial L(c)}{\partial c} = 0$ is the globally optimal point. It is worth noting that $\lim_{c\downarrow0} \frac{\partial L(0)}{\partial c} = 0$. However, in practice, the troublesome point, zero, will never be reached.

For this problem, if we are ever at the degenerate solution, $\hat{x} = \hat{\lambda} = 0$ which would happen when $c = 0$, we would still get convergence to the optimal solution. At this point, the loss is not differentiable, but the direction computed by the least squares solution still allows us to converge. Note that there is no solution to the following linear system.

\begin{equation*}
\begin{bmatrix}
2 & 0 \\
-1 & 0\\
\end{bmatrix}
\begin{bmatrix}
d_x \\
\end{bmatrix} = -
\begin{bmatrix}
1 \\
\end{bmatrix}
\end{equation*}

The least squares solution is given by solving the following least squares problem.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & (-2d_x + 1)^2 + d_x^2
\end{array}
\end{equation*}

The optimal solution is $d_x = 0.4$. Note that this value is not the average of the derivative approached from the left and the right, $0.5$ and $0$ respectively.

§ EXAMPLES

§.§ Path-planning

The goal of path-planning is to find the shortest path connecting $a$ to $b$ that avoids $m$ circles centered at $p_j$ with radius $r_j$. After discretizing the problem into $x_1, \dots, x_T$, the problem becomes

\begin{equation*}
\begin{array}{ll}
\text{minimize} & L \\
\text{subject to} & x_1 = a, x_T=b \\
& \|x_{i+1} - x_i\|_2 \leq \frac{L}{T}, \quad i = 1, \dots, T \\
& \|x_i - p_j\|_2 \geq r_j, \quad i=1, \dots, n, \quad j=1, ..., m. \\
\end{array}
\end{equation*}

The non-convexity in this problem comes from the obstacle avoidance constraint. In order to apply the nexos framework, we introduce new variables. The new variables will allow us to solve the convex optimization problem in closed-form. The equivalent problem is

\begin{equation*}
\begin{array}{ll}
\text{minimize} & L \\
\text{subject to} & x_1 = a, x_T=b \\
& \|\Delta_i\|_2 \leq \frac{l_i}{T}, \quad i = 1, \dots, T-1 \\
& \Delta_i = x_{i+1} - x_i, \quad i = 1, \dots, T-1 \\
& L \mathbf{\ones} = l + s\\
& s \geq 0 \\
& \|x_i - p_j\|_2 \geq r_j, \quad i=1, \dots, n, \quad j=1, ..., m. \\
\end{array}
\end{equation*}

where $L \in \mathbf{R},  l \in \mathbf{R}^{T-1}, s \in \mathbf{R}^{T-1}, x_i \in \mathbf{R}^{2} \quad i= 1, \dots, T, \Delta_i \in \mathbf{R}^{2} \quad i= 1, \dots, T-1$ are the variables.

Now we can rewrite this problem as

\begin{equation*}
\begin{array}{ll}
\text{minimize} & f(z) + i(z) \\
\end{array}
\end{equation*}

where $f$ is a convex function and $i$ is the indicator function of the non-convex set $\mathcal{X}$. $z$ is the concatenation of all the variables.

f(z) = L + \mathcal{I}(L \mathbf{\ones} = l + s) + \sum_{i=1}^{T-1} \mathcal{I}(\Delta_i = x_{i+1} - x_i)

i(z) = \sum_{i=1}^T \sum_{j=1}^m \mathcal{I}(\|x_i - p_j\|_2 \geq r_j) + \mathcal{I}(s \geq 0) + \sum_{i=1}^{T-1} \mathcal{I}(\|\Delta_i\|_2 \leq \frac{l_i}{T})

The benefit of splitting the problem like this is that evaluating the prox of $f$ admits a closed-form solution. Similarly, evaluating the prox of $i$ is a projection that is separable; it has a closed-form solution. It involves a second-order cone projection on the $(\Delta_i,l_i)$ variables, a clipping of $s$, and a projection to avoid the obstacles of the $x_i$ variables.

We assume that there is no overlap in the obstacles. The problem of projecting a point onto the set of all points that is not in any obstacle is simple. By the assumption, a point can only lie in one obstacle circle. The projection of a point $x$ onto the region outside an obstacle with center $p$ and raduis $r$ is given by

$$\pi(x) = p + \frac{r (x - p)}{\|x-p\|_2}$$
if $x$ is in the obstacle and $\pi(x) = x$ otherwise.

The prox of $f$ can be evaluated in closed-form since it is an equality-constrained quadratic program. Thus, we can pick a value for $\mu$ and apply Algorithm 2 from the nexos paper.

§.§ Half-relaxed sparse PCA

Consider the sparse pca problem

\begin{equation*}
\begin{array}{ll}
\text{minimize} & x^T (-A) x \\
\text{subject to} & x^T x \leq 1 \\
& ||x||_0 \leq k
\end{array}
\end{equation*}

We can relax the cardinality constraint and move it into the objective. This is still a non-convex problem.  This is the non-convex optimization problem with weakly convex objective that we want to solve.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & x^T (-A) x + \rho ||x||_1 \\
\text{subject to} & x^T x \leq 1
\end{array}
\end{equation*}

We can convexify it by adding a penalty term with $y$ as a parameter. If $\lambda \leq \lambda_{max}(A)$ then this problem is convex. If the inequality is strict, the problem becomes strongly convex.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & x^T (-A) x + \rho ||x||_1 + \frac{1}{2 \gamma} ||x-y||_2^2\\
\text{subject to} & x^T x \leq 1
\end{array}
\end{equation*}

§.§ Sparse Portfolio

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \gamma (w^{+})^T \Sigma w^{+} - \mu^T w^{+}\\
\text{subject to} & w^{+} =  w_t + z \\
& \mathbf{1}^T z + \kappa^T |z| + \nu^T (w^{+})_- \leq 0 \\
& \|z\|_0 \leq k
\end{array}
\end{equation*}

We can convexify it by adding a penalty term with $y$ as a parameter. If $\gamma \leq \lambda_{max}(A)$ then this problem is convex. If the inequality is strict, the problem becomes strongly convex.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \gamma (w^{+})^T \Sigma w^{+} - \mu^T w^{+} + \frac{1}{2 \gamma} ||z-y||_2^2\\
\text{subject to} & w^{+} =  w_t + z \\
& \mathbf{1}^T z + \kappa^T |z| + \nu^T (w^{+})_- \leq 0
\end{array}
\end{equation*}

§ RELATED WORK

§ OUR METHOD

§.§ Different non-convex cases we want to solve

We focus on a few cases of non-convex problems.

1. The first case is when the objective is non-convex but weakly convex, and the constraints are convex. Roughly speaking, a convex function $f$ is one where we can add a quadratic penalty to make it convex. A non-convex quadratic function would be one example. Consider the following optimization problem where $Q \nsucceq 0$

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2}  x^T Q x + c^T x \\
\text{subject to} & Ax = b \\
& Gx \leq h \\
\end{array}
\end{equation*}

To convexify this non-convex problem we add the Moreau penalty $\frac{1}{2 \gamma} ||x-y||_2^2$ For small enough $\gamma$, this problem is strongly convex and therefore has a unique optimal solution if feasible.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2}  x^T Q x + c^T x + \frac{1}{2 \gamma} ||x-y||_2^2\\
\text{subject to} & Ax = b \\
& Gx \leq h \\
\end{array}
\end{equation*}

2. The second case deals with sparsity constraints. Here $Q$ might be positive semidefinite or not. The non-convex optimization problem is the following.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2} x^T Q x + c^T x \\
\text{subject to} & Ax = b \\
& Gx \leq h \\
& \|x\|_0 \leq k
\end{array}
\end{equation*}

To convexify this problem, we drop the cardinality constraint and add the Moreau penalty. In practice, the output $\hat{x}(y)$ will not be sparse even after training. We can achieve good results by rounding our solution and taking the top $k$ components in certain cases. However, when we do this, the solution may not be optimal given the support of $x$. We can simply solve a new convex optimization problem zeroing out the entries of $x$ that aren't included in the support. i.e. given the set $\mathcal{I}$ of the top-$k$ support indices of $x$ we would have the following optimization problem.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2}  x^T Q x + c^T x \\
\text{subject to} & Ax = b \\
& Gx \leq h \\
& x_i = 0 \quad \forall i \notin \mathcal{I}
\end{array}
\end{equation*}

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2}  x^T Q x + c^T x + \frac{1}{2 \gamma} ||x-y||_2^2\\
\text{subject to} & Ax = b \\
& Gx \leq h \\
\end{array}
\end{equation*}

Note that all of this optimization problems, we can include a parameter $M \succeq 0$ to reweight the Moreau penalty term.

3. The third case deals with non-convex constraints. Non-convex QCQPs are a very common type of problem that is difficult to solve. Distance geometry problems can be formulated as non-convex QCQP feasibility problems. One typical strategy is to create the SDP relaxation by dropping a rank constraint. Our method proposes to try to learn the optimal objective matrix $C$ that parametrizes the SDP such that the solution fulfills the desired rank constraint. This is likely a separate project. Non-convex QCQPs can be rewritten as rank-constrained SDPs. Here, we could also consider feasibility problems.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \Tr(CX) \\
\text{subject to} & \Tr(A_i X) = b_i \quad  i = 1, \dots m \\
& X \succeq 0 \\
& \text{rank}(X) = k
\end{array}
\end{equation*}

We convexify this problem by dropping the rank constraint and adding the Moreau penalty with $Y$ as a parameter. This method seems like a natural way to convexify constraints.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \Tr(CX) + \frac{1}{2 \gamma} ||X-Y||_2^2\\
\text{subject to} & \Tr(A_i X) = b_i \quad  i = 1, \dots m \\
& X \succeq 0 \\
\end{array}
\end{equation*}

§.§ Learning y

We are given $N$ different training problems that are similar enough that they come from the same family.

The ith problem is parametrized by $\theta_i$. Offline, we solve the original sparse-pca, non-convex problem to get $x^\star_i$. Now we piece together a new neural network architecture. The first model consists of standard neural network layers. The input will be $\theta$ and the output is $y$ (and $M$ if we reshape the Moreau penalty). The last layer is the convex optimization problem parametrized by $y$. The output is $\hat{x}$. The loss is $\sum_{i=1}^{N} ||\hat{x}_i - x^\star_i||_2^2$. The last layer is just an activation function, and it is a continuous function mapping from $y$ to $\hat{x}$. The convex optimization problem is always feasible. Convex QPs can be solved much more quickly than non-convex QPs. Initial results show that we can make improvements on the order of 4 magnitudes or more.

§.§ Our Contribution

* We provide the general framework to use machine learning to solve non-convex optimization problems

* We are the first to propose the penalization of the dual variables for imitation learning

* We prove that implicit gradient descent on the loss function is guaranteed to converge to the local minimum

* We provide generalization bounds for out-of-sample performance.

* TODO: We provide a rigorous definition for what constitutes a family of problems.

* TODO: We showcase that our method works well on real-world applications

§ THEORETICAL RESULTS

§.§ Single Problem Convergence Proof

In this section, we prove for a single problem, that gradient descent for the unconstrained problem, $\min L(y)$ achieves linear convergence to the optimal solution. $y$ is the parameter to the following convex optimization problem.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2} x^T Q x + c^T x + \frac{1}{2 \gamma} ||x-y||_2^2\\
\text{subject to} & Ax = b \\
& Gx \leq h
\end{array}
\end{equation*}

Without loss of generality, our proof of convergence is shown for the problem where $q$ is the parameter. We also assume $Q \succ 0$.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2} x^T Q x + q^T x\\
\text{subject to} & Ax = b \\
& Gx \leq h
\end{array}
\end{equation*}

$$L(q) = \frac{1}{2}\|\hat{x}(q) - x^\star\|_2^2 + \frac{1}{2}\|\hat{\lambda}(q) - \lambda^\star\|_2^2 +  \frac{1}{2}\|\hat{\nu}(q) - \nu^\star\|_2^2$$

where $(\hat{x}(q), \hat{\lambda}(q), \hat{\nu}(q))$ are defined as the primal-dual optimal solution. $(x^\star, \lambda^\star, \nu^\star)$ is meant to represent the optimal solution to the original non-convex problem. What this means is that $x^\star$ is the non-convex solution. $\lambda^\star$ and $\nu^\star$ are not the solutions to the dual of the non-convex problem. Rather they are the dual solution to the convexified primal problem that has solution $x^\star$. i.e. we find some $q^\star$ such that $\hat{x}(q^\star) = x^\star$. Then $\lambda^\star = \hat{\lambda}(q^\star)$ and $\nu^\star = \hat{\nu}(q^\star)$.

Among the points, $q$, that induce non-degenerate solutions, the gradient, $\frac{\partial L}{\partial q}$, is zero only at the globally optimal point.

First, note that there is only one optimal solution to the QP, $(\hat{x}, \hat{\lambda}, \hat{\nu}) = (x^\star, \lambda^\star, \nu^\star)$. We prove the contrapositive, that if $(\hat{x}, \hat{\lambda}, \hat{\nu})$ is not globally optimal then $\frac{\partial L}{\partial q} \neq 0$ for $q$ such that the solution $(\hat{x}, \hat{\lambda}, \hat{\nu})$ is non-degenerate. We label the indices that correspond to tight inequality constraints as $\mathcal{A}$ and the slack inequality constraints as $\mathcal{I}$. That is, $(G \hat{x} - h)_i = 0 \quad \forall i \in \mathcal{A}$ and $(G \hat{x} - h)_i < 0 \quad \forall i \in \mathcal{A}$. Since the solution is non-degenerate, we have $\hat{\lambda}_i > 0 \quad \forall i \in  \mathcal{A}$ and $\hat{\lambda}_i = 0 \quad \forall i \in  \mathcal{I}$. Note that the sets $\mathcal{A}$ and $\mathcal{I}$ refer to the tight solutions for the parametric problem and not the optimal solution.

Since the assumption is that the solution is non-degenerate, the differential matrix is invertible. Thus, we have a unique to the following linear system.

\[\begin{cases}
A d_x = \hat{\nu} - \nu^\star \\
Q d_x + G^T D(\hat{\lambda}) d_{\lambda} + A^T d_{\nu} = \hat{\lambda} - \lambda^\star\\
G d_x + D(G \hat{x} - h) d_{\lambda} = \hat{\lambda} - \lambda^\star
\end{cases}
\]

Case 1: If $\hat{\lambda} = \lambda^\star$ and $\hat{\nu} = \nu^\star$ then, $\hat{x} = x^\star$. This follows from the stationarity condition and the assumption that $Q \succ 0$. We have that $\hat{x}$ is uniquely determined by $\hat{x} = -Q^{-1}(q + A^T \hat{\nu} + G^T \hat{\lambda})$. It must also be true that $x^\star = -Q^{-1}(q + A^T \hat{\nu} + G^T \hat{\lambda})$.

Case 2: If $\hat{\nu} \neq \nu^\star$ then $d_x \neq 0$. This is seen from the first constraint.

Case 3: If $\hat{\nu} = \nu^\star$ and $\hat{\lambda}_{\mathcal{A}} \neq \lambda^\star_{\mathcal{A}}$ then $d_x \neq 0$ from the third constraint. We can rewrite this in block form as

\begin{equation*}
G d_x +
\begin{bmatrix}
(G \hat{x} - h)_\mathcal{I} & 0 \\
0 & 0\\
\end{bmatrix}
\begin{bmatrix}
d_\lambda^{\mathcal{I}} \\
\end{bmatrix} = -
\begin{bmatrix}
\hat{\lambda}_{\mathcal{I}} - \lambda^\star_{\mathcal{I}} \\
\hat{\lambda}_{\mathcal{A}} - \lambda^\star_{\mathcal{A}}
\end{bmatrix}
\end{equation*}

It is clear that $d_x \neq 0$ from the bottom constraint of this block form since we require $G_{\mathcal{A}} d_x = \hat{\lambda}_{\mathcal{A}} - \lambda^\star_{\mathcal{A}}$. Where $G_{\mathcal{A}}$ is the rows of $G$ that correspond to the tight constraints.

Case 4:  If $\hat{\nu} = \nu^\star$ and $\hat{\lambda}_{\mathcal{A}} = \lambda^\star_{\mathcal{A}}$ then $\hat{\lambda}_{\mathcal{I}} = \lambda^\star_{\mathcal{I}} = 0$. Thus we are in case 1.

Any solution $(\hat{x}, \hat{\lambda}, \hat{\nu})$ must fall under one of these cases.

It is easy to see that $\frac{\partial L}{\partial q} = 0$ at $(x^\star, \lambda^\star, \nu^\star)$ since the right hand side is zero and the differential matrix is invertible.

For a small enough step size, $\eta$, when moving from one regime to another or to the boundary between the two, the loss will decrease.

Case 1: we step from one regime to the boundary of that regime and another. In this case, the proof is simple since the by strong convexity and smoothness of the current regime, the loss is guaranteed to decrease for small enough step size. The fact that the solution is degenerate is not relevant here.

Case 2: we step from one regime into another. Here, the update to $q$ is linear, but the update to the solutions $(\hat{x}, \hat{\lambda}, \hat{\nu})$ is composed of two linear pieces. This happens because within a regime, $(\hat{x}, \hat{\lambda}, \hat{\nu})$ are linear in $q$. We denote the relevant points as $q^k$ and $q^{k+1}$ where $q^{k+1} = q^k + \eta \frac{\partial L(q^k)}{\partial q}$. The corresponding solutions are $(\hat{x}^k, \hat{\lambda}^k, \hat{\nu}^k)$ and $(\hat{x}^{k+1}, \hat{\lambda}^{k+1}, \hat{\nu}^{k+1})$. At some point in between the two iterates, the regime switches. We denote this point as $k + \frac{1}{2}$ and write the points as $q^{k + \frac{1}{2}}$, $\hat{x}(k + \frac{1}{2})$, $\hat{}\lambda^{k + \frac{1}{2}}$, $\hat{\nu}^{k + \frac{1}{2}}$. From case 1, we know that $L(q^{k + \frac{1}{2}}) < L(q^k)$.

We claim that $L(q^{k+1}) \leq L(q^{k+\frac{1}{2}})$. Assume that the regime switches on index $i$. Without loss of generality, say that $(\hat{\lambda}(q^k))_i > 0$. Thus by strict complementarity, $(G\hat{x}(q^k) - h)_i = 0$. By definition of the degenerate solution, $(\hat{\lambda}(q^{k+\frac{1}{2}}))_i = 0$ and $(G\hat{x}(q^{k+\frac{1}{2}}) - h)_i = 0$. Once we step further after the degenerate point, the solution will no longer be degenerate. Since the optimal solutions are linear within each regime, it must be the case that $\lambda(q^{k+1}) = 0$. Hence, by strict complementarity, $(G\hat{x}(q^k) - h)_i < 0$. For a very small step size, this essentially frees up this constraint so that $\hat{x}(q^{k+\frac{1}{2}})$ can be closer to $x^\star$. This means that $L(q^{k+1}) \leq L(q^{k+\frac{1}{2}})$.

Thus we've shown that $L(q^{k+1}) < L(q^k)$ for a step size small enough.

At a degenerate point, the solution to the least squares problem is a descent direction on the loss.

For the optimal solution, $x^\star$, to the non-convex problem, gradient descent with small enough step size is guaranteed to solve $L(y)$ with linear convergence.

The loss function $L(y)$ is Lipschitz continuous and differentiable almost everywhere. It is subdifferentiable everywhere. Since $\hat{x}(y)$ is piecewise linear in $y$, the loss $L(y)$ is piecewise quadratic in $y$.

Within regime: Since each quadratic piece is strongly convex and and L-smooth, gradient descent convergences linearly to the local min of a piece of the piecewise quadratic or will transition to a new regime. Claim: The only regime where we converge to the local min of the current quadratic function is the optimal regime. Proof:

Switching regimes: Suppose the point $x$

Approaching zero gradient: We've shown that within a regime, the only point where the gradient is zero is at the globally optimal solution. However, it is possible that the gradient approaches zero at the boundary of the regime. However, in this case, if we pick a constant step size, then we will either reach the degenerate point or will cross the degenerate boundary into a new regime.

Degenerate solutions: The degenerate solutions are never hit with probability one. However, if we happen to be at a degenerate point, then the least squares solution to determine the direction is a descent direction. The least squares solution is a linear combination of the two derivatives $\frac{\partial L}{\partial q}$ when perturbed to be on the 2 regimes.

Corollary: linear convergence will work if $x^\star$ is not optimal for the non-convex problem. The only requirement was that there exists a $q$ such that $x^\star = \hat{x}(q)$. This is important because in many cases, we cannot find the best solution to the non-convex problem.

Prove it works with diagonal penalties? In this case, $\hat{x}(q)$ is no longer piecewise linear. We show an example where the loss is no longer piecewise convex, although convergence can still be achieved

§.§ Multiple problem convergence

1. Linear prediction

Consider the case of multiple problems where we linearly predict $q_i$ from the parameter of the $i$th problem $\theta_i$. So we have $q_i = W \theta_i$. Let $t$ be the size of each $\theta$, i.e. $\theta_i \in \mathbb{R}^t$. The loss function is now $L(W) = \frac{1}{2}\sum_{i=1}^N \|\hat{x}(W \theta_i) - x^\star_i \|_2^2 + \|\hat{\lambda}(W \theta_i) - \lambda^\star_i \|_2^2 + \|\hat{\nu}(W \theta_i) - \nu^\star_i \|_2^2$.

2. Linear neural network

3. Single hidden nonlinear layer

§.§ Generalization Bound

Use the fact that $\hat{x}(y)$ is Lipschitz continuous to bound Rademacher complexity.

§ MISCELLANEOUS PARTS

We could also use the FBE for the regression-based loss. In this case, this eliminates the second order information of the objective. So the parametrized problem becomes the following with $d$ being the parameter.

\begin{equation*}
\begin{array}{ll}
\text{minimize} & \frac{1}{2}  x^T  x + d^T x \\
\text{subject to} & Ax = b \\
& Gx \leq h \\
\end{array}
\end{equation*}