Okapi BM25をスパース行列のまま計算するPythonライブラリを作った

情報検索で使われる単語の重み付け方法のひとつにOkapi BM25というものがあります。文献によって細かな違いはありますが，今回は Wikipediaに載ってるやつを使うことにします。

それぞれの文書におけるそれぞれの語の重みを表す行列を計算するクラスを作りたいのですが，いろいろと省略すると，以下のような形式の計算を行う必要が出てきます。

$\displaystyle W_{i,j}=\frac{A_{i,j}}{A_{i,j}+b_{i}}$

ここで， $A$ は $N\times M$ の疎行列で $\forall A_{i,j}\geq0$ ， $b$ は $N\times1$ の密なベクトルで $\forall b_{i}>0$ です。

この計算は，NumPyのブロードキャストルールに従えば以下のように書けるはずです。

$\displaystyle W=\frac{A}{A+b}$

残念ながら，この通りに書くだけでは都合の悪いことがあります。 $A$ は疎行列で，展開するとメモリにやさしくないので疎行列のままにしたいのですが，分母は $b$ が密なベクトルなので $A+b$ を計算すると密な行列になってしまいます。ところが，分子は疎行列の $A$ なので計算結果の $W$ は明らかに疎行列です。そういうわけで，密な行列になってしまう $A+b$ の計算を回避しながら $W$ の計算を行いたいです。

式変形でうまく回避できないか模索したのですが，あまりシンプルな方法が見つからなかったので，ここではscipy.sparseの内部データ構造を利用した計算方法を紹介します。

最初に，疎行列をscipy.sparse.csr_matrixに変換します (scikit-learnで使うことを前提にしている場合，疎行列にはcsr_matrixを選択することが多いと思います)。csr_matrixの内部データ構造については以下のページが詳しいです。

scipy.sparseの内部データ構造 – はむかず！

csr_matrixでは， $A$ の各行の非ゼロ要素の数は以下のようにして求められます。

sz = A.indptr[1:] - A.indptr[0:-1]

次に $b_i$ の値を，対応する $A$ の各行の非ゼロ要素の数だけ繰り返します。numpy.repeatを使うことによって実現できます。

rep = numpy.repeat(b, sz)

最後に， $A$ の非ゼロ要素の部分のみで $\displaystyle W=\frac{A}{A+b}$ の計算を行ってscipy.sparse.csr_matrixを再構築します。

data = A.data / (A.data + rep)
return scipy.sparse.csr_matrix((data, A.indices, A.indptr), shape=A.shape)

動作例を書いておくと，以下のような感じになります。

A = scipy.sparse.csr_matrix([[4, 0, 0], [5, 6, 0]], dtype=scipy.float_)
b = scipy.array([7, 8], dtype=scipy.float_)
sz = A.indptr[1:] - A.indptr[0:-1] # => array([1, 2], dtype=int32)
rep = scipy.repeat(b, sz) # => array([7, 8, 8])
data = A.data / (A.data + rep) # => array([ 0.36363636,  0.38461538,  0.42857143])
ans = scipy.sparse.csr_matrix((data, A.indices, A.indptr), shape=A.shape)
ans.todense() # => matrix([[0.36363636, 0., 0.], [0.38461538, 0.42857143, 0.]])