はじめに

Differentiable Convex Optimization Layersについて調べたのでメモ．注意として凸最適問題やそのソルバーには一切精通していないため，お気持ちとライブラリの簡単な使い方だけ説明する．気持ちはどうでもいいから使い方だけという人は後半へ．

Differentiable Convex Optimization Layersとは

NeurIPS2019で発表された論文．凸最適問題を制約条件などを入力として最適解を返す関数と見て，制約条件などに存在するパラメータに関する微分を計算できるような統一的な枠組みを提案． PyTorchやTensorflowをベースとしたライブラリも公開していて，簡単に利用することが可能．

TL;DR

凸最適化問題をパラメータを入力として最適解を返す関数 $\mathcal{S}$ として考える．この時，元の問題を凸錐問題に変える写像と凸錐問題を解くソルバー，得られた凸錐問題の解を元の問題の解に変える写像の3つの成分に分解し，affine-solver-affineの合成関数として表現する．あとは連鎖律により $\mathcal{S}$ の導関数を計算することができるというもの．

凸最適問題

一般に凸最適問題は次のように表現される．

$\displaystyle \text{minimize}\ f_0(x;\theta)\\ \text{subject to}\ f_i(x;\theta)\leq0,\ i=1,\dots,m_1.\ g_i(x;\theta)=0,\ i=1,\dots,m_2$

$x\in\mathbb{R}^n$ は最適化される変数で $\theta\in\mathbb{R}^p$ はパラメータ．関数 $f_i:\mathbb{R}^n\rightarrow\mathbb{R}$ は凸で $g_i:\mathbb{R}^n\rightarrow\mathbb{R}$ はアフィン．解 $x^\ast\in\mathbb{R}^\ast$ は全ての制約を満たす任意の目的関数を最小とするベクトル．上記の凸最適問題はパラメータを入力として最適解を返す一種の関数 $\mathcal{S}:\mathbb{R}^p\rightarrow\mathbb{R}^n$ と言える．さらに今回はパラメータ $\theta$ は $x^\ast$ を最小とする学習パラメータとして考える． Differentiable Convex Optimization Layersは凸最適問題がDPP-compliant programである時， $theta$ に関する $\mathcal{S}$ の導関数を与える．

Disciplined convex programming

DCPは凸最適問題を構成する文法のようなもので，関数，"atom"，それらを合成するためのルールからなる． "atom"はアフィンや，凸，もしくは凹関数のような曲率がよく知られた関数で引数に対する単調性を持つ．合成ルールは以下の凸解析の理論に基づく． $h:\mathbb{R}^k\rightarrow\mathbb{R}$ を凸で， $I_1\subseteq{1,2,\dots,k}$ で順序づけられた引数に対して非減少，同様に $I_2$ に対して非増加な関数とする．また， $g_i:\mathbb{R}^n\rightarrow\mathbb{R}$ は $i\in I_1$ に対して凸で， $i\in I_2$ に対して凹， $i\in(I_1\cap I_2)^c$ に対してアフィンであるとする．すると合成関数 $f(x)=h(g_1(x),g_2(x),\dots,g_k(x))$ は凸となる． DCPは合成関数が上述のように凸である限りatomの合成を許す．全てのDCPは凸最適問題である（ただし逆は成り立たない）．

Cone programs

凸錐問題は次のように表現される(ただしいくつかの等価な表現方法が存在する)．

$\displaystyle \text{minimize}\ c^\top x\\ \text{subject to}\ b-Ax\in\mathcal{K}$

$x\in\mathbb{R}^n$ は変数で $\mathcal{K}\in\mathbb{R}^m$ は空でない閉じた凸錐集合で $A\in\mathbb{R}^{m\times n},b\in\mathbb{R}^m,c\in\mathbb{R}^n$ は"problem data"(要はパラメータということかと)．ここではこの問題が唯一の解を持つと仮定する． Differentiable Convex Optimization Layersは凸最適問題を何かしらのsolverを使って解く必要があるが，ここでは"conic solver"を使う限られた場合を考える． conic solverは凸錐問題を対象としたものでproblem data $(A,b,c)$ を解 $x^\ast$ に写像する関数 $s:\mathbb{R}^{m\times n}\times \mathbb{R}^m\times\mathbb{R}^n\rightarrow\mathbb{R}^n$ を作る．

Differentiating through disciplined convex programs

まず $\theta\in\mathbb{R}^p$ でパラメタライズされた変数 $x\in\mathbb{R}^n$ に関する錐問題のDCPを考える．解への写像は $\mathcal{S}:\mathbb{R}^p\rightarrow\mathbb{R}^n$ として定義することができ，その導関数は $\mathrm{D}^\top\mathcal{S}$ として記述する．さらに $\mathcal{S}$ を $R\circ s\circ C$ の合成関数として考える．ただし，パラメータから錐問題のデータ $(A,b,c)$ への写像 $C$ ，錐問題のソルバー $s$ ，錐問題の解 $\tilde{x}^ast$ から解 $x^ast$ への写像 $R$ とする． $C$ と $R$ がアフィンであればこれはaffine-solver-affine (ASA)の形となる．

よって $\theta$ に関する $\mathcal{S}$ の導関数は連鎖律から次のように計算される．

$\displaystyle \mathrm{D}^\top\mathcal{S}(\theta)=\mathrm{D}^\top C(\theta)\mathrm{D}^\top s(A,b,c)\mathrm{D}^\top R(\tilde{x}^\ast)$

affineの部分は変数の置き換えなどの変換として表現され，普通に微分が可能なので割愛(論文には例が載っているので参照)．問題はソルバーをどのように微分するかで，これはかなり深く研究されているようで理解が追いつかなかった．詳細はAppendix Bに載っているので参照．

Diciplined parametrized programming

DPPは集合関数もしくはatomからDCPによってパラメタライズされた計画( $\mathcal{S}$ のこと)に対する文法である(DPPによって生成された計画をdisciplined parametrized programと呼ぶ)． DCPのようにDPPは凸関数のよく知られた合成定理に基づいていて，disciplined parametrized programに現れるどのような関数も全てaffine，convex，またはconcaveであることを保証する． DCPと違い，DPPは生成された計画がASAであることを保証する．

disciplined parametrized programは以下の形の最適化問題として表現される．

$\displaystyle \text{minimize}\ f_0(x,\theta)\\ \text{subject to}\ f_i(x,\theta)\leq\tilde{f}_i(x,\theta),\ i=1,\dots,m_1,\ g_i(x,\theta)=\tilde{g}_i(x,\theta),\ i=1,\dots,m_2$

変わらず $x\in\mathbb{R}^n$ は変数で $\theta\in\mathbb{R}^p$ はパラメータ． $f_i$ は凸関数で[tex:\tilde{f}i]は凹関数， $g_i$ と[tex:\tilde{g}i]はアフィン．この表現はatomをノード，変数，制約，パラメータを葉とする木として捉えることができ，変数やパラメータなどの状態によってparameter-affine(葉に変数を持たないaffine)，parameter-free(パラメタライズされていない)，variable-free(変数を持たない)などと呼ばれる．

全てのDPP programはDCPであるがその逆は成り立たない．DPPはパラメータを含む表現に以下の二つの制限を加えることでASAの形に還元された計画を生成する．

DCPの中に現れる各部分表現をconvex，concave，affine，または定数に分類する．ただし，パラメータは全て定数として見なされる．DPPの中ではパラメータは変数と同様affineとして分類される．
DCPの中で $x,y$ が定数(variable-free)であれば積 $\phi_\text{prod}(x,y)=xy$ はaffineとなる．DPPのもとでは以下の少なくとも一方が成り立つとき積はaffineとなる．
$x$ または $y$ が定数
どちらかがparameter-affineでもう一方がparameter-free

例として以下の問題を考える．

$\displaystyle \text{minimize}\ \|Fx-g\|_2+\lambda\|x\|_2\\ \text{subject to}\ x\geq 0$

変数は $x\in\mathbb{R}^n$ でパラメータは $F\in\mathbb{R}^{m\times n},g\in\mathbb{R}^m,\lambda\gt 0$ ．以下のように $\|\cdot\|_2$ ，積，差，和がatomであるためこの問題はDPP-compliantとなる．

$F$ がparameter-affine， $x$ がparameter-freeであるため $F,x$ はaffineとなるため， $\phi_{\text{prod}}(F,x)=Fx$ はaffine．
$Fx$ と $-g$ はaffineであり，affineの和はaffineであるため $Fx-g$ はaffine．
$\|\cdot\|_2$ は凸で凸とaffineの合成は凸であるから $\|Fx-g\|_2$ は凸．
$\lambda$ はparameter-affineで $\|x\|_2$ はparameter-freeであるためその積はaffine．さらに $\lambda$ は非負であるため $\|x\|_2$ は増加関数で $\|x\|_2$ は凸となるため， $\phi_{\text{prod}}(\lambda,\|x\|_2)$ は凸．
凸の和は凸であるため目的関数は凸．

Non-DPP transformations of parameters

多くの場合Non-DPP表現はDPP-compliantに書き直すことができる． $p_i$ をパラメータとする以下の例を考える．

$\phi_{\text{prod}}(p_1,p_2)$ は引数がパラメタライズされているためDPPではない．ただし， $p_1p_2$ を変数 $s$ を導入して $p_1s$ とし，制約 $s=p_2$ を加えることでDPPにすることができる．
$\log|p_1|$ は $\log$ は増加凹関数で $|\cdot|$ が凸関数であるためDPPではない．ただし， $|p_1|$ を表現する新たなパラメータ $p_2$ を使って $\log p_2$ と書き直すことが可能．

cvpyx

Differentiable Convex Optimization Layersのライブラリ．pipで入る． githubのリポジトリはこちら．

凸問題の作り方

まずcvxpyをimportする．

import cvxpy as cp

変数はcp.Variable()で定義可能で，引数には変数の次元をとる．同様にパラメータはcp.Parameter()で定義される．例えば，次のように変数とパラメータを定義する．

x = cp.Variable(2)
A = cp.Parameter([2,3])
b = cp.Parameter(3)

注意として，変数，またはパラメータとして定義された(プログラムでいう)変数x,A,bは具体的な値を持っていない．次に目的関数を決める． cp.Minimize()もしくはcp.Maximize()を使って最小化または最大化問題を定義する．例えば $\min_x\|Ax-b\|_2$ という問題を定義するときには次のようにかける．

objective = cp.Minimize(cp.pnorm(A @ x - b, p=2))

基本的な四則演算は通常の演算子で，和などはcvxpyの関数を使って記述する．制約条件はリストにとにかく突っ込んでいく．例として $x\geq0$ という制約条件を考える．

constraints = [x >= 0]

最後に変数やパラメータ，目的関数，制約条件を一つの凸問題として定義する．定義にはcp.Problem()で定義．第一引数には目的関数を，第二引数には制約条件を渡す．

problem = cp.Problem(objective, constraints)
assert problem.is_dpp()

定義した問題がdppかどうかはis_dppで確認できる(戻り値はbool)． Differentiable Convex Optimization Layersは(現状?)DPPしか扱うことができないため，このチェックを抜けられるように問題を頑張って定義する必要がある．

導関数の求め方

cvxpyはPyTorchまたはTensorflowの微分ライブラリを利用して導関数を計算している．自分はPyTorch派なのでここではPyTorchを使った例を考える．まず以下をimport．

import torch
from cvpylayers.torch import CvxpyLayer

Differentiable Convex Optimization LayersをPyTorchのlayerとして定義するには２行目のCvxpyLayerを使う． CvxpyLayer()はcp.Problem()で定義された凸問題とそれに関するパラメータと変数を渡すことで定義できる．

cvxpylayer = CvxpyLayer(problem, paramters=[A, b], variables=[x])

定義されたcvxpylayerはAとbに対応する次元を持つtorch.Tensorを引数にとり，xに対応する次元を持つproblemの解であるtorch.Tensorを返す．

A_tch = torch.randn(3, 2, requires_grad=True)
b_tch = torch.randn(3, requires_grad=True)

solution, = cvxpylayer(A_tch, b_tch)

cvxpylayerの戻り値はリストで返ってくることに注意．ここまでの例はgithubのreadmeに載っているので参照．

さらなる例としてrelu関数を凸最適問題として定義して解くことを考える． reluは $\min_x\|x-y\|_p,\ x\geq 0$ という凸問題として考えることができるため次のように書くことができる．

import cvxpy as cp
import torch 
from cvxpylayers.torch import CvxpyLayer

dim = 10

x = cp.Variable(dim)
y = cp.Parameter(dim)
constraints = [x >= 0]
objective = cp.Minimize(cp.pnorm(x - y, p=2))
problem = cp.Problem(objective, constraints)
assert problem.is_dpp()

cvxpylayer = CvxpyLayer(problem, parameters=[y], variables=[x])
y_tch = torch.randn(dim, requires_grad=True)

# solve the problem
solution, = cvxpylayer(y_tch)

# compute the gradient of the sum of the solution with respect to y
solution.sum().backward()

ちなみにcvxpylayerはバッチ入力に対応しているため，cvxpylayer以降を次のように書き換えることも可能．

batch_size = 5
y_tch = torch.randn(batch_size, dim, requires_grad=True)

# solve the problem
solution, = cvxpylayer(y_tch)

# compute the gradient of the sum of the solution with respect to y
solution.sum().backward()

ただしバッチで入れる場合には引数が全てバッチ次元を持つ必要があるので注意． githubのexamplesの中には入力 $x$ に線形変換を加えてrelu関数を食わせる $z=\max(Wx+b,0)$ を，cvxpyを使ってreluを解きつつPyTorchの枠組みでパラメータW,bをend-to-endでバッチ学習する例も載っているので参照．

まとめ

理論も実装も物凄い完成度．論文では(NeurIPSに投稿するためもあってか)ニューラルネットと組み合わせてend-to-endで学習可能という方向性を見せられているが，ニューラルネットの中間層で凸最適問題を解くというユースケースがあまり思いつかないのが少しなんとも言えないところ．ただ凸最適に明るくなくても簡単に使うことができるためこれからいろんな応用が開けることを期待．