行列の幾何平均計算の実装

正定値Hermite行列に対する幾何平均をNumPyで実装

ブラインド音源分離手法の多チャネル非負値行列因子分解(MNMF)や独立半正定値テンソル分析（IPSDTA）には，半正定値Hermite行列に対する幾何平均の計算が登場します．Pythonの数値計算ライブラリのNumPyでこれらの手法を実装してみました．その際の工夫について共有します．

正定値Hermite行列に対する幾何平均の定義

2つの正定値Hermite行列$\bm{A},\bm{B}\in\mathbb{C}^{N\times N}$が与えられたとき，$\bm{A}$と$\bm{B}$の幾何平均は，次の式を満たす正定値Hermite行列$\bm{X}\in\mathbb{C}^{N\times N}$を求めることを考えます．

$$\bm{X}\bm{A}^{-1}\bm{X} = \bm{B}~~~(1)$$

この式の解は，次のように定義される行列の幾何平均で与えられます．

$$\begin{align} \bm{X} = \bm{A}\#\bm{B} &:= \bm{A}^{\frac{1}{2}}\left(\bm{A}^{-\frac{1}{2}}\bm{B}\bm{A}^{-\frac{1}{2}}\right)^{\frac{1}{2}}\bm{A}^{\frac{1}{2}}~~~(2) \\ &= \bm{A}\left(\bm{A}^{-1}\bm{B}\right)^{\frac{1}{2}}~~~(3) \\ &= \left(\bm{A}\bm{B}^{-1}\right)^{\frac{1}{2}}\bm{B}~~~(4) \end{align}$$

式(1)の両辺の逆行列を計算し整理すると，次のように変形も可能です．

$$\bm{X}\bm{B}^{-1}\bm{X} = \bm{A}$$

つまり，$\bm{B}\#\bm{A} = \bm{A}\#\bm{B}$が成立します．以降は，式(2)~(4)を用いて説明を進めますが，$\bm{A}$と$\bm{B}$を入れ替えても同様の議論が可能であることに注意してください．

NumPyで正定値Hermite行列の幾何平均を実装する際の問題点

式(2)〜(4)からわかるように，行列の幾何平均を計算する際には，行列の平方根の実装が必要になります．
NumPyのnumpy.ndarrayで行列の平方根を計算する方法の一つは，scipy.linalg.sqrtmを利用することです．この関数を使うと，平方根の計算の実装が不要で便利な一方，並列処理ができない欠点があります^[1]．応用先のMNMFやIPSDTAでは，周波数ごとに幾何平均を計算するため，scipy.linalg.sqrtmをforループで回す必要があり，計算時間のボトルネックとなります．また，この関数は一般の行列の平方根を計算するために用意されており，正定値Hermite行列の性質を使って計算していません．そこで，並列処理が可能で，正定値Hermite行列の性質を使った効率的な幾何平均の実装をNumPyでできないか考えてみました．

正定値Hermite行列に対する平方根の実装

幾何平均の具体的な実装に入る前に，式(2)に登場する正定値Hermite行列に対する平方根の計算とその実装を確認します．
正定値Hermite行列$\bm{A}\in\mathbb{C}^{N\times N}$に対する行列の平方根は，

$$\bm{A}^{\frac{1}{2}} = \bm{Z}\left( \begin{array}{cccc} \sqrt{\lambda_{1}} & 0 & \cdots & 0 \\ 0 & \sqrt{\lambda_{2}} & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \sqrt{\lambda_{N}} \end{array} \right)\bm{Z}^{-1} \\ \bm{Z} = \left( \begin{array}{cccc} \bm{z}_{1} & \bm{z}_{2} & \cdots & \bm{z}_{N} \end{array} \right)\in\mathbb{C}^{N\times N}$$

で定義されます．ただし，$\lambda_{n}$は$\bm{A}$の固有値分解により得られる固有値で，$\bm{z}_{n}$は$\lambda_{n}$対応する固有ベクトルです．$\bm{A}$は正定値Hermite行列のため，$\lambda_{n}$は常に正の値となり，$\sqrt{\lambda_{n}}$が実数の範囲で計算できます．

正定値Hermite行列に対する固有値と固有ベクトルは，numpy.linalg.eighで計算可能です．この関数は，並列処理が可能な利点があります．また，一般の正方行列に対する固有値と固有ベクトルを計算するnumpy.linalg.eigとは異なり，正定値Hermite行列の性質を利用した効率的な計算が可能です．

まとめると，正定値Hermite行列に対する平方根は，次のように実装できます．

import numpy as np

def make_matrix(K: int, N: int) -> np.ndarray:
    """Make Hermitian positive definite matrix.

    Args:
        K (int): Number of matrices.
        N (int): Size of matrix.

    Returns:
        np.ndarray: Hermitian positive definite matrix of shape (K, N, N).
    """
    vec = np.random.randn(K, 1, N, 2 * N) + 1j * np.random.randn(K, 1, N, 2 * N)
    cov = vec * vec.transpose(0, 2, 1, 3).conj()

    return np.mean(cov, axis=-1)

def sqrt_hermite(X: np.ndarray) -> np.ndarray:
    """Square root of Hermitian positive definite matrix.

    Args:
        X (np.ndarray): Hermitian positive definite matrix of shape (K, N, N).

    Returns:
        np.ndarray: Square root of X of shape (K, N, N).
    """
    lamb, Z = np.linalg.eigh(X)

    Z_hermite = Z.swapaxes(-2, -1)
    Z_hermite = Z_hermite.conj()
    Lamb = np.sqrt(lamb)[..., np.newaxis] * np.eye(lamb.shape[-1])

    return Z @ Lamb @ Z_hermite

正定値Hermite行列に対する幾何平均の実装（その1）

正定値Hermite行列に対する平方根の実装を踏まえると，式(2)の定義で幾何平均を計算する場合は，numpy.linalg.eighを使用して次のように実装できます．

def geo_mean_hermite_type1(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """Geometric mean of Hermitian positive definite matrices.

    Args:
        A (np.ndarray): Hermitian positive definite matrix of shape (K, N, N).
        B (np.ndarray): Hermitian positive definite matrix of shape (K, N, N).

    Returns:
        np.ndarray: Geometric mean of A and B of shape (K, N, N).
    """
    A_sqrt = sqrt_hermite(A)
    A_sqrt_inv = np.linalg.inv(A_sqrt)
    ABA = sqrt_hermite(A_sqrt_inv @ B @ A_sqrt_inv)

    return A_sqrt @ ABA @ A_sqrt

一般化固有値分解

式(3)と(4)の定義で幾何平均を計算する場合，$\bm{A}^{-1}\bm{B}$や$\bm{A}\bm{B}^{-1}$の平方根を計算する必要があります．正定値Hermite行列に対する平方根の実装で説明した方法で実装したいところですが，$\bm{A}$と$\bm{B}$がいずれも正定値Hermite行列の場合でも，$\bm{A}^{-1}\bm{B}$は一般には正定値Hermite行列になりません．そのため，numpy.linalg.eigを用いて，次の固有値分解を解く必要があります．

$$\bm{A}^{-1}\bm{B}\tilde{\bm{z}} = \tilde{\lambda}\tilde{\bm{z}}~~~(5)$$

この固有値問題が計算できれば，次のように$\bm{A}^{-1}\bm{B}$の平方根が計算可能です．

$$\left(\bm{A}^{-1}\bm{B}\right)^{\frac{1}{2}} = \tilde{\bm{Z}}\left( \begin{array}{cccc} \sqrt{\tilde{\lambda}_{1}} & 0 & \cdots & 0 \\ 0 & \sqrt{\tilde{\lambda}_{2}} & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \sqrt{\tilde{\lambda}_{N}} \end{array} \right)\tilde{\bm{Z}}^{-1}$$

この方法で実装する場合，numpy.linalg.eigを用いた並列処理が可能ですが，正定値Hermite行列の性質を利用できていません．そこで，式(5)の両辺に$\bm{A}$を乗じてみます．

$$\bm{B}\tilde{\bm{z}} = \tilde{\lambda}\bm{A}\tilde{\bm{z}}~~~(6)$$

行列$\bm{A}$と$\bm{B}$が与えられたときに，式(6)を満たすスカラー$\tilde{\lambda}$とベクトル$\tilde{\bm{z}}$を求める問題は，一般化固有値問題と呼ばれます．式(6)は，$\bm{A}$が正則な場合(＝逆行列が存在)，式(5)の形に変形することで固有値問題に帰着できます．しかし，今回は，$\bm{A}$と$\bm{B}$が正定値Hermite行列であることを利用して実装する方法を考えます．

半正定値Hermite行列のための一般化固有値分解

式(6)を解くために，$\bm{A}$を次のようにCholesky分解します．

$$\bm{A}=\bm{L}\bm{L}^{\mathsf{H}}$$

NumPyではnumpy.linalg.choleskyで$\bm{L}$を並列処理で求めることができます．この分解を利用すると，正定値Hermite行列を用いた一般化固有値問題(6)は，正定値Hermite行列を用いた固有値問題に帰着できます．

$$\begin{align} \bm{C}\bm{y} &= \mu\bm{y}~~~(7) \\ \bm{C} &:= \bm{L}^{-1}\bm{B}\bm{L}^{-\mathsf{H}} \\ \bm{y} &:= \bm{L}^{\mathsf{H}}\bm{z} \end{align}$$

ただし，$\mu$は半正定値Hermite行列$\bm{C}$に対する固有値で，$\bm{y}$は固有ベクトルです．固有値問題(7)は，正定値Hermite行列に対する平方根の実装と同様にnumpy.linalg.eighを用いて固有値と固有ベクトルを計算できます．一般化固有値問題(6)に対する固有値と固有ベクトルは，求めた$\bm{y}$を用いて

$$\begin{align} \tilde{\lambda}_{n} &= \mu_{n} \\ \tilde{\bm{z}}_{n} &= \bm{L}^{-\mathsf{H}}\bm{y}_{n} \end{align}$$

で得られます．

以上を踏まえると，正定値Hermite行列のための一般化固有値分解は，次のように実装できます．

def gen_eig_hermite(A: np.ndarray, B: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Generalized eigenvalue decomposition for Hermitian positive definite matrices.

    Args:
        A (np.ndarray): Eigenvalues of Hermitian positive definite matrix of shape (K, N, N).
        B (np.ndarray): Eigenvectors of Hermitian positive definite matrix of shape (K, N, N).

    Returns:
        np.ndarray: Geometric mean of A and B of shape (K, N, N).
    """
    L = np.linalg.cholesky(A)
    L_inv = np.linalg.inv(L)
    L_inv_hermite = np.swapaxes(L_inv, -2, -1)
    L_inv_hermite = L_inv_hermite.conj()

    C = L_inv @ B @ L_inv_hermite
    lamb, Y = np.linalg.eigh(C)
    Z = L_inv_hermite @ Y

    return lamb, Z

正定値Hermite行列に対する幾何平均の実装（その2）

これまでの説明をまとめると式(3)に基づく行列の幾何平均は

$$\bm{A}\left(\bm{A}^{-1}\bm{B}\right)^{\frac{1}{2}} = \bm{A}\tilde{\bm{Z}}\left( \begin{array}{cccc} \sqrt{\tilde{\lambda}_{1}} & 0 & \cdots & 0 \\ 0 & \sqrt{\tilde{\lambda}_{2}} & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \sqrt{\tilde{\lambda}_{N}} \end{array} \right)\tilde{\bm{Z}}^{-1}$$

と整理できます．

def geo_mean_hermite_type2(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """Geometric mean of Hermitian positive definite matrices.

    Args:
        A (np.ndarray): Hermitian positive definite matrix of shape (K, N, N).
        B (np.ndarray): Hermitian positive definite matrix of shape (K, N, N).

    Returns:
        np.ndarray: Geometric mean of A and B of shape (K, N, N).
    """
    lamb, Z = gen_eig_hermite(A, B)
    lamb = np.sqrt(lamb)
    Lamb = lamb[..., np.newaxis] * np.eye(Z.shape[-1])
    AB = Z @ Lamb @ np.linalg.inv(Z)

    return A @ AB

まとめ

MNMFやIPSDTAのための正定値Hermite行列の計算の実装方法を説明しました．今回紹介した方法は，並列処理が可能であり，正定値Hermite行列の性質を利用した効率的な実装を利用しています．

脚注

1. 各行列を1次元に変形して，numpy.apply_along_axisを利用すれば，擬似的に並列計算は可能です．