Frank-Wolfe with away steps

はじめに

何か学習したいパラメータ $\boldsymbol{x}$ について定義された（例えば二乗誤差）損失関数 $f(\boldsymbol{x})$ を $L_1$ ノルム制約の下で最適化する問題は，次のように定義できます．

$$
\arg\min_{\boldsymbol{x}} f(\boldsymbol{x}) \quad \text{s.t.} \quad ||\boldsymbol{x}||_1 \le D
$$
またこの行列版として，低トレースノルム制約の下での行列の最適化問題

$$
\arg\min_{X} f(X) \quad \text{s.t.} \quad ||X||_{\rm Tr} \le \tau
$$
は，低ランク性の凸緩和として推薦システムなどでは良く扱われ，得られる解が低ランクになる傾向があります．前の記事ではこういった凸最適化問題を解く際に使えるアルゴリズムとして Frank-Wolfe アルゴリズムを紹介しました．

Frank-Wolfe アルゴリズムは毎回の更新に必要な計算が実行可能領域 $\mathcal{K}$ 上での線形最適化とステップ幅の line search のみと軽く，Gradient Descent と同じく繰り返し回数 $k$ について $O(1/k)$ の収束レートを達成しますが，Accerelated Gradient Descent の $O(1/k^2)$ に比べれば遅くなります．本稿では Frank-Wolfe の収束を改善する手法である away step を紹介します．また，トレースノルム制約の下での最適化に適用する際の課題を述べます．

Frank-Wolfe with away steps

解く問題として，

$$
\arg\min_{\boldsymbol{x}} f(\boldsymbol{x}) \quad \text{s.t.} \quad \boldsymbol{x} \in \mathcal{K}:={\rm conv}( \mathcal{A} )
$$
を扱います．ここで， $\mathcal{A}$ は何かの有限集合， ${\rm conv}(\mathcal{A})$ をその凸包とします．$\mathcal{A}$ の凸包とは $\mathcal{A}$ を包む最小の凸集合のことです．例えば $L_1$ ノルム球は第 $i$ 要素だけ 1 のベクトル $\boldsymbol{e}_i$ を用いて，

\{ \boldsymbol{x} \in \mathbb{R}^N : ||\boldsymbol{x} ||_1 \le 1 \} = {\rm conv}( \{ \boldsymbol{e}_1, \boldsymbol{e}_2, \dots, \boldsymbol{e}_N \} )

のように $N$ 個の要素からなる有限集合 $\mathcal{A}$ の凸包として扱えます．

なぜ Frank-Wolfe が遅いのかを視覚的に表した図が [1] の Fig.1 に示されていますが，最適解 $\boldsymbol{x}^\star$ が実行可能領域の境界にあるとき，その近辺でジグザグに動くことがあるようです．Frank-Wolfe の更新方法は"その場で $f$ の一次近似が最も小さくなる方向へ向かう"ですが， away step の考えは"$f$ の一次近似が最大になる向きの反対"に更新することで，ジグザグな動きをすっ飛ばそうというものです．毎回，通常の更新方向 $\boldsymbol{d}_{\rm FW}$ と away step 方向の更新 $\boldsymbol{d}_{\rm AS}$ を計算し，$\boldsymbol{d}_{\rm AS}$ 方向に動いたほうが $f$ が小さくなるならそっちに行く，という処理を追加します．例えば擬似コードは次のようになります（わかりやすさのために [1] の擬似コードをいじったものです）．

\begin{align}
&{\rm input} \; \boldsymbol{x}_0 \in \mathcal{K} \\
&{\rm for} \; k=0,1,2,\dots \\
&\quad \boldsymbol{s}_k = \arg\min_{\boldsymbol{s} \in \mathcal{A}} \; \boldsymbol{s} \cdot \nabla f( \boldsymbol{x}_k ) \quad {\rm and} \quad \boldsymbol{d}_{\rm FW} = \boldsymbol{s}_k - \boldsymbol{x}_k \\
&\quad \boldsymbol{v}_k = \arg\max_{\boldsymbol{v} \in \mathcal{A}} \; \boldsymbol{v} \cdot \nabla f( \boldsymbol{x}_k ) \quad {\rm and} \quad \boldsymbol{d}_{\rm AS} = \boldsymbol{x}_k - \boldsymbol{v}_k \\
&\quad {\rm if} \quad \nabla f(\boldsymbol{x}_k) \cdot \boldsymbol{d}_{\rm FW} \le \nabla f(\boldsymbol{x}_k) \cdot \boldsymbol{d}_{\rm AS} \\
&\quad \quad \boldsymbol{d} = \boldsymbol{d}_{\rm FW} \quad {\rm and} \quad \gamma_\max = 1\\
&\quad {\rm else} \\
&\quad \quad \boldsymbol{d} = \boldsymbol{d}_{\rm AS} \quad {\rm and} \quad \gamma_\max = \arg\max_{\gamma \ge 0} \gamma \quad {\rm s.t.} \quad \boldsymbol{x} + \gamma \boldsymbol{d}_{\rm AS} \in \mathcal{K} \\
&\quad \gamma_k = \arg\min_{\gamma \in [0,\gamma_\max]} f(\boldsymbol{x}_k + \gamma \boldsymbol{d}) \\
&\quad \boldsymbol{x}_{k+1} = \boldsymbol{x}_k + \gamma_k \boldsymbol{d}
\end{align}

ここで $\boldsymbol{s}_k$ と $\boldsymbol{v}_k$ の探索範囲が $\mathcal{K}$ でなく $\mathcal{A}$ になっています．これらは本来 $\mathcal{K}$ 上での内積の最小化（最大化）の解ですが，線形計画問題なのでどっちみち $\mathcal{K}$ の端っこ = $\mathcal{A}$ のどれかの点が解になります．

さてここで $\gamma_k$ の line search は解けるものと仮定したとき，$\boldsymbol{d}_{\rm FW}$ 方向への更新の際，普通の Frank-Wolfe アルゴリズムと同様に $\boldsymbol{x}_k\in\mathcal{K}$ と $\boldsymbol{s}_k\in\mathcal{K}$ の内分点として $\boldsymbol{x}_{k+1} = (1-\gamma_k)\boldsymbol{x}_k + \gamma_k \boldsymbol{s}_k$ と計算することで $\boldsymbol{x}_{k+1}\in\mathcal{K}$ を保証しています．この式からもわかる通り，この場合は $\gamma_k$ の最大値は $\gamma_\max=1$ で問題ありません．一方で，$\boldsymbol{d}_{\rm AS}$ 方向への更新の際は $\boldsymbol{x}_{k+1} = (1+\gamma_k)\boldsymbol{x}_k -\gamma_k \boldsymbol{v}_k$ と内分点の形をしておらず，変な値を $\gamma_\max$ とすると $\boldsymbol{x}_{k+1}$ が $\mathcal{K}$ の外に出て行きます．そこで $\gamma_\max$ の探し方が問題になりますが，$\mathcal{K}$ に依存する話なので一概にこれだ！と決めることが出来ません．また $\gamma_\max$ を探すために重たい計算が必要になっては意味がありません．

[1] ではさらに pairwise away step なるものも提案し，実験ではただの away step より良いと示していますが，$\gamma_{\rm max}$ を何とかして求めないといけないという点は一緒です．

実装の詳細

$\gamma_\max$ を効率的に求めるために $\mathcal{A}$ が有限集合であることを使います．Frank-Wolfe アルゴリズムでは $\boldsymbol{x}_k$ は常に $\mathcal{K}$ に入っている（そうなるように $\gamma_\max$ を決める）ため， $\boldsymbol{x}_0 \in \mathcal{A}$ として端点からスタートすれば， $\sum_{\boldsymbol{v} \in \mathcal{S}_k} \alpha_{\boldsymbol{v}}=1$ となる $\alpha_{\boldsymbol{v}}\ge 0$ を使って $\boldsymbol{x}_k = \sum_{\boldsymbol{v} \in \mathcal{S}_k} \alpha_{\boldsymbol{v}} \boldsymbol{v}$ と解 $\boldsymbol{x}_k$ を端点の凸結合の形で書けます．ここで $\mathcal{S}_k \subseteq \mathcal{A}$ です．

$\mathcal{A}$ は有限集合なので $|\mathcal{A}|$ 次元のベクトル $\boldsymbol{\alpha}$ や集合 $\mathcal{S}_k$ をプログラム中で保持しておくことが出来ます．このことを使って away step を書き換えると [1] に載っている下記アルゴリズムが得られます．

\begin{align}
&{\rm input} \; \boldsymbol{x}_0 = \boldsymbol{v} \in \mathcal{A} \quad {\rm and} \quad {\rm set} \; \alpha_{\boldsymbol{v},0} = 1 \\
&{\rm for} \; k=0,1,2,\dots \\
&\quad \boldsymbol{s}_k = \arg\min_{\boldsymbol{s} \in \mathcal{A}} \; \boldsymbol{s} \cdot \nabla f( \boldsymbol{x}_k ) \quad {\rm and} \quad \boldsymbol{d}_{\rm FW} = \boldsymbol{s}_k - \boldsymbol{x}_k \\
&\quad \boldsymbol{v}_k = \arg\max_{\boldsymbol{v} \in \mathcal{S}_k} \; \boldsymbol{v} \cdot \nabla f( \boldsymbol{x}_k ) \quad {\rm and} \quad \boldsymbol{d}_{\rm AS} = \boldsymbol{x}_k - \boldsymbol{v}_k \\
&\quad {\rm if} \quad \nabla f(\boldsymbol{x}_k) \cdot \boldsymbol{d}_{\rm FW} \le \nabla f(\boldsymbol{x}_k) \cdot \boldsymbol{d}_{\rm AS} \\
&\quad \quad \boldsymbol{d} = \boldsymbol{d}_{\rm FW} \quad {\rm and} \quad \gamma_\max = 1\\
&\quad {\rm else} \\
&\quad \quad \boldsymbol{d} = \boldsymbol{d}_{\rm AS} \quad {\rm and} \quad \gamma_\max = \alpha_{\boldsymbol{v}_k,k}/(1-\alpha_{\boldsymbol{v}_k,k}) \\
&\quad \gamma_k = \arg\min_{\gamma \in [0,\gamma_\max]} f(\boldsymbol{x}_k + \gamma \boldsymbol{d}) \\
&\quad \boldsymbol{x}_{k+1} = \boldsymbol{x}_k + \gamma_k \boldsymbol{d}\\
&\quad {\rm update} \quad \boldsymbol{\alpha}_{k+1} \quad {\rm and} \quad \mathcal{S}_{k+1}
\end{align}

$\boldsymbol{\alpha}_{k+1}$ と $\mathcal{S}_{k+1}$ の計算は場合分けが多いので別の関数に分割していますが，下記のとおりです．なお [1] では see text とあり正直よくわからなかったので，著者の公開している MATLAB 実装を見て補完しました．

\begin{align}
&{\rm if} \quad \nabla f(\boldsymbol{x}_k) \cdot \boldsymbol{d}_{\rm FW} \le \nabla f(\boldsymbol{x}_k) \cdot \boldsymbol{d}_{\rm AS} \\
&\quad {\rm if} \quad \gamma_k = 1 \\
&\quad\quad \mathcal{S}_{k+1} = \{ \boldsymbol{s}_k \} \\
&\quad {\rm else} \\
&\quad\quad \mathcal{S}_{k+1} = \mathcal{S}_{k} \cup \{ \boldsymbol{s}_k \} \\
&\quad \alpha_{\boldsymbol{s}_k,k+1} = (1-\gamma_k) \alpha_{\boldsymbol{s}_k,k} + \gamma_k \\
&\quad \alpha_{\boldsymbol{v},k+1} = (1-\gamma_k) \alpha_{\boldsymbol{v},k} \quad \forall \boldsymbol{v} \in \mathcal{S}_k \setminus \{\boldsymbol{s}_k\} \\
&{\rm else} \\
&\quad {\rm if} \quad \gamma_k = \gamma_\max \\
&\quad\quad \mathcal{S}_{k+1} = \mathcal{S}_k \setminus \{ \boldsymbol{v}_k \} \\
&\quad\quad \alpha_{\boldsymbol{v}_k,k+1} = 0 \\
&\quad\quad \alpha_{\boldsymbol{v},k+1} = (1+\gamma_k) \alpha_{\boldsymbol{v},k} \quad \forall \boldsymbol{v} \in \mathcal{S}_k \setminus \{ \boldsymbol{v}_k \}\\
&\quad {\rm else} \\
&\quad\quad \mathcal{S}_{k+1} = \mathcal{S}_{k} \\
&\quad\quad \alpha_{\boldsymbol{v}_k,k+1} = (1+\gamma_k) \alpha_{\boldsymbol{v}_k,k} - \gamma_k \\ 
&\quad\quad \alpha_{\boldsymbol{v},k+1} = (1+\gamma_k) \alpha_{\boldsymbol{v}_k,k} \quad \forall \boldsymbol{v} \in \mathcal{S}_k \setminus \{ \boldsymbol{v} \}
\end{align}

ごりごりと場合分けしていますが，要するに $\boldsymbol{x}_k$ の更新に合わせて $\boldsymbol{x}_k = \sum_{\boldsymbol{v} \in \mathcal{S}_k} \alpha_{\boldsymbol{v},k} \boldsymbol{v}$ となるように更新しています．なので，$\boldsymbol{v}_k \in \mathcal{S}_k \iff \alpha_{\boldsymbol{v}_k} > 0$ となります．さて，通常の Frank-Wolfe な更新をしたときは $\boldsymbol{x}_k\in \mathcal{K}$ である限り $\boldsymbol{x}_{k+1}\in\mathcal{K}$ は保証されますが，では away step による更新をしたとき $\boldsymbol{x}_{k+1} \in \mathcal{K}$ は保証されるのでしょうか．away step 更新したときの $\boldsymbol{x}_{k+1}$ を書き下すと，次式のようになります．

\begin{align}
\boldsymbol{x}_{k+1} =
(1+\gamma_k)\boldsymbol{x}_k - \gamma_k \boldsymbol{v}_k
&= (1+\gamma_k)\sum_{\boldsymbol{w} \in \mathcal{S}_k \setminus \boldsymbol{v}_k} \alpha_{\boldsymbol{w},k} \boldsymbol{w} +
(1+\gamma_k)\alpha_{\boldsymbol{v}_k,k} \boldsymbol{v}_k -
\gamma_k \boldsymbol{v}_k \\
&=(1+\gamma_k)\sum_{\boldsymbol{w} \in \mathcal{S}_k \setminus \boldsymbol{v}_k} \alpha_{\boldsymbol{w},k} \boldsymbol{w} +
(\alpha_{\boldsymbol{v}_k,k}+\gamma_k \alpha_{\boldsymbol{v}_k,k}-\gamma_k) \boldsymbol{v}_k
\end{align}

係数の和を取ると，

\begin{align}
(1+\gamma_k)\sum_{\boldsymbol{w} \in \mathcal{S}_k \setminus \boldsymbol{v}_k} \alpha_{\boldsymbol{w},k}+
(\alpha_{\boldsymbol{v}_k,k}+\gamma_k \alpha_{\boldsymbol{v}_k,k}-\gamma_k)
&=
(1+\gamma_k)(1-\alpha_{\boldsymbol{v}_k,k})+
(\alpha_{\boldsymbol{v}_k,k}+\gamma_k \alpha_{\boldsymbol{v}_k,k}-\gamma_k) \\
&=
1+\gamma_k - \alpha_{\boldsymbol{v}_k,k} - \gamma_k\alpha_{\boldsymbol{v}_k,k} +
\alpha_{\boldsymbol{v}_k,k}+\gamma_k \alpha_{\boldsymbol{v}_k,k}-\gamma_k =1
\end{align}

あとは $0\le\alpha_{\boldsymbol{v}_k,k}+\gamma_k\alpha_{\boldsymbol{v}_k,k}-\gamma_k$ なら凸結合として成り立ちますが，これは $\gamma_k \le \alpha_{\boldsymbol{v}_k,k}/(1-\alpha_{\boldsymbol{v}_k,k})$ より保証されます．
よって $\boldsymbol{x}_{k+1}$ は $\mathcal{A}$ の要素の凸結合として書け，${\rm conv}(\mathcal{K})$ に入っていると確認できます．

性能の理論保証

このアルゴリズムの収束性の理論保証ですが，[1] の定理 1 によると，$f$ が $\mu$-強凸，$\nabla f$ が $L$-Lipschitz， $\forall \boldsymbol{x},\boldsymbol{y} \in {\rm conv}(\mathcal{A}),\;||\boldsymbol{x}-\boldsymbol{y}||_2 \le M$ のとき Pyramidal width $\delta$ という $\mathcal{K}$ のパラメータを用いて

$$
f(\boldsymbol{x}_{k+1})-f(\boldsymbol{x}^\star) \le O\left(\exp(-k \delta^2 \mu/4L M^2)\right)$$
となり，なんと $k$ について指数的に減っていきます．また，強凸性が言えないときは $O(1/k)$ と，もとの Frank-Wolfe と変わらないようです．Pyramidal width は定義が複雑で，まだ直感的に理解できていないため詳細は省略しますが， $N$ 次元の $L_1$ ノルム球や確率単体

\{\boldsymbol{x}\in \mathbb{R}^N : ||\boldsymbol{x}||_1=1 \}

では $1/\sqrt{N}$ になるとのことです [1,2]．

トレースノルム正則化に使えない理由

トレースノルム球は下記の通りランク 1 行列の凸包であることが知られています [4]．

\{ X \in \mathbb{R}^{N\times M} : ||X||_{\rm Tr} \le \tau \} = {\rm conv}( \{ \tau \boldsymbol{u} \boldsymbol{v}^\top : ||\boldsymbol{u}||_2 = ||\boldsymbol{v}||_2 = 1 \} )

本稿では急に有限集合とか言い出したのでお察しかと思いますが，これは有限集合の凸包ではなく，$\gamma_{\rm max}$ の計算が論文通りに効率的にできない（ $\boldsymbol{\alpha}_k$ と $\mathcal{S}_k$ を持てない）ため破綻します．

なんとかして

\gamma_\max = \arg\max \; \gamma \quad {\rm s.t.} \quad ||X+\gamma D_{\rm AS} ||_{\rm Tr} \le \tau 

を重たい SVD を避けてお手軽に求める方法は無いか......となるのですが，正確に解く方法はちょっと思いつきません．近似であれば，仮に $X_k$ が低ランクなら各 $k$ での SVD は速くできるため，SVD して $||X_k||_{\rm Tr}$ を計算し，$\gamma$ の自明な上界である $\gamma_\max = (||X_k||_{\rm Tr} - \tau)/(||X_k||_{\rm Tr} + \tau)$ を使うことが考えられます．ただこれは近似が荒すぎたのか，手元で実験した際には away step ですぐに動かなくなり，最適化が止まる（ $\gamma_\max = 0$ になり $\gamma_k = 0$ ）という現象が起きることがありました．もしも近似で $\gamma_\max$ を決めるなら，$\gamma$ を二分探索し，その際は [7] のように SVD よりは簡単な計算で得られる $|| X+\gamma D_{\rm AS} ||_{\rm Tr}$ の上界を使う......など，さらに凝った手法が必要になるかと思います．

実は [3] で In-Face Frank Wolfe という away step をトレースノルム球に適用して，さらに低ランクな解を導くよう拡張されたものを示しています．In-Face Frank Wolfe は $O(1/k)$ の収束レートに関する理論保証があり（実際はハイパーパラメータが関わりもう少し複雑です），実験的には Frank Wolfe より数倍～数十倍速くなっています．肝になるアイデアは，(i)変数を $X_k = U_k S_k V_k^\top$ と最初から SVD しておき，実際は特異ベクトル $U_k,V_k$ と特異値の並んだ対角行列 $S_k$ を更新すること，(ii) $X_k$ の SVD が分かっているなら $X_{k+1} = X_k + \tau \boldsymbol{u}\boldsymbol{v}^\top$ の\ SVD は理論上は速く計算できるということを使います．しかし上記の $\gamma_\max$ 問題は部分的にしか解決しておらず，場合によっては二分探索で $\gamma_\max$ を近似すると書いてあります．また，実装もかなり複雑になり， rank-one SVD update という別のテクニックが必要になります．

おわりに

本稿では Frank-Wolfe をスピードアップしたいと思い [1] を紹介しましたが，トレースノルム制約に適用するにはなかなか上手くいかないという話をしました． away-step 以外にも Frank-Wolfe を速くしよう！という話はあり，目的関数が強凸かつ平滑で上位 $r$ 個の特異値について SVD したら [5]...目的関数が強凸で実行可能領域が強凸だったら...関数が強凸かつ最適解が実行可能領域の縁じゃくて内部にあることが分かっていたら...などがあるようです（[6] によくまとめられてる）．本稿で紹介したものも含め，いずれも目的関数が Lipschitz だ，とか強凸だ，とかの性質が必要です．さらにその関数の性質に関わるパラメータ（の上界）がアルゴリズム中でステップ幅などの計算に用いられるため，事前にわかってないといけないのが難点です．

参考文献

[1]On the Global Linear Convergence of Frank-Wolfe Optimization Variants, Simon Lacoste-Julien and Martin Jaggi, 2015

[2]Recent Advances in Frank-Wolfe Optimization
, Simon Lacoste-Julien, 2017

[3]An Extended Frank-Wolfe Method with "In-Face" Directions, and its Application to Low-Rank Matrix Completion, Robert M. Freund, Paul Grigas and Rahul Mazumder, 2015

[4]Rank, Trace-Norm and Max-Norm, Nathan Srebro and Adi Shraibman, 2005

[5]Linear Convergence of a Frank-Wolfe Type Algorithm over Trace-Norm Balls, Zeyuan Allen-Zhu, Elad Hazan, Wei Hu and Yuanzhi Li, 2017

[6]Faster Rates for the Frank-Wolfe Method over Strongly-Convex Sets, Dan Garber and Elad Hazan, 2015

[7]Some simple estimates for singular values of a matrix, Liqun Qi, 1984