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

対象とする問題

リーマン多様体 $\mathcal{M}$ 上の無制約最適化問題

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

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

リーマン多様体上の共役勾配法 [Absil et al. 2009]

リーマン多様体 $\mathcal M$ 上の共役勾配法のアルゴリズムは， Absil らの多様体上の最適化の書籍 "Optimization Algorithms on Matrix Manifolds" にて紹介されている．

$x^0$ を初期化する．
$d^0=−\mathop{\mathrm{grad}} f(x^0)$ で探索方向を初期化する．
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)$ で探索方向を更新する．

end for

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

リーマン多様体上の共役勾配法 [Sato & Iwai 2015]

上記で紹介したリーマン多様体上の共役勾配法は，ベクトル輸送が非増加

\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]．

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\}

$x^0$ を初期化する．
$d^0=−\mathop{\mathrm{grad}} f(x^0)$ で探索方向を初期化する．
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)$ で探索方向を更新する．

end for

リーマン多様体上の共役勾配法 [Zhu & Sato 2020]

$\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$ 上の自然な表現として，

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

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

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}$ を用いた次の共役勾配法が得られる．

$x^0$ を初期化する．
$d^0=−\mathop{\mathrm{grad}} f(x^0)$ で探索方向を初期化する．
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)$ で探索方向を更新する．

end for

Python での実装

使用例

問題

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

\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$$

で計算される．

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

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

上の無制約最適化問題

\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)}$ を用いて，

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

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

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)}$ は

\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)$ は，

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

で計算できる．詳細は "Optimization Algorithms on Matrix Manifolds" の 4.8 節にある．

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

内積

\left\langle U, V\right\rangle_X = \mathop{\mathrm{tr}}\left(U^\top V\right)

レトラクション

\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)$ は逆写像の計算が可能であり，

{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.

ベクトル輸送

\mathcal{T}_U(V) = \mathrm{P}_{T_{R_X(U)} \mathop{\mathrm{St}}(n, r)}(V)

実行例

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

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$ として実行した．

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 の結果と比較することで確かめられる．

_, 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.]

実装

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

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')

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

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

共役勾配法のクラス

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 のソルバークラス

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

参考文献