この記事はレコメンドエンジン連載の第3回目になります。前回までの記事はこちらを御覧ください。
レコメンド#2 Sparkで機械学習モデルを高速分散推論させる
はじめまして、Marketing Solution Division所属の野尻と申します。19年度にARISEに新卒入社してから約1年間レコメンドエンジンの開発を担当しています。
今回は商品間の類似距離を計算する際に近似近傍探索×PySparkを用いることで、大量の商品に対しての計算時間を当初の12時間から50分まで、大幅に削減したお話をします。
- 背景と課題
- 最近傍探索について
- 近似近傍探索について
- 転置インデックスについて
- 直積量子化について
- Faissの利用法
- Index作成
- 学習
- 探索
- その他の注意点
- FaissとGPUを使用するまでの経緯
- まとめ
背景と課題
第2回とほぼ同じ内容ですが、簡単にご説明します。
本シリーズはECサイトのお客様に「レコメンドエンジン導入」の支援をしたときの話です。
レコメンドエンジンを導入するにあたり、大きく2つの課題を抱えておりました。
- 数千万から数億件と、商品数が多いこと
- 商品の入れ替わりが激しい為、大半の商品がコールドスタート状態であること
コールドスタートについては「商品タイトルと商品説明文」を用いて自然言語処理を行い（※第1回参照）、fastTextを用いて商品を数百次元のベクトル空間に埋め込むことで対処しました（※第2回参照）。
しかし、数百次元のベクトルを持つ数千万件の商品の類似度組み合わせ計算を行う際に、全ての組み合わせ計算を行うとメモリに乗り切らず、処理速度も現実的な時間に収まらないという課題を抱えていました。
最近傍探索について
本題である近似近傍探索についてお話する前に、最近傍探索についてご説明させていただきます。
最近傍探索とは、とあるベクトルに対しもともと保持していたN個のベクトルから1番距離が近いベクトルを探索する問題のことです。
例としては、「アイドリッシュセブン」「アイドルマスターシャイニーカラーズ」「鬼滅の刃」といったコンテンツのタイトルがベクトル空間に埋め込まれていた場合に、「アイドリッシュセブン」に対して近いタイトルのコンテンツを探したい！という場合に用いられます。
その際に使用される距離は、2点間の直線距離であるユークリッド距離や、ベクトル同士がどれぐらい同じ角度を向いているかを計算するCosine類似度が用いられることが多いです。
最近傍、つまり自身と1番近いベクトルを探索する際にはN個のベクトルとの距離を逐一計算する必要がありますが、その際にネックとなるのがベクトルの個数の増加による計算時間の増加です。
ベクトルが2つだった場合は2×2、3つに増えた場合には3×3…といったように2乗オーダーで増加して行くため、先程の課題で挙げたように1000万商品に対しては1000万×1000万で100兆、1億商品に対しては、1億×1億で1京と爆発的に増えていくため、現実的でない計算時間が必要となってしまいます。
そのため、本プロジェクトでは次項でご説明する近似近傍探索を用いることで、現実的な計算時間に収めることに成功しました。
近似近傍探索について
近似近傍探索とは、最近傍探索に比べ厳密な近傍点を求めず、近傍点との距離を近似計算することで、計算時間を抑えながら近傍探索を行う手法です。
本プロジェクトでは近似近傍探索を高速化するための方法として3つ採用しています。
- 転置インデックス
- 直積量子化
- GPUの使用
また、これらの方法を組み込むに当たり、Faissというライブラリを使用しています。
Faissとは、Facebook社が開発を行っている近似近傍探索のOSSで、転置インデックスと直積量子化も実装されています。
コア部分がC++とCudaで実装されていますが、SWIGによりPythonインターフェースが用意されており、Python上でも動かすことが可能となっています。
公式URL: facebookresearch/faiss
転置インデックスについて
転置インデックスとは、全ベクトルを予め決められた数のクラスタに分割し、各クラスタ内で重心点を求めそれにクラスタのIDを紐付けたものを保持しておき、またそれに紐付いたクラスタ内の全ベクトルを別で保持しておく方法です。
この方法のメリットとして、検索速度の向上が挙げられます。とあるベクトルに対して全ベクトルとの距離を計算するのではなく、ひとまず各クラスタの重心点に対して距離を計算し一番近い重心点を持つクラスタを見つけ、そのクラスタ内でまた距離を計算するという処理により検索を効率よく行うことができます。
Faissにおける動作としては、元のベクトルデータをそのまま用いてk-meansを行っており、train部分で行われる1つ目の処理になります。k-meansに関してはネット上に文献が多くあるため、ここでは割愛します。また、実際のコードサンプルは後述します。
注意点として、trainの時点ではまだベクトルが登録されておらず、実際に登録されるのはその後のaddという処理の部分になります。こちらも詳細は後述したいと思います。
直積量子化について
直積量子化とは、分割したベクトルごとにベクトル量子化を行う方法のことです。
まずベクトル量子化とは、観測対象となるベクトルの集合を有限個の代表ベクトル（コードブック）で表現することです。具体的な手法はk-meansで求めた各クラスタの重心点（コードブック）を保持しておき、とあるベクトルがどのクラスタの重心点に近いかを求め、一番近いクラスタの重心点のIDをベクトルとして扱う手法です。
これを行うことで、観測ベクトルの表現に必要なデータサイズが、N次元 × 観測ベクトル数から1次元 × 代表点の個数まで圧縮できます。
k-meansのクラスタ数（K）は大きければ大きいほど保持する重心点が増え、その分とあるベクトルに対してより近い重心点を持つクラスタを提示することができ精度が上がりますが、大きくした分各重心点との近さを求める回数が増え量子化が遅くなってしまうので注意してください。
また、量子化後のベクトルは基本的に8bitのuchar（0~255の整数）で扱われます。Faissではnbitsというパラメータに該当しますが、現状だとGPUの場合だとnbits=8しか指定できないため注意してください（将来的に4, 5, 6, 8が利用できるように開発中のようです）
さらに元々のN次元ベクトルをM分割し、N/M次元に分割されたサブベクトルそれぞれに対してベクトル量子化する手法を直積量子化と言います。
直積量子化を行うことで、データ量の圧縮を行うことができ、その後の探索時間の大幅な短縮が可能となります。
何分割するかは、分割した後のサブベクトルの次元を指定することで決まります。元の次元数÷分割後のサブベクトルの次元数=分割数となり、この分割後のサブベクトルの次元数が大きいほど、データサイズは小さくなるが近似の精度が下がります。逆に次元数が小さいほど、近似の精度は上がりますがデータサイズは大きくなります。
例として、128次元の32bitの浮動小数点のベクトルだとデータサイズは4096bit（128×32）必要になりますが、直積量子化を行い128次元を4次元のサブベクトルに32分割し圧縮を行うと32次元×8bitで256bitとなり、1/16にデータサイズを圧縮することが可能です。
これらのパラメータは精度と処理時間に直結するため、許容できるデータサイズや精度、実行時間等とご相談しながら決定してください。
Faissの利用法
実際にFaissで利用するときのサンプルコードや注意点を説明していきます。
trainを行うときのサンプルコードは以下になります。
import faiss
res = faiss.StandardGpuResources()
dim = 128
nlist = 1024
M = 32
nbits = 8
metric = faiss.METRIC_L2
ivfpq_config = faiss.GpuIndexIVFPQConfig()
ivfpq_config.Index作成
インデックス作成部のパラメータについて説明していきます。
- res
resはGPUを使用するときのリソースの設定です。
一時的に割り当てるメモリのサイズ等を設定する必要がありますが、FaissではStandardGpuResources()というものが事前に用意されており、GPUのメモリサイズに応じて一時的に割り当てるメモリサイズを変更してくれる処理が実装されています。よって基本的にはStandartGpuResources()を指定しておけば良いと思います。
もし割り当てるメモリサイズを下げたい等のチューニングが必要な場合はご自身で設定する必要があります。
- dim
dimはindexに学習しようとしているベクトルの次元数を指定します。
そのまま次元数を入力するだけで大丈夫ですが、ベクトルの次元数における注意点として、後述する分割後の次元数で割り切れる次元数でないとともそも量子化を行えない点に注意してください。
計算時間を抑えたかったから分割後の次元数を32や64にしたかったのに、割り切れないから泣く泣く小さい次元数にして、計算時間が長くなってしまった…といったことにならないように要注意です！
- nlist
nlistは前述した転置インデックスのk-meansを行った後の分割後のクラスタ数です。
nlistは大きければ大きいほど1クラスタ当たりのベクトル数が減り、その分探索時に計算対象となるベクトルの数が減るため速度が早くなります。しかし、その分本当はクエリのベクトルに近かったのに検索対象とならないクラスタにいたために、探索結果にそのベクトルが出現しない確率が上がってしまいます。よってnlistの値は速度と精度のトレードオフになります。
検索対象とするクラスタ数は、後述のnprobeというパラメータで設定出来るため、そことの兼ね合いもあることにも注意してください。
nlistについては、公式のissueで通常これで良いという計算式が提示されているため、それを用いて算出すると良いと思います。
参考issue: How to select a suitable ‘nlist’? · Issue #112 · facebookresearch/faiss
- M
Mは分割後のサブベクトルの次元数です。作成するindexがGpuIndexIVFPQの場合は[1, 2, 3, 4, 8, 12, 16, 20, 24, 28, 32, 48, 56, 64, 96]の中から指定する必要があります。
直積量子化の項でお話しましたが、分割後のサブベクトルの次元数が大きいほど、データサイズは小さくなるが近似の精度が下がります。逆に次元数が小さいほど、近似の精度は上がりますがデータサイズは大きくなります。トレードオフなので、処理時間と相談の上決めると良いと思います。
またdim / Mが割り切れる値でないと動作しないため、注意してください。
- nbits
nbitsは分割後のサブベクトルのbit数です。
現状GPUで動かす場合は8bitのみとなっています。よって基本的にnbits = 8と指定してください。
（後述するivfpq_configでivfpq_config.interleavedLayout = Trueと設定すると4, 5, 6をnbitsに設定することが出来るようですが、開発中のためTrueにしないでくださいとの記述がありました。）
- metric
metricは距離計算時のアルゴリズムです。
選択出来る主なアルゴリズムは内積やL2ノルムです。この2つはほとんどのアルゴリズムでサポートしていると書かれているので、この2つから選択するのが一般的だと思います。
内積: metric = faiss.METRIC_INNER_PRODUCT
L2ノルム: metric = faiss.METRIC_L2
- ivfpq_config
ivfpq_configはindexに対するその他の設定です。
設定は4つあります。
1.userFloat16LookupTables
これは計算した残差の距離のテーブルに対して16bit浮動小数点を用いるかどうかの設定です。Mを大きくした結果、GPUのメモリにサブベクトルが乗り切らなかった場合にTrueに設定するようです。
2.usePrecomputedTables
これは事前計算されたテーブルを使用するかどうかという設定です。事前計算されたテーブル=コードブックのことだと思われます。デフォルトではFalseに設定されていますが、コードブックは使用するためivfpq_config.usePrecomputedTables = Trueと設定する必要がありそうです。
ただし、公式wikiにはIf you see cudaMalloc errors, disable precomputed tables.との記載があるため、GPUのメモリが足りずエラーが出てしまった場合にはFalseと設定する必要がありそうです。
3.interleavedLayout
こちらについては、GPUでは「開発中なので使用しないでください。」と記載されています。よってGPUを用いる場合はデフォルトのFalseのままで良いと思われます。
CPUの方で確認すると、Whether or not our index uses an interleaved by 32 layoutと記載されていました。
interleavingとは、メモリなどの記録媒体にデータを記録する際に予めデータを一定量で分割しておき、各媒体に同時に送ることで書き込み速度を向上させる手法のことです。
おそらく予めベクトルを32個つづで分割しておき、各メモリに同時に書き込むことでメモリへの書き込み速度を向上させるオプションなのではないかと考えられます。
この変数を用いたエラー文の実装のみされているようなので、今後の開発に期待ですね。
4.useMMCodeDistance
これはコードブックを使用せずに、強制的に距離計算を行われる設定で、This is for debugging purprosesと書いてあったためデバッグ用の設定のようです。
実際にTestGpuIndexIVFPQ.cppのコード内で使用されていました。
以上がindex作成時のパラメータの説明になります。
学習
次にtrain部の説明と注意点です。
index.train(vector_array)
trainメソッドの内部で実際のベクトルを用いて2つの処理を行っています。
1つ目の処理は、先述した転置インデックス用のk-meansです。ベクトルがnlistで指定した分のクラスタに分割され、各ベクトルで重心点のベクトルを保持します。
2つ目の処理は、ベクトルを登録するときの直積量子化に用いるコードブックの作成です。こちらは上記の処理で作成した転置インデックスの各クラスタの重心点とクエリベクトルの残差を直積量子化しています。残差を用いることで、近似精度を高める工夫を行っています。
注意点として、学習に用いるベクトルは32bit浮動小数点の二次元配列にのみ対応しているという点があります。たとえベクトルの数が1つの場合でも、二次元配列にする必要があるため注意してください。また、float64やfloat16でベクトルの配列を作成していた場合は、trainの手前でfloat32に変換する必要があります。
print(vec_array.dtype)
# float64
float32_vec_array = vec_array.astype(np.float32)
print(float32_vec_array.shape)
# (100, 128)
index.train(float32_vec_array)次にベクトルの登録です。
ベクトルの登録は、index.add()という関数で行います。train時に作成したコードブックを用いて、クエリのベクトルの直積量子化を行い、量子化後のベクトルを登録しています。また、登録する際には登録された順に0始まりで連番の整数IDが付与されます。
こちらも上記と同じように32bitの浮動小数点にしか対応していないため注意してください。
float32_vec_array = vec_array.astype(np.float32)
index.train(float32_vec_array)
index.add(float32_vec_array)また、各ベクトルに対してユーザーが独自で定義したIDを紐付けたい場合はindex.add_with_ids()を使用します。
ベクトルと同じ長さのIDの配列を用意し、ベクトルと一緒に登録します。
float32_vec_array = vec_array.astype(np.float32)
index.train(float32_vec_array)
index.add_with_ids(float32_こちらの注意点としては、64bitの整数のみ対応している点です。int32等だった場合は、add_with_idsの前にint64への変換が必要です。
int64_index_array = index_array.astype(np.int64)
探索
次に探索部についての説明と注意点です。
探索はindex.search()という関数で行えます。こちらがコードサンプルになります。
k = 100
index.nprobe = 50
print(type(queries))
# numpy.ndarray
print(queries.shape)
# (10000, 128)
distance, index = index.search(queries, k)では、パラメータについての説明です。
- k
kは探索する際に、クエリに近いベクトルをいくつ取得するか、というパラメータです。
k = 100と指定した場合は、探索結果が100個得られます。kは大きければ大きいほど多くの結果を取得しますが、その分探索時間が長くなる上に、探索後のデータサイズも大きくなるため、取得したい数と処理時間のバランスで設定してください。また公式wikiにはk <= 2048との記載があるため、2048以下で設定してください。
- nprobe
nprobeは、探索する際にクエリに対して一番近い重心点を持つクラスタを見つけた後、その周りに存在するクラスタに対して、何個のクラスタを探索しに行くか、というパラメータです。
重心点が一番近いクラスタ外にもクエリに対して近いベクトルを登録しているクラスタは存在する可能性があるため、大きければ大きいほどよりクエリに対して近いベクトルを発見できる可能性が高まります。しかし精度が上がる分探索時間も長くなるため、処理時間とのバランスでこちらも設定してください。存在するクラスタ数以上は見に行くことが出来ないため、nprobe <= nlistとなっているのと、kと同じくnprobe <= 2048と公式wikiに記載されているため、そちらにも注意してください。
次に実際の動作についてです。
まずindexにクエリベクトルが与えられるとコードブックを用いて量子化を行い、量子化後のベクトルと各クラスタの重心点との距離を計算します。そして一番近い重心点を持つクラスタを特定し、その重心点とクエリとの残差を計算します。そしてそのクラスタ内に登録されているベクトルの残差との距離を計算し、より距離が近いものを取得します。
返り値としては、距離とIDを近い順に格納した2つの配列が返却されます。
探索を行う際の注意点として、kで指定した分のベクトルが取得できなかった場合にはIDの配列に-1の値が挿入される仕様のため、ID配列内の-1は除外するといった処理が必要になります。
その他の注意点
その他の注意点として、faissのGPU用のindexを一度ファイルとして保存したい場合は、CPU用のindexに変換する必要があります。
converted_index_for_cpu = faiss.index_gpu_to_cpu(index)
faiss.write_index(converted_その後保存したインデックスを読み込んで、再びGPUにインデックスを載せたい！となった場合は、GPU用のインデックスに変換する必要があります。
cpu_index = faiss.read_index('hoge.faiss')
gpu_index = faiss.index_cpu_to_all_gpus(一般的なユースケースとして、ベクトルの学習・登録をGPUで行い、その後のAPI等に用いる際にいはCPUで動作させるといったケースがあると思われますが、その際にはこちらの方法で変換・保存を行ってください。
本プロジェクトではバッチで事前計算する必要があったため、複数のGPUマシンを用意し、PySparkを用いてFaissの探索を分散処理させています。こちらについては第2回の「Sparkで機械学習モデルを高速分散推論させる」という記事内で用いているfastTextをFaissに置き換えたような話になるため、そちらの記事をご参考ください。
FaissとGPUを使用するまでの経緯
ここまでFaissとGPUを使用した近似近傍探索について解説してきましたが、そもそもどのような経緯でFaissとGPUを使用したかを簡単にご説明したいと思います。
背景と課題で挙げたとおり、本プロジェクトでは数百次元のベクトルを持つ数千万件の商品の類似度組み合わせ計算を行う際に、全ての組み合わせ計算を行うとメモリに乗り切らず、処理速度も現実的な時間に収まらないという課題を抱えていました。
一番最初のアプローチとしては、ダメ元で全ベクトル×全ベクトルの距離をPySparkのUDFを用いて計算を行っていました。しかし、Sparkがメモリエラーで落ち、想定通り計算できませんでした。対向システムの関係上、バッチで事前計算しておく必要があったため、どのようにしたらレコメンドエンジンのバッチ処理として現実的な時間で計算が行えるかを考えました。
そこで出た案が以下の2つでした。
- 事前にk-meansを行いある程度クラスタリングした上で、各クラスタ内でベクトル同士の距離計算を行う
- Faiss等のライブラリを用いて近似近傍探索を行う
1については、同じく計算が終わることはなく途中で断念しました。
そして次にFaissを利用しました。
近似近傍探索を行うライブラリはいくつか存在しますが、
- GPUが利用できること
- ベクトルIDのマッピング機能があること
- インデックスファイルがシングルファイルなため、PySparkのUDFで使用しやすいこと
- バッチ検索が高速なため、PySparkのPandas UDFと相性が良かったこと
以上4点の理由からFaissを選択しました。
2のFaissのCPUによる近似近傍探索では処理を完了することはできたものの、約12時間の計算時間を要することが判明し、毎日レコメンドリストを連携する必要がある本プロジェクトにおいては実用的では無いと判断し、こちらも断念しました。
さらなる高速化を目指した結果、同じFaissでもGPUを用いて実装することで、実用的な時間に計算時間を収めることにチャレンジしました。
それまでMarketing Solution Division内ではGPUを使用した前例が存在せず、そもそも我々が利用できる基盤上にGPUを構築することができなかったため、基盤運用チームにGPUを利用できる環境を用意していただくところからのスタートとなりました。
その後環境を用意していただきGPUを用いてFaissを実行したところ、CPUでは約12時間必要としていた処理を、約50分に短縮することができました。
皆様のプロジェクトで、大規模データで近似近傍探索を短い計算時間で行いたいとなった場合は、FaissとGPUを用いることを一考してみてはいかがでしょうか。
まとめ
今回は、GPUで近似近傍探索を行うことで大規模なデータでも処理時間を大幅に削減することができたというお話でした。近年は大規模データを扱う機会が増え、計算時間にお困りの方も多いのではないかと思います。しかし、クラウド上でGPUを使用できるサービスも登場し、気軽に使用できる機会も増えたのではないでしょうか。プロジェクトの予算にもよるとは思いますが、もし計算時間でお困りの方がいらっしゃいましたら、GPUを用いての計算時間短縮を試してみてください。大幅な時間短縮になれば、CPUとGPUのコスト差を上回り、計算時間短縮&コスト削減になる可能性があります。
3回に渡ったレコメンドの記事は今回で最後となります。稚拙な文章ではありましたが、記事を読んでくださった皆様ありがとうございました。
