Top-k retrievalのアルゴリズムを書いてみた(tiny topk)

最近top-k retrievalの話を少し聞いたので、簡単にコードを書いてみた。いつものように恥もなくgithubで公開している。

cpp-ToyBox-TinyTopK(github)

Top-k retrievalを簡単に説明すると、転置インデックスに対してdisjunctiveなクエリで問い合わせて(OR検索)スコアの上位k件を取得したいという話である。詳細は"[IR] 転置インデックスとtop-k query - tsubosakaの日記"がとてもわかり易いので、そちらを参照していただきたい。

Top-k retrievalではvector space modelでのdot product(もちろんcosineやBM25なども同様)を考えるので、同じ原理でコサイン距離でのk-nearest neighborや線形分類器のtop-kマルチラベル分類を行うことができたりと、(実用的に使う場面があるかどうかは別として)それなりに適用できる範囲は広いのではないかと思う。

今回のtiny topkは"Evaluation Strategies for Top-k Queries over Memory-Resident Inverted Indexes"と同じように、インデックスを全てメモリ上に載せることを前提としているので、あまり多くのドキュメントを扱えないので注意が必要。posting listを圧縮して持つ実装にすればそれなりの量を扱えるはず。検索アルゴリズムとしてはnaiveなTAAT、naiveなDAAT、WANDを実装した。max_scoreのアルゴリズムも実装してみたかったのだが、ちょうど良さそうなデータ構造が思いつかなかった。

以下、benchmarkディレクトリにあるプログラムで速度を計測した結果である。なおデータとしては手元にあったLIBSVM datasetsのa1aとnews20、rcv1.multiclassを用いた。そのため、実際のドキュメントのデータとは異なり、あまり参考にはならないかもしれない。

$ time ./bench_tinytopk ~/data/a1a.t ~/data/a1a.txt 
document num: 30956, query num: 1605
TAAT    3.354257 sec    2.089880 msec/req
DAAT    29.734377 sec   18.526092 msec/req
WAND    21.928297 sec   13.662490 msec/req

real    0m55.686s
user    0m52.679s
sys     0m2.820s


$ time ./bench_tinytopk ~/data/news20 ~/data/news20.t
document num: 15935, query num: 3993
TAAT    2.894727 sec    0.724950 msec/req
DAAT    40.442991 sec   10.128473 msec/req
WAND    46.316534 sec   11.599433 msec/req

real    1m33.483s
user    1m26.665s
sys     0m6.572s


$ time ./bench_tinytopk ~/data/rcv1_test.multiclass ~/data/rcv1_train.multiclass_1000 
document num: 518571, query num: 1000
TAAT    70.841805 sec   70.841805 msec/req
DAAT    530.754437 sec  530.754437 msec/req
WAND    587.605763 sec  587.605763 msec/req

real    23m47.418s
user    21m32.297s
sys     2m10.936s

全体的に悲しいくらいの速度しか出なかった。三つのアルゴリズムの比較では、naiveなTAATが同じくnaiveなDAATやWANDを大きく引き離す結果となった。DAATとWANDでは各posting listのイテレータが指す中で最小idのドキュメントを探索するためにヒープ木を使っているのだが、そこらへんの実装がいまいちなのかもしれない。DAATとWANDを比較すると、スキップが発生して高速なはずのWANDが三つ中一つのデータでのみ優れた結果となった。こちらもposting listのスキップに、単純にstd::lower_boundを使っているだけのため、実装が残念なだけのような気がする。