此论文通过隐式方法将优化问题表达为神经网络中的一层,它的内部不是传统意义上的神经元堆叠,不是包含全连接层、激活函数等结构,而是一个特殊的“可学习的可微优化层”

通过理论推导得到它的前向传播和反向传播的公式表达,并不需要定义它的内部结构

2. 相关研究

使用神经网络解决受限类别的优化问题的整体思路一般可以分为四种

基于能量的学习方法(Energy-based learning methods)

总体思路为:在训练过程中将能量函数在观测数据流形附近调低,而在其他地方调高。缺点:一些问题中可能没有参测数据;可能会出现不稳定问题

解析法(Analytically)

如果能找到优化问题的解析解,那么梯度一般也可以解析地计算出来。缺点:大多数问题没有解析解

展开(Unrolling)

将优化问题的迭代求解过程近似展开为神经网络的结构。缺点:展开操作会增加网络的复杂性和深度;对于有约束的问题可能难以展开

最小化操作求导(Argmin differentiation)

类似本文的操作,对argmin问题进行微分

3. OptNet:在神经网络中解决优化问题

本文重点研究二次规划问题

min⁡z 12z⊤Qz+q⊤zs.t. Az=b,Gz≤h\begin{align*} \min _z \ & \frac{1}{2} z^\top Qz + q^\top z \\ \text{s.t.} \ & Az = b,Gz \leq h \end{align*}

其中z∈Rnz \in \mathbb{R} ^n是优化变量,Q∈Rn,n⪰0Q \in \mathbb{R} ^{n,n} \succeq 0,q∈Rnq \in \mathbb{R} ^n,A∈Rm,nA \in \mathbb{R} ^{m,n},b∈Rmb \in \mathbb{R} ^m,G∈Rp,nG \in \mathbb{R} ^{p,n},h∈Rph \in \mathbb{R} ^p

3.1 反向传播

拉格朗日函数可以写为

L(z,ν,λ)=12z⊤Qz+q⊤z+ν⊤(Az−b)+λ⊤(Gz−h)L(z,\nu,\lambda) = \frac{1}{2} z^\top Qz + q^\top z+ \nu ^\top (Az - b)+ \lambda ^\top (Gz-h)

其中ν\nu和λ≥0\lambda \geq 0是对偶变量。可以写出它的KKT条件(分别为平稳性条件、原始可行性和互补松弛性)

Qz∗+q+A⊤ν∗+G⊤λ∗=0Az∗−b=0D(λ∗)(Gz∗−h)=0\begin{gather*} Qz^* + q + A^\top \nu ^* + G^\top \lambda^* =0 \\ Az^*-b=0 \\ D(\lambda ^*)(Gz^*-h) = 0 \end{gather*}

其中D(⋅)D(\cdot)表示由向量展开为对角矩阵,z∗,ν∗,λ∗z^*,\nu ^*,\lambda ^*代表最优原始变量和最优对偶变量

对这些条件求微分得到

ΔQz∗+QΔz∗+Δq+ΔA⊤ν∗+A⊤Δν∗+ΔG⊤λ∗+G⊤Δλ∗=0ΔAz∗+AΔz∗−Δb=0D(Gz∗−h)Δλ∗+D(λ∗)(ΔGz∗+GΔz∗−Δh)=0\begin{gather*} \Delta Q z^* + Q \Delta z^* + \Delta q + \Delta A^\top \nu ^* + A^\top \Delta \nu^* + \Delta G^\top \lambda ^* + G^\top \Delta \lambda ^* =0 \\ \Delta A z^* + A \Delta z^* - \Delta b =0 \\ D(Gz^*-h)\Delta \lambda ^* +D(\lambda ^*)(\Delta G z^* + G \Delta z^* - \Delta h)= 0 \end{gather*}

写成矩阵形式为

[QG⊤A⊤D(λ∗)GD(Gz∗−h)0A00][Δz∗Δλ∗Δν∗]=[−Δq−ΔA⊤ν∗−ΔG⊤λ∗−ΔQz∗−D(λ∗)ΔGz∗+D(λ∗)Δh−ΔAz∗+Δb]\begin{bmatrix} Q && G^\top && A^\top \\ D(\lambda ^*)G && D(Gz ^*-h) && 0 \\ A && 0 && 0 \end{bmatrix} \begin{bmatrix} \Delta z^* \\ \Delta \lambda^* \\ \Delta \nu^* \\ \end{bmatrix} = \begin{bmatrix} - \Delta q - \Delta A^\top \nu ^* - \Delta G^\top \lambda ^* -\Delta Q z^* \\ -D(\lambda ^*)\Delta G z^* + D(\lambda ^*) \Delta h \\ - \Delta A z^* + \Delta b \end{bmatrix}

KKT条件是优化问题的必要条件,即如果x∗x^∗是最优解,则它必须满足KKT条件

对于凸优化问题,KKT条件是充要条件,即满足KKT条件的解一定是最优解

这一层的输入为ziz_i,部分参数是由ziz_i决定的,优化变量为zz,输出为z∗z^*,参数为q,b,h,Q,A,Gq,b,h,Q,A,G。想要进行反向传播,需要计算得到损失对于这些参数的梯度

式(7)的推导

将KKT条件的矩阵形式写为紧凑形式

K[Δz∗Δλ∗Δν∗]=rK \begin{bmatrix} \Delta z^* \\ \Delta \lambda^* \\ \Delta \nu^* \\ \end{bmatrix} =r

令上一层传回来的梯度为

g=∂ℓ∂z∗g= \frac{\partial \ell}{\partial z^*}

那么微分可以写为

Δℓ=g⊤Δz∗\Delta \ell = g^\top \Delta z^*

定义

e=[g00]e = \begin{bmatrix} g \\ 0 \\ 0 \end{bmatrix}

那么

Δℓ=e⊤[Δz∗Δλ∗Δν∗]\Delta \ell = e^\top \begin{bmatrix} \Delta z^* \\ \Delta \lambda^* \\ \Delta \nu^* \\ \end{bmatrix}

结合紧凑版KKT条件,上式可以进一步写为

Δℓ=e⊤K−1r\Delta \ell = e^\top K^{-1} r

为了避免显式求解e⊤K−1e^\top K^{-1},引入变量

d=[dzdλdν]d = \begin{bmatrix} d_z \\ d_\lambda \\ d_\nu \\ \end{bmatrix}

它满足

K⊤d=−eK^\top d = -e

这里K⊤K^\top为

K⊤=[QG⊤D(λ∗)A⊤GD(Gz∗−h)0A00]K^\top = \begin{bmatrix} Q && G^\top D(\lambda ^*) && A^\top \\ G && D(Gz ^*-h) && 0 \\ A && 0 && 0 \end{bmatrix}

展开d=(K⊤)−1(−e)d= (K^\top)^{-1} (-e)就可以得到论文中式(7)的形式

式(8)的推导

因为用dd取代了(K⊤)−1(−e)(K^\top)^{-1} (-e),因此Δℓ\Delta \ell可以表示为

Δℓ=−d⊤r\Delta \ell = -d^\top r

展开为

Δℓ=dzTΔQz∗+dzTΔq+dzTΔGTλ∗+dzTΔATν∗+dλTD(λ∗)ΔGz∗−dλTD(λ∗)Δh+dνTΔAz∗−dνTΔb\begin{aligned} \Delta \ell &= d_z^T \Delta Q z^* + d_z^T \Delta q + d_z^T \Delta G^T \lambda^* + d_z^T \Delta A^T \nu^* \\ &\quad + d_\lambda^T D(\lambda^*) \Delta G z^* - d_\lambda^T D(\lambda^*) \Delta h \\ &\quad + d_\nu^T \Delta A z^* - d_\nu^T \Delta b \end{aligned}

下面逐项拆解梯度

对qq的梯度

含Δq\Delta q的项为

dzTΔqd_z^T \Delta q

因此

∂ℓ∂q=dz\frac{\partial \ell }{\partial q} = d_z

对bb的梯度

含Δb\Delta b的项为

−dνTΔb- d_\nu^T \Delta b

因此

∂ℓ∂b=−dν\frac{\partial \ell }{\partial b} = -d_{\nu}

对hh的梯度

含Δh\Delta h的项为

−dλTD(λ∗)Δh- d_\lambda^T D(\lambda^*) \Delta h

因为 D(λ⋆)D(\lambda^\star) 是对角矩阵,所以

−dλTD(λ∗)Δh=−(D(λ∗)dλ)⊤Δh- d_\lambda^T D(\lambda^*) \Delta h = -(D(\lambda^*)d_\lambda)^\top \Delta h

因此

∂ℓ∂h=−D(λ∗)dλ\frac{\partial \ell }{\partial h} = -D(\lambda^*)d_\lambda

对QQ的梯度

含ΔQ\Delta Q的项为

dzTΔQz∗d_z^T \Delta Q z^*

利用迹运算

dzTΔQz∗=tr⁡(dzTΔQz∗)=tr⁡(z∗dzTΔQ)d_z^T \Delta Q z^* = \operatorname{tr}(d_z^T \Delta Q z^*) = \operatorname{tr}(z^* d_z^T \Delta Q )

因此正常应该为

∂ℓ∂Q=dzz∗T\frac{\partial \ell }{\partial Q} = d_z z^{*T}

但为了结果保持对称性,梯度写为

∂ℓ∂Q=12(dzz∗T+z∗dzT)\frac{\partial \ell }{\partial Q} = \frac{1}{2} (d_z z^{*T} + z^* d_z^T)

对AA的梯度

含ΔA\Delta A的项为

dzTΔATν∗+dνTΔAz∗d_z^T \Delta A^T \nu^* + d_\nu^T \Delta A z^*

因此

∂ℓ∂A=dνz∗T+ν∗dzT\frac{\partial \ell }{\partial A} =d_\nu z^{*T} + \nu^* d_z^T

对GG的梯度

含ΔG\Delta G的项为

dzTΔGTλ∗+dλTD(λ∗)ΔGz∗d_z^T \Delta G^T \lambda^* + d_\lambda^T D(\lambda^*) \Delta G z^*

因此

∂ℓ∂G=D(λ∗)dλz∗T+λ∗dzT\frac{\partial \ell }{\partial G} =D(\lambda^*)d_\lambda z^{*T} + \lambda^* d_z^T

这样便得到了式(8)的形式

这样不需要显示求解显式求解e⊤K−1e^\top K^{-1},而是通过求解

K⊤d=−eK^\top d = -e

得到dd的值,继而求得反向传播时的参数梯度

3.2 正向传播

由于传统求解器Gurobi和CPLEX只能运行在cpu上,在进行批量运算时耗时较长,作者根据原始-对偶内点法开发了能跑在gpu上的求解器qpth,具体公式推导不再展开

值得一提的是,在正向传播时会顺便把反向传播时需要的d求出,所以计算耗时主要集中在正向传播,反向传播开销极低