【Python】リーマン多様体上の共役勾配法の動向 (2021)

リーマン多様体上の無制約最適化問題に対する共役勾配法によるアプローチのちょっとした歴史 (〜2020)．
また，その Python での実装．

# 対象とする問題
リーマン多様体 $\mathcal{M}$ 上の無制約最適化問題

$$\mathop{\text{minimize}}_{x\in\mathcal M}\quad f(x)$$

を考える．ここで，$f:\mathcal M \to \mathbb R$ は滑らかな写像とする．

# リーマン多様体上の共役勾配法 [[Absil et al. 2009](https://press.princeton.edu/absil)]
リーマン多様体 $\mathcal M$ 上の共役勾配法のアルゴリズムは， Absil らの多様体上の最適化の書籍 "[Optimization Algorithms on Matrix Manifolds](https://press.princeton.edu/absil)" にて紹介されている．

1. $x^0$ を初期化する．
1. $d^0=−\mathop{\mathrm{grad}} f(x^0)$ で探索方向を初期化する．
1. for $k=0,1,\ldots$ do
  * ステップ幅 $t_k$ を計算する．
  * $x^{k+1}=R_{x^k}\left(t_kd^k\right)$ で更新する．
  * $\beta_k=\frac{\left\langle\mathop{\mathrm{grad}}f(x^{k+1}),\mathop{\mathrm{grad}}f(x^{k+1})\right\rangle_{{x^{k+1}}}}{\left\langle\mathop{\mathrm{grad}}f(x^{k}),\mathop{\mathrm{grad}}f(x^{k})\right\rangle_{{x^{k}}}}$ を計算する．
  * $d^{k+1}=−\mathop{\mathrm{grad}} f(x^{k+1})+\beta_k{\mathcal T_{t_kd^k}}(d^k)$ で探索方向を更新する．
1. end for

この手法については，以下の過去記事で詳細に述べている．

* [線型方程式と共役勾配法 - Qiita](https://qiita.com/Mrrmm252/items/a8366c3a57879d9d60be)
* [【Python】リーマン多様体上の共役勾配法 - Qiita](https://qiita.com/Mrrmm252/items/d48a2eee45d046b707ee)

# リーマン多様体上の共役勾配法 [[Sato & Iwai 2015](https://arxiv.org/pdf/1302.0125.pdf)]
上記で紹介したリーマン多様体上の共役勾配法は，ベクトル輸送が非増加

```math
\left\|\mathcal T_{t_kd^k}(d^k)\right\|_{x^{k+1}} \leq \left\|d^k\right\|_{x^k}
```

である場合に収束することが証明されている [Ring \& Wirth 2012]．この条件を常に満たすように次の係数 $s_k$ を導入した共役勾配法が提案された [Sato & Iwai 2015]．

```math
s_k = \min\left\{\frac{\left\|d^k\right\|_{x^k}}{\left\|\mathcal T_{t_kd^k}(d^k)\right\|_{x^{k+1}}},1\right\}
```

1. $x^0$ を初期化する．
1. $d^0=−\mathop{\mathrm{grad}} f(x^0)$ で探索方向を初期化する．
1. for $k=0,1,\ldots$ do
  * ステップ幅 $t_k$ を計算する．
  * $x^{k+1}=R_{x^k}\left(t_kd^k\right)$ で更新する．
  * $s_k$ を計算する．
  * $\beta_k=\frac{\left\langle\mathop{\mathrm{grad}}f(x^{k+1}),\mathop{\mathrm{grad}}f(x^{k+1})\right\rangle_{x^{k+1}}}{\left\langle\mathop{\mathrm{grad}}f(x^{k}),\mathop{\mathrm{grad}}f(x^{k})\right\rangle_{x^{k}}}$ を計算する．
  * $d^{k+1}=−\mathop{\mathrm{grad}} f(x^{k+1})+\beta_ks_k\mathcal T_{t_kd^k}(d^k)$ で探索方向を更新する．
1. end for

# リーマン多様体上の共役勾配法 [[Zhu & Sato 2020](http://www.optimization-online.org/DB_FILE/2020/05/7798.pdf)]
$\mathbb R^n$ 上での共役勾配法の変数の更新式は $x^{k} = x^{k+1} - t_kd^k$ と変形できる．
この式をあるレトラクション $R^\text{bw}$ を用いてリーマン多様体上へ拡張すると，

$$x^k = R^\text{bw}_{x^{k+1}}(-t_kd^k)$$

とかける．これより，$d^k\in T_{x^k}\mathcal M$ の $T_{x^{k+1}}\mathcal M$ 上の自然な表現として，

```math
d^k \approx -t_k^{-1}{R^\text{bw}_{x^{k+1}}}^{-1}(x^k)
```

が得られる．このとき，ベクトル輸送の非増加条件をみたすための係数 $s_k$ は，

```math
s_k = \min\left\{\frac{\left\|t_kd^k\right\|_{x^k}}{\left\|{R^\text{bw}_{x^{k+1}}}^{-1}(x^k)\right\|_{x^{k+1}}},1\right\}
```

となる．以上より，2つのレトラクション $R^\text{fw}$, $R^\text{bw}$ を用いた次の共役勾配法が得られる．

1. $x^0$ を初期化する．
1. $d^0=−\mathop{\mathrm{grad}} f(x^0)$ で探索方向を初期化する．
1. for $k=0,1,\ldots$ do
  * ステップ幅 $t_k$ を計算する．
  * $x^{k+1}=R^\text{fw}_{x^k}\left(t_kd^k\right)$ で更新する．
  * $s_k$ を計算する．
  * $\beta_k=\frac{\left\langle\mathop{\mathrm{grad}}f(x^{k+1}),\mathop{\mathrm{grad}}f(x^{k+1})\right\rangle_{x^{k+1}}}{\left\langle\mathop{\mathrm{grad}}f(x^{k}),\mathop{\mathrm{grad}}f(x^{k})\right\rangle_{x^{k}}}$ を計算する．
  * $d^{k+1}=−\mathop{\mathrm{grad}} f(x^{k+1})-\beta_ks_kt_k^{-1}{R^\text{bw}_{x^{k+1}}}^{-1}(x^k)$ で探索方向を更新する．
1. end for

# Python での実装

## 使用例

### 問題
次の Brockett cost function の最小化問題を考える．

```math
\begin{align}
\mathop{\mathrm{minimize}}_{X\in \mathbb R^{n \times r}} &\quad f(X) :=\mathop{\mathrm{tr}}\left(X^\top A X N\right)\\
\text{subject to} &\quad X^\top X = I_r
\end{align}
```

ここで，$A\in\mathbb{R}^{n\times n}$ は対称行列，$N\in\mathbb{R}^{r\times r}$ は正値の対角行列 $N = \mathop{\mathrm{diag}}(\mu_1, \mu_2, \ldots, \mu_r)$ で，$0<\mu_1\leq\mu_2\leq\dots\leq\mu_r$ をみたす．最適解 $X_*$ は $A$ の小さい順から $r$ 個の固有値に対応する固有ベクトルを並べた行列となっている．また，ユークリッド空間における $f$ の勾配は，

$$\nabla f(X) = 2AXN$$

で計算される．

この最適化問題は，シュティーフェル多様体

```math
\mathop{\mathrm{St}}(n, r) := \left\{X\in\mathbb{R}^{n\times r} \mid X^\top X = I_r\right\}
```

上の無制約最適化問題

```math
\mathop{\mathrm{minimize}}_{X \in \mathop{\mathrm{St}}(n, r)} \quad f(X)
```

へと変換することができる．このとき，目的関数 $f$ のシュティーフェル多様体上での勾配 $\mathop{\mathrm{grad}}f(X)$ は，点 $X$ における接空間 $T_X \mathop{\mathrm{St}}(n, r)$ への直交射影 $\mathrm{P}_{T_X \mathop{\mathrm{St}}(n, r)}$ を用いて，

```math
\mathop{\mathrm{grad}}f(X) = \mathrm{P}_{T_X \mathop{\mathrm{St}}(n, r)}(\nabla f(X))
```

で計算できることが知られている．ここで，接空間 $T_X \mathop{\mathrm{St}}(n, r)$ は

```math
T_X \mathop{\mathrm{St}}(n, r) = \left\{U\in\mathbb R^{n\times r}\mid X^\top U + U^\top X = O\right\}
```

であるので，直交射影 $\mathrm{P}_{T_X \mathop{\mathrm{St}}(n, r)}$ は

```math
\mathrm{P}_{T_X \mathop{\mathrm{St}}(n, r)}(U) = U - \frac12 X \left(X^\top U + U^\top X\right) = U - X\mathop{\mathrm{sym}}\left(X^\top U\right)
```

となる．これより，目的関数 $f$ のシュティーフェル多様体上での勾配 $\mathop{\mathrm{grad}}f(X)$ は，

```math
\mathop{\mathrm{grad}} f(X) = 2AXN - XX^\top AXN - XNX^\top AX
```

で計算できる．詳細は "[Optimization Algorithms on Matrix Manifolds](https://press.princeton.edu/absil)" の 4.8 節にある．

### シュティーフェル多様体上の演算

* 内積
  
  ```math
  \left\langle U, V\right\rangle_X = \mathop{\mathrm{tr}}\left(U^\top V\right)
  ```

* レトラクション
  
  ```math
  \begin{align}
  R^{\mathrm{qr}}_X(U) &= \mathop{\mathrm{qf}}(X+U)\\
  R^{\mathrm{ct}}_X(U) &= \left(I_n - \frac12 \left(P_X UX^\top - XU^\top P_X\right)\right)^{-1} \left(I_n + \frac12 \left(P_X UX^\top - XU^\top P_X\right)\right) X
  \end{align}
  ```  
  
  ここで，$\mathop{\mathrm{qf}}(X)$ は $X$ の QR 分解 $X = \mathcal{Q}\mathcal{R}$ の $\mathcal Q$ 行列を返す関数 ($\mathcal R$ の対角成分は正)，$P_X = I_n - XX^\top$ である．
  $R^\mathrm{ct}_X(U)$ は逆写像の計算が可能であり，

  ```math
  {R^\mathrm{ct}_X}^{-1}(Y) = 2Y\left(I_r + X^\top Y\right)^{-1} + 2X\left(I_r + Y^\top X\right)^{-1} - 2X.
  ```

* ベクトル輸送
  
  ```math
  \mathcal{T}_U(V) = \mathrm{P}_{T_{R_X(U)} \mathop{\mathrm{St}}(n, r)}(V)
  ```

### 実行例
対称行列 $A$ を以下で生成する．

```python
import numpy as np


n = 10
r = 3
np.random.seed(0)

A = np.random.randn(n, n)
A = 0.5 * (A + A.T)
```

上記で紹介した ２ 通りの共役勾配法による Brockett cost function を最小化するクラスを `Brockett`, `BrockettInvR` で実装した．詳細な実装はページ末に記載する．
以下のコードで，初期点を $X_0 = \left(I_r, O\right)^\top$ として実行した．

```python
x0 = np.eye(n, r)

brockett = Brockett(A=A, r=r, max_iter=1000, tol=1e-6)
xopt = brockett.optimize(x0, verbosity=1) # Terminated after 133 iterations.

brockett_invr = BrockettInvR(A=A, r=r, max_iter=1000, tol=1e-6)
xopt_invr = brockett_invr.optimize(x0, verbosity=1) # Terminated after 139 iterations.
```

得られた解の正しさは，numpy.linalg.eigh の結果と比較することで確かめられる．

```python
_, V = np.linalg.eigh(A)
print(np.sum(xopt * V[:, :r][:, ::-1], axis=0)) # [1. 1. 1.]

_, V = np.linalg.eigh(A)
print(np.sum(xopt_invr * V[:, :r][:, ::-1], axis=0)) # [1. 1. 1.]
```

## 実装
リーマン多様体の抽象クラス

```python
from abc import ABCMeta, abstractmethod
from typing import Union

import numpy as np


class Manifold(metaclass=ABCMeta):
    @abstractmethod
    def inner_product(self, x: np.ndarray,
                      u: np.ndarray, v: np.ndarray) -> float:
        raise NotImplementedError('The function retraction is not implemented')

    @abstractmethod
    def norm(self, x: np.ndarray, u: np.ndarray) -> float:
        raise NotImplementedError('The function norm is not implemented')

    @abstractmethod
    def retraction(self, x: np.ndarray, v: np.ndarray) -> np.ndarray:
        raise NotImplementedError('The function retraction is not implemented')

    def distance(self, x: np.ndarray,
                 y: np.ndarray) -> Union[float, np.ndarray]:
        raise NotImplementedError('The function distance is not implemented')

    def exp(self, x: np.ndarray, v: np.ndarray) -> np.ndarray:
        raise NotImplementedError('The function exp is not implemented')

    def log(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        raise NotImplementedError('The function log is not implemented')

    def parallel_transport(self, x: np.ndarray,
                           y: np.ndarray, u: np.ndarray) -> np.ndarray:
        raise NotImplementedError(
            'The function parallel_transport is not implemented')

    def vector_transport(self, x: np.ndarray,
                         y: np.ndarray, u: np.ndarray) -> np.ndarray:
        raise NotImplementedError(
            'The function vector_transport is not implemented')

    def inverse_retraction(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        raise NotImplementedError(
            'The function inverse_retraction is not implemented')

    def gradient(self, x: np.ndarray, d: np.ndarray) -> np.ndarray:
        raise NotImplementedError('The function gradient is not implemented')
```

シュティーフェル多様体のクラス

```python
import numpy as np


class Stiefel(Manifold):
    def inner_product(self, x: np.ndarray,
                      u: np.ndarray, v: np.ndarray) -> float:
        return np.vdot(u, v)

    def norm(self, x: np.ndarray, u: np.ndarray) -> float:
        return np.linalg.norm(u)

    def retraction(self, x: np.ndarray, v: np.ndarray) -> np.ndarray:
        q, r = np.linalg.qr(x + v)
        return q * np.sign(np.sign(np.diag(r)) + .5)

    def vector_transport(self, x: np.ndarray,
                         y: np.ndarray, u: np.ndarray) -> np.ndarray:
        ytu = y.T @ u
        return u - y @ (ytu + ytu.T) * 0.5

    def inverse_retraction(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        xty = x.T @ y
        inv_matrix = np.linalg.inv(np.identity(xty.shape[0]) + xty)
        return 2 * y @ inv_matrix + 2 * x @ inv_matrix.T - 2 * x
```

共役勾配法のクラス

```python
from abc import ABCMeta, abstractmethod
from typing import Tuple

import numpy as np


class ConjugateGradient(metaclass=ABCMeta):
    def __init__(
            self,
            manifold: Manifold,
            initial_step: float = 1.0,
            armijo_param: float = 1e-4,
            max_iter: int = 300,
            tol: float = 1e-6,
            min_step_size: float = 1e-16,
            extended_output: bool = False
    ):
        self.manifold = manifold
        self.initial_step = initial_step
        self.armijo_param = armijo_param
        self.max_iter = max_iter
        self.tol = tol
        self.min_step_size = min_step_size
        self.extended_output = extended_output

        self.objective_values_ = []

    @abstractmethod
    def _f(self, x: np.ndarray) -> float:
        raise NotImplementedError('The function _f is not implemented')

    @abstractmethod
    def _df(self, x: np.ndarray) -> np.ndarray:
        raise NotImplementedError('The function _df is not implemented')

    def _beta(self, x1: np.ndarray, x2: np.ndarray,
              df1: np.ndarray, df2: np.ndarray, d: np.ndarray) -> float:
        inner1 = self.manifold.inner_product(x1, df1, df1)
        inner2 = self.manifold.inner_product(x2, df2, df2)
        return inner2 / inner1

    def _stopping_criterion(self, x: np.ndarray) -> float:
        d = self.manifold.norm(x, self._df(x))
        return d

    def _step_size(self, x: np.ndarray, d: np.ndarray, g: float, fx: float,
                   fxold: float = None) -> Tuple[float, np.ndarray]:
        """
        Armijo condition with back tracking.
        """
        if fxold is not None:
            t = 2 * (fx - fxold) / g
            t *= 2
        else:
            t = self.initial_step / self.manifold.norm(x, d)
        res = self.manifold.retraction(x, t * d)
        while self._f(res) > fx + self.armijo_param * t * g:
            t *= 0.5
            res = self.manifold.retraction(x, t * d)
            # break when step size 't' becomes too small
            if t <= self.min_step_size:
                return 0, x
        return t, res

    def _vector_transport(self, x1: np.ndarray, x2: np.ndarray,
                          d: np.ndarray, step_size: float) -> np.ndarray:
        v = self.manifold.vector_transport(x1, x2, d)
        norm_d = self.manifold.norm(x1, d)
        norm_v = self.manifold.norm(x2, v)
        s = norm_d / norm_v
        return v if s > 1 else s * v

    def optimize(self, x: np.ndarray, verbosity: int = 0) -> np.ndarray:
        fx = self._f(x)

        # store objective function value if necessary
        if self.extended_output:
            self.objective_values_.append(fx)

        # first iteration
        z = x.copy()
        df1 = self._df(z)
        d = -df1.copy()
        g = self.manifold.inner_product(x, df1, d)
        step_size, res = self._step_size(z, d, g, fx, None)

        fxold = fx
        fx = self._f(res)

        # store objective function value if necessary
        if self.extended_output:
            self.objective_values_.append(fx)

        # main loop
        i = 0
        for i in range(self.max_iter - 1):
            # direction
            df2 = self._df(res)
            beta = self._beta(z, res, df1, df2, d)
            d = -df2 + beta * self._vector_transport(z, res, d, step_size)

            g = self.manifold.inner_product(res, df2, d)
            if g >= 0:
                if verbosity >= 2:
                    print(f'Got an ascent direction (g = {g}), '
                          'reset the search direction')
                d = -df2.copy()
                g = self.manifold.inner_product(res, df2, d)

            df1 = df2.copy()

            # update
            z = res.copy()
            step_size, res = self._step_size(res, d, g, fx, fxold)

            fxold = fx
            fx = self._f(res)

            # store objective function value if necessary
            if self.extended_output:
                self.objective_values_.append(fx)

            if step_size == 0:
                if verbosity >= 2:
                    print(f'Got a zero step size.')
                break

            # break if convergence
            if self._stopping_criterion(res) < self.tol:
                break
        if verbosity >= 1:
            print(f'Terminated after {i + 2} iterations.')
        return res


class ConjugateGradientInvR(ConjugateGradient, metaclass=ABCMeta):
    @abstractmethod
    def _f(self, x: np.ndarray):
        raise NotImplementedError('The function _f is not implemented')

    @abstractmethod
    def _df(self, x: np.ndarray):
        raise NotImplementedError('The function _df is not implemented')

    def _vector_transport(self, x1: np.ndarray, x2: np.ndarray,
                          d: np.ndarray, step_size: float) -> np.ndarray:
        v = - self.manifold.inverse_retraction(x2, x1) / step_size
        norm_d = self.manifold.norm(x1, d)
        norm_v = self.manifold.norm(x2, v)
        s = norm_d / norm_v
        return v if s > 1 else s * v
```

Brockett cost function のソルバークラス

```python
import warnings

import numpy as np
from sklearn.utils.validation import check_symmetric


class Brockett(ConjugateGradient):
    def __init__(
            self,
            A: np.ndarray,
            r: int,
            initial_step: float = 1.0,
            armijo_param: float = 1e-4,
            max_iter: int = 300,
            tol: float = 1e-6,
            min_step_size: float = 1e-16,
            extended_output: bool = False
    ):
        if A.ndim != 2:
            raise ValueError(
                f'Expected 2D array, got {A.ndim}D array instead.')
        super().__init__(
            manifold=Stiefel(),
            initial_step=initial_step,
            armijo_param=armijo_param,
            max_iter=max_iter,
            tol=tol,
            min_step_size=min_step_size,
            extended_output=extended_output
        )
        n = A.shape[0]
        if r > n:
            warnings.warn(f'Given r = {r} is larger than {n}. '
                          f'Use r = {n} instead.')
            r = n
        self.A = check_symmetric(A)
        self.N = np.arange(1, r + 1) / r

    def _f(self, x: np.ndarray) -> float:
        return np.sum((self.A @ x) * (x * self.N))

    def _df(self, x: np.ndarray) -> np.ndarray:
        AX = self.A @ x
        AXN = AX * self.N
        return 2 * AXN - x @ (x.T @ AXN) - x @ (AXN.T @ x)


class BrockettInvR(ConjugateGradientInvR):
    def __init__(
            self,
            A: np.ndarray,
            r: int,
            initial_step: float = 1.0,
            armijo_param: float = 1e-4,
            max_iter: int = 300,
            tol: float = 1e-6,
            min_step_size: float = 1e-16,
            extended_output: bool = False
    ):
        if A.ndim != 2:
            raise ValueError(
                f'Expected 2D array, got {A.ndim}D array instead.')
        super().__init__(
            manifold=Stiefel(),
            initial_step=initial_step,
            armijo_param=armijo_param,
            max_iter=max_iter,
            tol=tol,
            min_step_size=min_step_size,
            extended_output=extended_output
        )
        n = A.shape[0]
        if r > n:
            warnings.warn(f'Given r = {r} is larger than {n}. '
                          f'Use r = {n} instead.')
            r = n
        self.A = check_symmetric(A)
        self.N = np.arange(1, r + 1) / r

    def _f(self, x: np.ndarray) -> float:
        return np.sum((self.A @ x) * (x * self.N))

    def _df(self, x: np.ndarray) -> np.ndarray:
        AX = self.A @ x
        AXN = AX * self.N
        return 2 * AXN - x @ (x.T @ AXN) - x @ (AXN.T @ x)
```

参考までに，Pymanopt にも共役勾配法の実装がある．使用例については過去記事を参照されたい．

* [【Python】リーマン多様体上の共役勾配法 - Qiita](https://qiita.com/Mrrmm252/items/d48a2eee45d046b707ee)

# 参考文献
* [Absil, P.-A., Mahony, R., & Sepulchre, R. (2009). Optimization algorithms on matrix manifolds. Princeton University Press.](https://press.princeton.edu/absil)
* [Ring, W., & Wirth, B. (2012). Optimization methods on Riemannian manifolds and their application to shape space. SIAM Journal on Optimization, 22(2), 596-627.](https://epubs.siam.org/doi/pdf/10.1137/11082885X?casa_token=Gwdo0RnxCIkAAAAA:RT76NFZgsO2f0EcLRV6UzY_WrKS_KFmaS4DTK1CxCKpddToD5kNvwb8jHmUp1hpLDGKZefEuEr2Zow)
* [Sato, H., & Iwai, T. (2015). A new, globally convergent Riemannian conjugate gradient method. Optimization, 64(4), 1011-1031.](https://arxiv.org/pdf/1302.0125.pdf)
* [Zhu, X., & Sato, H. (2020). Riemannian conjugate gradient methods with inverse retraction. Computational Optimization and Applications, 77(3), 779-810.](http://www.optimization-online.org/DB_FILE/2020/05/7798.pdf)