Differentiating Through a Cone Program

Introduction

A cone program is an optimization problem in which the objective is to minimize a linear function over the intersection of a subspace and a convex cone. Cone programs include linear programs, second-order cone programs, and semidefinite programs. Indeed, every convex optimization problem can be expressed as a cone program [Nem07].

Cone programs can be efficiently solved by a number of methods, including interior-point methods [NN94] and operator splitting methods such as the alternating directions method of multipliers [OCPB16]. Cone programs have found applications in many areas, including control [BEFB94], machine learning [HTF09, BPC+11], finance [Mar52, BBD+17], supply chain management [BT04, BTGNV05], and energy management [MBBW19], to name just a few. Cone programs are widely used in conjunction with several popular software systems for convex optimization, which reformulate a convex optimization problem expressed in a domain specific language into an equivalent cone program. These include YALMIP [Löf04], CVX [GB14], CVXPY [DB16a, AVDB18], CVXR [FNB19], and Convex.jl [UMZ+14].

An optimization problem can be viewed as a (possibly multi-valued) function mapping the problem data, i.e., the numerical data defining the problem, to the (primal and dual) solution. This solution map is in general set-valued. In neighborhoods where the solution map is single-valued, it is an implicit function of the problem data. In these neighborhoods it becomes meaningful and interesting to discuss how perturbations in the problem data affect the solution. The point of this paper is to calculate the effects of such perturbations for cone programs, in an efficient way.

The use of implicit differentiation to study the sensitivity of solution mappings of optimization problems dates back several decades, with the works of Fiacco [FM68] and Robinson [Rob80] marking significant milestones. The book [BS00] provides a thorough treatment of the subject, and [DR09] is a good reference for implicit functions more generally. In recent applications, the framework of implicit functions has been used to differentiate through quadratic programs [AK17], stochastic optimization [DAK17], physics simulators [dSA+18], control algorithms [AJS+18], and games [LFK18].

Automatic differentiation (AD).

Contemporary interest in applications of implicit differentiation is partly due to the availability of high-quality, modern, open-source AD software. AD is a family of techniques for algorithmically computing exact derivatives of compositions of differentiable functions. Techniques for AD have been known since at least the 1950s [BKSF59]; see also [Wen64, Gri89, GW08]. In fact, AD is essentially an efficient way of computing the chain rule, which can be traced back to a centuries-old manuscript by Leibniz [Lei76] (see [RF10] for a detailed history).

There are two main variants of AD. Reverse-mode AD computes the derivative of a composition of atomic differentiable functions by computing the sensitivity of an output with respect to the intermediate variables (without materializing the matrices for the intermediate derivatives). In this way, reverse-mode can efficiently compute the derivatives of scalar-valued functions. Forward-mode AD computes the derivative by calculating the sensitivity of the intermediate variables with respect to an input variable [GW08, §2].

Reverse-mode AD was implemented as early as the 1980s [Spe80]. Its rediscovery by the machine learning community (where it is known as backpropagation [RHW88]) and the modern popularity of deep neural networks have led to the development of many software libraries for reverse-mode AD. Examples include TensorFlow [ABC+16, AMP+19], PyTorch [PGC+17], autograd [MDA15], and Zygote [Inn19]. The library JAX [FJL18] supports both reverse-mode and forward-mode AD. This requires representing the derivative of each atomic function as an abstract linear map, i.e., as methods that apply the derivative and its adjoint [DB16b]; implementing just reverse-mode AD only requires the adjoint. Many of the atomic functions included in these libraries are not differentiable at all points in their domains. At non-differentiable points, libraries compute heuristic quantities instead of derivatives. For a discussion on non-differentiability as it relates to AD, see [GW08, §14].

This paper.

In this paper, we give conditions that guarantee the existence of the derivative of the solution map for a cone program, and we provide an expression for the derivative at points where these conditions are satisfied. As in [BMB18], our formulation involves expressing the cone program as the problem of finding a zero of a particular function, specifically the residual map for a homogeneous self-dual embedding of the program [YTM94, OCPB16]. We also show how to efficiently compute the derivative and its adjoint, which involves computing projections onto convex cones, solving a linear system, and exploiting sparsity. In §3, we present an open-source Python package that furnishes the derivative of a cone program as an abstract linear map.

The solution map and its derivative

We consider a (convex) conic optimization problem in its primal (P) and dual (D) forms (see, e.g., [BV04, §4.6.1] or [BTN01, §1.4]):

Here $x\in{\mbox{\bf R}}^{n}$ is the primal variable, $y\in{\mbox{\bf R}}^{m}$ is the dual variable, and $s\in{\mbox{\bf R}}^{m}$ is the primal slack variable. The set $\mathcal{K}\subseteq{\mbox{\bf R}}^{m}$ is a nonempty, closed, convex cone with dual cone $\mathcal{K}^{*}\subseteq{\mbox{\bf R}}^{m}$. The problem data are $A\in{\mbox{\bf R}}^{m\times n}$, $b\in{\mbox{\bf R}}^{m}$, $c\in{\mbox{\bf R}}^{n}$, and the cone $\mathcal{K}$. (In the sequel, however, we will consider the cone as fixed.) In the following we let $N=m+n+1$, and use $\Pi$ to denote the projection onto ${\mbox{\bf R}}^{n}\times\mathcal{K}^{*}\times{\mbox{\bf R}}_{+}$. Finally, we define

Evidently $\mathcal{Q}$ is a proper subspace of the space of $N\times N$ skew symmetric matrices.

We call $(x,y,s)$ a solution of the primal-dual conic program (2) if

For given problem data, the corresponding primal-dual conic program (2) may have no solution, a unique solution, or multiple solutions. We focus on the case when it has a unique solution. We define the solution map $\mathcal{S}:{\mbox{\bf R}}^{m\times n}\times{\mbox{\bf R}}^{m}\times{\mbox{\bf R}}^{n}\to{\mbox{\bf R}}^{n+2m}$ as the function mapping $(A,b,c)$ to vectors $(x,y,s)$ that satisfy (4). We express $\mathcal{S}$ as the composition $\phi\circ s\circ Q$, where

$Q:{\mbox{\bf R}}^{m\times n}\times{\mbox{\bf R}}^{m}\times{\mbox{\bf R}}^{n}\to\mathcal{Q}$ maps the problem data to the corresponding skew-symmetric matrix in $\mathcal{Q}$,

$s:\mathcal{Q}\to{\mbox{\bf R}}^{N}$ furnishes a solution of the homogeneous self-dual embedding, and

$\phi:{\mbox{\bf R}}^{N}\to{\mbox{\bf R}}^{n}\times{\mbox{\bf R}}^{m}\times{\mbox{\bf R}}^{m}$ maps a solution of the homogeneous self-dual embedding to the solution of the primal-dual pair (2).

At a point $(A,b,c)$ where $\mathcal{S}$ is differentiable, the derivative of the solution map is

by the chain rule. In the remainder of this section, we describe the functions $Q$, $s$, and $\phi$ and their derivatives, along with sufficient conditions for their differentiability.

Skew-symmetric mapping.

Homogeneous self-dual embedding.

The homogeneous self-dual embedding of (2) uses the variable $z\in{\mbox{\bf R}}^{N}$. We partition $z$ as $(u,v,w)\in{\mbox{\bf R}}^{n}\times{\mbox{\bf R}}^{m}\times{\mbox{\bf R}}$. The normalized residual map introduced in [BMB18] is the function $\mathcal{N}:{\mbox{\bf R}}^{N}\times\mathcal{Q}\to{\mbox{\bf R}}^{N}$, defined by

For a given $Q$, $z$ can be used to construct the solution of the primal-dual pair (2) if and only if $\mathcal{N}(z,Q)=0$ and $w>0$ [BMB18].

The normalized residual map is an affine function of $Q$, hence its derivative $\mathsf{D}_{Q}\mathcal{N}$ always exists, and is given by

We now turn to $\mathsf{D}_{z}\mathcal{N}(z,Q)$. $\mathcal{N}$ is differentiable at $z$, with $w\neq 0$, whenever $\Pi$ is differentiable at $z$, in which case one can directly verify that

where $e=(0,0,\ldots,1)\in{\mbox{\bf R}}^{N}$. In particular, when $z$ is a solution of the primal-dual pair (2), the second term on the right hand side of (7) vanishes and we have

Implicit function theorem.

Suppose that $z$ is a solution of the primal-dual pair (2) for a given $Q$, and that $\Pi$ is differentiable at $z$. Then $\mathcal{N}$ is differentiable at $z$, $\mathcal{N}(z,Q)=0$ and $w>0$. Now, suppose that $\mathsf{D}_{z}\mathcal{N}(z,Q)$ is invertible. It follows from the implicit function theorem (see, e.g., [Din07] and [DR09]) that there exists a neighborhood $V\subseteq\mathcal{Q}$ of $Q$ on which the solution $z=s(Q)$ of $\mathcal{N}(z,Q)=0$ is unique. Furthermore, $s$ is differentiable on $V$, $\mathcal{N}(s(Q),Q)=0$ for all $Q\in V$, and

Solution construction.

The function $\phi:{\mbox{\bf R}}^{N}\to{\mbox{\bf R}}^{n+2m}$, given by

constructs a solution $(x,y,s)$ of the primal-dual pair (2) from a solution $z=(u,v,w)$ of the homogeneous self-dual embedding ($\Pi_{\mathcal{K}^{*}}$ denotes the projection onto $\mathcal{K}^{*}$). If $\Pi_{\mathcal{K}^{\star}}$ is differentiable with derivative $\mathsf{D}\Pi_{\mathcal{K}^{\star}}(v)\in{\mbox{\bf R}}^{m\times m}$, then $\phi$ is also differentiable, with derivative

Note that a solution of the homogeneous self-dual embedding can also be constructed from a solution of the primal-dual pair [BMB18].

Implementation

In this section we detail how to form the derivative of a cone program as an abstract linear map. We also describe our Python package that implements these methods, and as an example use it to differentiate a semidefinite program.

The matrix $A$ is often stored and manipulated as a sparse matrix. We assume that the sparsity pattern of $A$ is fixed, meaning we only have to consider the nonzero entries of $A$ when computing the derivative and its adjoint; this can provide significant speed-ups if $A$ is very sparse. Of course, one can still furnish the derivative with respect to every entry of $A$ by making $A$ dense.

Computing the derivative.

Applying the derivative $\mathsf{D}\mathcal{S}(A,b,c)$ to a perturbation $(\mathsf{d}A,\mathsf{d}b,\mathsf{d}c)$ corresponds to evaluating

where $\mathsf{d}A$ has the same sparsity pattern as $A$. We work from right to left. The first step is to form

where $M=((Q-I)\mathsf{D}\Pi(z)+I)/w$. One option is to form $M$ as a dense matrix, factorize it, and then back-solve. However, when $M$ is large, as is the case in many applications, this can be impractical. Instead, we suggest using LSQR [PS82] to solve

which only requires multiplication by $M$ and $M^{T}$, and is of particular interest for semidefinite cone programs. Next we partition $\mathsf{d}z$ as $\mathsf{d}z=(\mathsf{d}u,\mathsf{d}v,\mathsf{d}w)$. The final step is to compute

Computing the adjoint of the derivative.

Applying the adjoint of the derivative to a perturbation $(\mathsf{d}x,\mathsf{d}y,\mathsf{d}s)$ corresponds to evaluating

We again work right to left. The first step is to form

Next we form $g=-M^{-T}\mathsf{d}z$, again using LSQR. Then $\mathsf{d}Q$ is given by

Instead of explicitly forming $\mathsf{d}Q$, we only obtain its nonzero entries. Let its nonzero entries be indexed by $\Omega$; we compute

Integration into AD.

The calculations that we have described can be immediately be integrated into the forward and reverse-mode AD software described in §1. In a forward-mode AD system, one calculates the sensitivity of the intermediate variables with respect to perturbations; this operation corresponds to applying the derivative to those perturbations. In a reverse-mode AD system, one computes the sensitivity of a scalar function with respect to intermediate variables; this operation is given by applying the adjoint of the derivative to the derivative of the scalar function with respect to $x$, $y$, and $s$.

1 Python implementation

We provide a Python implementation of the ideas described in the paper, available at

We use the libraries NumPy for dense linear algebra [VDWCV11] and SciPy for sparse linear algebra and its LSQR implementation [JOP+01]. To solve the cone program, we use the numerical solver SCS [OCPB16]. Our implementation supports any cone that can be expressed as the Cartesian product of the zero cone, positive orthant, second-order cone, positive semidefinite cone, exponential cone, and dual exponential cone. Expressions for the derivative of the projection onto each of these cones are given in [AWK17] and [BMB18]; most of these have analytical expressions.

solve_and_derivative(A, b, c, cone_dict),

where A is a SciPy sparse matrix, b and c are NumPy arrays, and cone_dict is a dictionary representing the cone $\mathcal{K}$. This function returns a solution $(x,y,s)$ of the primal-dual pair and two functions, derivative(dA, db, dc) and adjoint_derivative(dx, dy, ds), which respectively apply the derivative and its adjoint (at (A, b, c)) to their inputs and return the result.

2 Example

We consider differentiating a semidefinite program

with optimization variable $X\in{\mbox{\bf S}}_{+}^{n}$ and problem data $C\in{\mbox{\bf S}}^{n}$, $A_{1},\ldots,A_{p}\in{\mbox{\bf S}}^{n}$, and $b\in{\mbox{\bf R}}^{p}$ (${\mbox{\bf S}}^{n}$ denotes the set of symmetric matrices in ${\mbox{\bf R}}^{n\times n}$, and ${\mbox{\bf S}}_{+}^{n}$ denotes the set of symmetric positive semidefinite matrices in ${\mbox{\bf R}}^{n\times n}$). This problem can be readily cast as a standard cone program (2), where the cone is the Cartesian product of a zero cone and a semidefinite cone.

We generated a feasible, bounded random instance of (11) with $p=100$ and $n=300$. We used our Python package to solve this instance and retrieve the derivative and its adjoint. The solve, which calls into SCS, took about 195.7 seconds on a machine with a six-core Intel i7-8700K. Next, we computed the derivative of the optimal value of (11) with respect to each of the $A_{i}$ by applying the adjoint of the derivative to the gradient of the objective. Applying the adjoint of the derivative took about 191.7 seconds, which is roughly equal to the time it took to solve the problem. Note that we calculated the derivative of the objective with respect to 4,515,000 elements (each of the $A_{i}$), which required the solution of a 90,401 $\times$ 90,401 linear system. This evidently would not have been practical had we not treated the linear mapping as an abstract operator.

Acknowledgements

We thank Brandon Amos and Zico Kolter, who concurrently and independently derived calculations similar to the ones in this paper, for many useful discussions. Shane Barratt is supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE-1656518. Walaa Moursi is partially supported by the Natural Science and Engineering Research Council of Canada Postdoctoral Fellowship.

References