はじめに

この記事では行列をいくつかのブロックに分割して行列-行列積を計算するやり方についてまとめる（元ネタはここ）

基本的には自分用の備忘録なので、適当なこと書いてたらすみません。

背景

線形方程式や固有値解析などの数値計算の文脈ではしばしば巨大なサイズの行列が現れる。たとえば、有限要素法や有限体積法などで偏微分方程式を離散化した際に現れる係数行列などはこの典型であり、自由度の数によっては$O(10^6 \times 10^6)$程度、あるいはそれ以上の大きな行列となる。ただし、このような場合はデータ行列を疎行列として表現できるため、実際に必要なメモリ消費量は非零要素数によって決まり、基本的には見かけのサイズが大きくてもメモリに格納できる。

一方、機械学習の文脈ではデータ行列が大規模な密行列となることが普通であり、メモリに保持できないことが多い。その結果、行列演算を満足に行うことができないため、アルゴリズムを工夫しない限り大規模なデータは処理できないことになってしまう。こういうとき、データ行列を分割して逐次的に処理する必要が出てくる。

行列-行列積のブロッキング

教科書的なやり方

具体的に行列-行列積で考える。2つの行列$\mathbf{A} \in \mathbb{R}^{n \times m},\ \mathbf{B} \in \mathbb{R}^{m \times l}$の積$\mathbf{C} = \mathbf{AB} \in \mathbb{R}^{n \times l}$は次式のように計算できる。

C_{ij} = \sum_{k=1}^{m} A_{ik} B_{kj} = \mathbf{a}_{i}^{\top} \mathbf{b}_{j}

ただし、

\mathbf{A} = \begin{bmatrix}
\mathbf{a}_{1}^{\top} \\
\mathbf{a}_{2}^{\top} \\
\vdots \\
\mathbf{a}_{n}^{\top} 
\end{bmatrix}
,\ 
\mathbf{B} = \begin{bmatrix}
\mathbf{b}_{1} &
\mathbf{b}_{2} &
\cdots &
\mathbf{b}_{l}
\end{bmatrix}

とする。すなわち、$\mathbf{C}$の各要素は$\mathbf{A}$の行ベクトル$\mathbf{a}_{i}^{\top} \in \mathbb{R}^{m}$と$\mathbf{B}$の列ベクトル$\mathbf{b}_{j} \in \mathbb{R}^{m}$の内積として計算できる。なので、$\mathbf{A}$を行方向に、$\mathbf{B}$を列方向に分割できるのであれば、それらを逐次読み出しながら計算することができる。

とはいえ、通常、データ行列の行と列の意味付けは共通しており、異なる方向に分割することはあまりない気もする。たとえば、行方向に空間、列方向に時間を割り当て、データがサンプリング時刻ごとに個別のファイルに保存されている場合、列方向の分割は容易だが、行方向の分割は複数のファイルをまたぐため、効率的ではない。

列方向の分割

次に、$\mathbf{A}$全体がメモリに保持され、$\mathbf{B}$は列方向に分割されているとする。この場合、行列-ベクトル積を繰り返し、列ベクトル単位で結果を取得することで逐次処理が可能である。

\begin{align}
\mathbf{c}_{i} &= \mathbf{A} \mathbf{b}_{i} \\
\mathbf{C} &= \begin{bmatrix}
\mathbf{c}_{1} &
\mathbf{c}_{2} &
\cdots &
\mathbf{c}_{l}
\end{bmatrix}
\end{align}

これは$\mathbf{B}$全体がメモリに保持できない場合に使用できる。

行方向の分割

列方向の場合と同様のことを行方向でも考えることができる。この場合、結果は行ベクトル単位で取得する。

\begin{align}
\mathbf{c}_{i}^{\top} &= \mathbf{a}_{i}^{\top} \mathbf{B} \\
\mathbf{C} &= \begin{bmatrix}
\mathbf{c}_{1}^{\top} \\
\mathbf{c}_{2}^{\top} \\
\vdots \\
\mathbf{c}_{n}^{\top} 
\end{bmatrix}
\end{align}

こっちは$\mathbf{A}$全体がメモリに保持できない場合に使用する。$\mathbf{A}$が列方向に分割されていたとしても、その転置行列$\mathbf{A}^{\top}$との積が必要な場合に使えるかもしれない。

rank-1 update

教科書的なやり方の逆パターンで、$\mathbf{A}$が列方向、$\mathbf{B}$が行方向に分割されているとする。この場合、ベクトル同士の外積（outer product）の和を逐次的に更新することで結果の全体を得ることができる。$\mathbf{a}_{i}\mathbf{b}_{i}^{\top}$のランクは常に1なので、これはrank-1 updateと呼ばれる。

\mathbf{A} = \begin{bmatrix}
\mathbf{a}_{1} &
\mathbf{a}_{2} &
\cdots &
\mathbf{a}_{m}
\end{bmatrix}
,\ 
\mathbf{B} = \begin{bmatrix}
\mathbf{b}_{1}^{\top} \\
\mathbf{b}_{2}^{\top} \\
\vdots \\
\mathbf{b}_{m}^{\top} 
\end{bmatrix}

とすると、

\mathbf{a}\mathbf{b}^{\top} = 
\begin{pmatrix}
a_{1} \\
a_{2} \\
\vdots \\
a_{n}
\end{pmatrix}
\begin{pmatrix}
b_{1} &
b_{2} &
\cdots &
b_{l}
\end{pmatrix}
=
\begin{pmatrix}
a_{1} b_{1} & a_{1} b_{2} & \cdots &  a_{1} b_{l}\\
a_{2} b_{1} & a_{2} b_{2} & \cdots &  a_{2} b_{l}\\
\vdots & \vdots & \ddots & \vdots \\
a_{n} b_{1} & a_{n} b_{2} & \cdots &  a_{n} b_{l}\\
\end{pmatrix}

であり、

\mathbf{C} = \mathbf{a}_{1}\mathbf{b}_{1}^{\top} + \mathbf{a}_{2}\mathbf{b}_{2}^{\top} + \cdots + \mathbf{a}_{m}\mathbf{b}_{m}^{\top}

となる。この場合、$\mathbf{C}$を零行列で初期化して$\mathbf{C}:= \mathbf{C} + \mathbf{a}_{i}\mathbf{b}_{i}^{\top}$で更新する。

余談だが、rank-1 updateを用いた計算としてはSherman-Morrisonの公式が個人的に興味深い。これは逆行列を逐次更新する際に有用なもので、準ニュートン法の一種であるBroyden法などに応用されている。

まとめ

行列-行列積をブロック行列に分割しながら行うやり方についてまとめた。上記はベクトル1本ずつに分けて処理するやり方になっているが、実際にはメモリに保持できる範囲で複数のベクトルをまとめて処理する方が効率的である。

最近この手のことについて色々調べたいのだが、数値線形代数というワードで検索してもよさげな日本語文献が見当たらないのが残念だったりする（調べ方が悪いだけ？）。線形方程式や固有値計算に関するマニアックな記述は見つかるのだが、機械学習分野に対応した密行列を逐次的に処理するアルゴリズムを解説した文献なんかがあると嬉しい。もちろん、"numerical linear algebra"とかで検索すればどこかの大学の講義資料なんかが見つかるので問題はないっちゃないんだけど、そろそろ日本語でそういうのが読めてもいいんじゃないかと思ったり。特に乱数を使った行列の近似計算（Random SamplingとかRandom Projectionとか）はニーズが高いと思うんだけどなぁ。