scannの実装を追う
katzscann
直積量子化を行うscannの論文を読んだあと実装を読んだがまったくわからなかったため実装を追ったメモを残す
ちなみに実装によくでてくるsoarは最近実装された論文
ぱっと見た感じMIPSに合わせていい感じ(AVQ的な)に2番目のクラスタを割り当てようって感じに見える
katzsearcherの作成例
searcher = scann.scann_ops_pybind.builder(normalized_dataset, 10, "dot_product")
.tree(num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
.score_ah(2, anisotropic_quantization_threshold=0.2)
.reorder(100)
.build()
目安
ScaNN performs vector search in three phases. They are described below:
- Partitioning (optional): ScaNN partitions the dataset during training time, and at query time selects the top partitions to pass onto the scoring stage.
- Scoring: ScaNN computes the distances from the query to all datapoints in the dataset (if partitioning isn't enabled) or all datapoints in a partition to search (if partitioning is enabled). These distances aren't necessarily exact.
- Rescoring (optional): ScaNN takes the best k' distances from scoring and re-computes these distances more accurately. From these k' re-computed distances the top k are selected.
All three phases can be configured through the ScaNNBuilder class. Before going into details, below are some general guidelines for ScaNN configuration.
このあたりで論文にないtree, score_ah, rescoreingなどが出てくる
ahはasymmetry hashingであるわけだが、Anisotropicはtreeのところにでてこないのか？
LSHにkmeans-tree？
katzと思ったら論文の最後にscannの実装は以下の2つを参考にしていることがわかる。
どちらもgoogleの論文。これを読む。
katzQuantization-based inner product (QUIP) search
MIPS Problemを量子化で対応しようとしている
Note that this approximation is ’asymmetric’ in the sense that only database vectors x are quantized, not the query vector q. One can quantize q as well but it will lead to increased approximation error. In fact, the above asymmetric computation for all the database vectors can still be carried out very efficiently via look up tables similar to [9], except that each entry in the kth table is a dot product between q(k) and columns of U(k).
asymmetry hashingも[Jégou, 2010]でのAsymmetric distance computationと認識して問題なさそう
これの6章がTree-Quantization Hybrids for Large Scale Searchとなっていてそれっぽい。実際scannでもtree()をするとTreeXHybrid~クラスが登場する。
The basic idea of tree-quantization hybrids is to combine tree-based recursive data partitioning with
QUIPS applied to each partition. At the training time, one first learns a locality-preserving tree such as hierarchical k-means tree, followed by applying QUIPS to each partition. In practice only a shallow tree is learned such that each leaf contains a few thousand points. Of course, a special case of tree-based partitioners is a flat partitioner such as k-means. At the query time, a query is assigned to more than one partition to deal with the errors caused by hard partitioning of the data. This soft assignment of query to multiple partitions is crucial for achieving good accuracy for high-dimensional data.
クエリはpartition centerの近くに割り当てられ、実験では全てで2000 partition, クエリに割り当てられるpartitionは100である。これらはbrute forceアプローチより高速と論文に書いてあるが、ここでいうbruto forceとは?? この論文内でbrute forceという単語が出てくるのはここだけなので、すべての検索対象についてasymmetricな距離の計算をしているという意味のはず。
が、となるとscannで出てくるscore_ahとscore_brute_forceの違いは？
→このときのbrute_forceは単純に性格な値を計算しているっぽい？
katzそもそも論文の位置づけとしては
[Jégou, 2010] → 直積量子化、目的関数はreconstruction error
[Guo, 2016] → 直積量子化、目的関数はMIPS
[Guo, 2020] → score awareな損失関数、scann
みたいな形か、、？
[Guo,2016]の場合は
[Wu,2017]はいまいち接点がわからず
↓の通りtree+score_ahでは明示的に指定しない場合residual_quantizationが使われそう
katztree, socre_ahなどは_factory_decoratorによって関数の引数をparamsフィールドに入れられ、各paramsの値を基にconfigが作られる。これをちゃんとみてなかったため、tree+score_ahが呼ばれた場合にresidual_quantizationがtrueに上書きされることを見落としていた。
def _factory_decorator(key):
"""Wraps a function that produces a portion of the ScaNN config proto."""
def func_taker(f):
"""Captures arguments to function and saves them to params for later."""
def inner(self, *args, **kwargs):
if key in self.params:
raise Exception(f"{key} has already been configured")
kwargs.update(zip(f.__code__.co_varnames[1:], args))
self.params[key] = kwargs
return self
inner.proto_maker = f
return inner
return func_taker
ah = self.params.get("score_ah")
bf = self.params.get("score_bf")
if ah is not None and bf is None:
if "residual_quantization" not in ah:
ah["residual_quantization"] = (
tree_params is not None and self.distance_measure == "dot_product")
そこからscannのコードを追っていく
︙
treeを設定していればTreeXHybridFactoryが呼ばれ、
更に追っていくと
partitionでは
hashingでは とそれぞれ学習コードが出てくるようである。
ここでpartitionのほうではデフォルトではAnisotropic VQはオフみたいだが、
ここでavqを指定することが可能。その後がよくわからない
kmeans-treeの各葉に対して中心を変更している？？
score_ahのほうは から
と学習コードが存在。
この中の
gmm.ComputeKmeansClustering(
chunked_dataset[i], opts.config().num_clusters_per_block(), ¢ers,
{.final_partitions = &subpartitions, .weights = weights}));`
は普通のkmeansでクラスタの中心を計算しているっぽいが、
それからしばらくした
ではkmeansで計算したresidualに対して再度割り当てるクラスを計算し直している、、？？
katzdebugしたいがmacでビルドができない、、
katz
3日前に最新のbazelに対応してた
katz
macで動かした記事もあるけどうまくいかない...
katz一旦諦めて、コードのbuildは見終わったのでsearchに行きたい