こんにちは！ SAIGの金子です。
普段はフューチャーのAIグループで開発を行っている他、nadareというハンドルネームでデータ分析コンペティションに参加しています。

アルゴリズムとデータ構造連載の1本目となる今回はアルゴリズムをテーマとして、faissを題材にソートされたTopKを取得するために使われているソートアルゴリズムについて紹介したいと思います。

faissとは

https://github.com/facebookresearch/faiss

faissはFacebook AI Researchによって開発されている近傍探索のライブラリです。近傍探索はベクトル同士の類似度(ユークリッド距離やコサイン類似度)を計算することで似たベクトルを取得します。主に文章や画像をベクトル化し類似アイテムやクエリに合ったものを検索することに応用されます。

TopKを取得するのに最適なアルゴリズムは？

faissは大量のベクトル(1M, 1G, 1P単位)から類似ベクトルを取得する際に活用されます。一般的なケースでは全てのベクトルとの類似度とその順位を正確に知る必要はなく、最も近いベクトルかあるいはTop100程度を知ることができればうれしい場合が多いです。そのため、faiss内ではTopKを取得するために基本的にはヒープソートが用いられています。

ヒープソートについて

ヒープソートは二分木を用いたソート方法です。N個のアイテムについてクエリとの類似度を計算してTopKを得たい場合を考えるとします。

ヒープソートでは

• まず容量Kのヒープを用意します.
• N個のアイテムについて順に類似度を計算しヒープに追加しヒープ内でソートを行います。
• ソート後は順位がKより大きいものをヒープから削除します。

1回の追加につきヒープのソートにかかる時間はlogKで、これをN回繰り返すため計算量はO(NlogK)になります。K << Nであるため、普通のソートのO(NlogN)と比べても計算時間が改善することが多いです。

ヒープキューvs他のソート

内積やユークリッド距離で近傍を取得する際はヒープキューを基本的には用いるのですが、ハミング距離で近傍探索を行う場合、faissではバケットソートも選択できます。

IndexBinaryFlatにてuse_heapをfalseにするとバケットソートによるtopKの取得ができます。私のPCの環境で2e6件のレコードに対し, 2e5回256bitのハミング距離を計算しtop1000を取得するコードをを実行したところ、ヒープソートで3分10秒、バケットソートで3分20秒とヒープソートの方が早かったです。バケットソートはO(N)なので意外ですが、各バケットに最低でもtopKの容量を用意する必要があり効率が良くなかったのではと考えられます。

余談ですが、上記の256bitのバイナリベクトルは300次元のベクトルをハッシュ化したもので、300次元のベクトル同士の内積計算を256bitのハミング距離の計算で近似しようとしたのですが、同様の条件で計算したところ6~8倍ほどしか高速化できませんでした。おそらく距離の計算を高速化してもソートの部分がボトルネックになったためこのような結果になったのだと思われます。

締め

ソートアルゴリズムはたいていの場合はクイックソートが使われがちですが、条件次第ではその他のソートの方が便利なことがあります。faissのコードは専門家によって高度にチューニングされており、学べる点がとても多いため読んでみると面白いと思います。