エンジニアの谷田です。
最大内積探索問題(Maximum Inner Product Search, 以下MIPS)ってご存知でしょうか？データベースに登録された多くのアイテムのベクトルのうち、クエリのベクトルとの内積を最大化するアイテムを探す問題です。行列分解を用いてユーザにアイテムをレコメンドするときなど、この探索が問題になってくることがあります。
MIPSを解く単純な方法は、与えられたクエリに対してデータベースのすべてのアイテムとの内積を計算することですが、これにはアイテム数に比例した時間がかかります。そのため、アイテムが何千万や何億といった数にのぼると、ベストなアイテムを返す処理を遅延なくリアルタイムで行うのは厳しくなってしまいます。
これを解決するために、局所性鋭敏型ハッシュ(Locality Sensitive Hashing, 以下LSH)という手法をもとにしたアルゴリズムが提案されてきました。LSHは実ベクトルから小さい次元数の2値ベクトル(ハッシュ値)を計算する手法で、類似したアイテムからは同じハッシュ値が得られやすいという特徴があります。ALSHやSimple-LSHなどの手法では、入力するベクトルを変換することで、コサイン類似度のためのLSHを用いて内積を最大化するアイテムの近似探索を高速に行えることが示されました。
しかし実際にこれらの手法を使ってみると、アイテムに対して出力されるハッシュ値に偏りがあるなどして、期待する性能が得られないことがあります。もともとのLSHでも同様の問題があり、これを解決するためにLearning to Hashと総称されるハッシュ関数を学習する手法が提案されてきました。ITQ(ITerative Quantization)もそのうちの手法のひとつで、データベースのアイテムからハッシュ関数を学習して、より上手にハッシュ値を計算できるようにします。
そこで、ALSHやSimple-LSHとITQを組み合わせることで、MIPSに対してもより良いハッシュ値の計算を行い、近似探索の性能を向上させることができるのではないかと考えたのが、今回の本題であるライブラリのアルゴリズムです。SITQと名付けたこのアルゴリズムでは、Simple-LSHと同様の方法でアイテムやクエリのベクトルを変換し、その上でITQを用いてハッシュ関数を最適化します。探索では、クエリのハッシュ値と似ているハッシュ値を持つアイテムを取ってきて、そのなかで実際に内積を計算して最も大きい値となったアイテムを取得します。実際に内積を計算するアイテム数はデータベースのすべてのアイテムに比べて遥かに少ないので、探索にかかる時間を短くすることが期待できます。
このアルゴリズムをPythonのライブラリとして公開しました。pip install sitqでインストールして使うことができます。ハッシュ関数の学習およびハッシュ値の計算を行うSitqと、内積を最大化するアイテムの近似探索を行うMipsのふたつのクラスがあります。
Sitqでは次のようにしてハッシュ値を計算することができます。
import numpy as np
from sitq import Sitq
# サンプルデータを作成
items = np.random.rand(10000, 50)
query = np.random.rand(50)
sitq = Sitq(signature_size=8)
# ハッシュ関数を学習
sitq.fit(items)
# アイテムのハッシュ値を取得
item_sigs = sitq.get_item_signatures(items)
# クエリのハッシュ値を取得
query_sig = sitq.get_query_signatures([query])[0]
Mipsでは次のようにしてクエリに対してアイテムを探索することができます。
import numpy as np
from sitq import Mips
# サンプルデータを作成
items = np.random.rand(10000, 50)
query = np.random.rand(50)
mips = Mips(signature_size=8)
# アイテムを登録
mips.fit(items)
# クエリに対して内積を最大化しそうなアイテムを探索
item_indexes, scores = mips.search(query, limit=10, distance=1)
ベンチマークでは、評価した他の近似探索アルゴリズムと比べてハッシュ値の偏りが小さく精度も高いという結果が得られました。実験結果はGitHubリポジトリに載せています。
また、MovieLens 20M Datasetデータセットに対してレコメンドを行うKaggle Kernelも作りました。レコメンドにはALSで学習したベクトルを使い、ブルートフォースとSITQを用いた場合の性能を出しています。
