文 / 軟件工程師 Philip Sun，Google Research

人們要從大量的文學(xué)作品中進(jìn)行搜索，一般需要使用與標(biāo)題、作者或其他易于機(jī)器索引的標(biāo)準(zhǔn)完全匹配的查詢條件。那么使用 SQL 等語(yǔ)言的關(guān)系型數(shù)據(jù)庫(kù)可以輕松完成此類任務(wù)。

但是，如果想要查詢?nèi)?‘內(nèi)戰(zhàn)期間的詩(shī)歌’ 這類比較抽象的內(nèi)容，就無(wú)法再依賴如對(duì)比短語(yǔ)之間共同單詞數(shù)量等簡(jiǎn)單的 相似性指標(biāo)。如，“科幻小說(shuō) (Science Fiction)”與“未來(lái) (Future)”的關(guān)系比與“地球科學(xué) (Earth Science)”的關(guān)系更緊密，盡管，后者有一個(gè)相同單詞  Science，而前者沒有。

機(jī)器學(xué)習(xí) (ML) 極大地提高了計(jì)算機(jī)理解語(yǔ)言語(yǔ)義以及響應(yīng)這些抽象查詢的能力?，F(xiàn)代 ML 模型可以將文本和圖像等輸入轉(zhuǎn)換為嵌入向量，對(duì)高維向量進(jìn)行訓(xùn)練，使更相似的輸入更緊密地聚集在一起。因此，對(duì)于給定的查詢，我們可以計(jì)算其嵌入向量，并找到嵌入向量與查詢最接近的文學(xué)作品。通過(guò)這種方式，ML 已經(jīng)將以前難以指定的抽象任務(wù)轉(zhuǎn)變?yōu)閲?yán)格的數(shù)學(xué)任務(wù)。

然而，計(jì)算難題依然存在：對(duì)于給定的查詢嵌入向量，如何快速找到數(shù)據(jù)集中最接近的嵌入向量呢？嵌入向量集合通常太大，無(wú)法進(jìn)行窮盡式搜索，并且通常嵌入向量的維數(shù)很高。這使得很難對(duì)向量本身進(jìn)行修剪。

在我們的 ICML 2020 論文“Accelerating Large-Scale Inference with Anisotropic Vector Quantization”中，我們通過(guò)關(guān)注如何壓縮數(shù)據(jù)集向量來(lái)實(shí)現(xiàn)快速近似距離計(jì)算解決這一問(wèn)題，并提出了一種新的壓縮技術(shù)，與以前的工作相比，這項(xiàng)技術(shù)可以大大提高準(zhǔn)確率。我們?cè)诮陂_源的向量相似性搜索庫(kù) (ScaNN) 中應(yīng)用了這項(xiàng)技術(shù)，與其他向量相似性搜索庫(kù)相比，我們的性能高出兩倍（在 ann-benchmarks.com 上測(cè)得）。

• Accelerating Large-Scale Inference with Anisotropic Vector Quantization
  https:///abs/1908.10396

• 向量相似性搜索庫(kù) (ScaNN)
  https://github.com/google-research/google-research/tree/master/scann

基于嵌入向量的搜索技術(shù)能夠有效響應(yīng)依賴于語(yǔ)義理解的查詢，這比簡(jiǎn)單的可索引屬性的查詢更為高效。在這種技術(shù)中，機(jī)器學(xué)習(xí)模型被訓(xùn)練成將查詢和數(shù)據(jù)庫(kù)項(xiàng)映射到公共向量嵌入空間，使得嵌入向量之間的距離帶有語(yǔ)義含義，即相似項(xiàng)更接近。

如上所示，雙塔神經(jīng)網(wǎng)絡(luò)模型是一種基于嵌入向量的特定類型搜索，其中查詢和數(shù)據(jù)庫(kù)項(xiàng)通過(guò)兩個(gè)相應(yīng)的神經(jīng)網(wǎng)絡(luò)映射到嵌入空間。在此示例中，模型對(duì)假設(shè)的文學(xué)數(shù)據(jù)庫(kù)的自然語(yǔ)言查詢做出響應(yīng)

要用這種方法響應(yīng)查詢，系統(tǒng)必須先將查詢映射到嵌入空間。然后是最近鄰搜索 (Nearest Neighbor Search) 問(wèn)題：它需要在所有數(shù)據(jù)庫(kù)嵌入向量中找到最接近查詢的嵌入向量。定義查詢數(shù)據(jù)庫(kù)嵌入向量相似性的最常見方法之一是通過(guò)它們的內(nèi)積；這種最近鄰搜索被稱為最大內(nèi)積搜索 (Maximum Inner-Product Search，MIPS)。

數(shù)據(jù)庫(kù)的大小很容易達(dá)到數(shù)百萬(wàn)甚至數(shù)十億量級(jí)，這種情況下 MIPS 通常是推理速度的計(jì)算瓶頸，進(jìn)行窮盡式搜索不切實(shí)際。而使用近似的 MIPS 算法雖然會(huì)犧牲一些精度，但能大幅提高蠻力搜索速度。

MIPS 的幾種最新解決方案均基于壓縮數(shù)據(jù)庫(kù)項(xiàng)，可以超越蠻力搜索，迅速計(jì)算出它們的內(nèi)積近似值。這種壓縮通常通過(guò)學(xué)習(xí)量化 (Learned Quantization) 來(lái)完成，其中向量的 碼本 (Codebook) 基于數(shù)據(jù)庫(kù)訓(xùn)練得出，并用于近似表示數(shù)據(jù)庫(kù)元素。

• 幾種最新解決方案
  https://github.com/erikbern/ann-benchmarks#evaluated

先前的向量量化方案對(duì)數(shù)據(jù)庫(kù)元素進(jìn)行了量化，目的是最小化每個(gè)向量 x 及其量化形式 x? 之間的平均距離。盡管這是一個(gè)有用的指標(biāo)，但為此進(jìn)行優(yōu)化并不等同于優(yōu)化最近鄰搜索精度。論文背后的關(guān)鍵理論是，編碼的平均距離 越大 ，實(shí)際上 MIPS 精度也可能越高。

我們對(duì)結(jié)果的分析如下所示。假設(shè)我們有兩個(gè)數(shù)據(jù)庫(kù)嵌入向量 x₁ 和 x₂，并且必須將每個(gè)量化到兩個(gè)中心之一：c₁ 或 c₂。我們的目標(biāo)是將每個(gè) x_i 量化為 x?_i ，以使內(nèi)積 <q, x?_i> 盡可能接近原始內(nèi)積 <q, x_i>?？蓪⒋诉^(guò)程可視化為使 x?i 在 q 上的投影幅度盡可能類似于 x?_i 在 q上的投影幅度。在傳統(tǒng)的量化方法（左）中，我們將為每個(gè) x_i 選擇最近的中心，這將導(dǎo)致兩個(gè)點(diǎn)的相對(duì)排名不正確：<q, x?₁> 大于 <q,  x?₂>，但是實(shí)際上 <q, x₁> 小于 <q, x₂>！如果我們改為將 x₁ 賦給 c₁，將 x₂ 賦給 c₂，則會(huì)得到正確的排名。如下圖所示。

目標(biāo)是將每個(gè) _xi 量化為 x?_i =  c₁ 或 x?_i = c₂：傳統(tǒng)的量化（左）導(dǎo)致該查詢的 x₁ 和 x₂ 的排序不正確。盡管我們的方法（右）選擇的中心距數(shù)據(jù)點(diǎn)較遠(yuǎn)，但這實(shí)際上降低了內(nèi)積誤差、提高了精度

事實(shí)證明，方向和幅度同樣重要——即使 c₁ 與 x₁ 的距離超過(guò) c₂ 與它的距離， c₁ 在幾乎與 x₁ 完全正交的方向上偏離 x₁，而 c₂ 的偏離是平行的（對(duì)于 x₂，情況相同，但方向相反）。平行方向上的誤差在 MIPS 問(wèn)題中危害更大，因?yàn)樗鼑?yán)重影響高內(nèi)積，根據(jù)定義，高內(nèi)積是 MIPS 試圖精確估計(jì)的內(nèi)積。

基于這種認(rèn)知，我們會(huì)更加嚴(yán)厲地懲罰與原始向量平行的量化誤差。由于其損失函數(shù)的方向依賴性，我們將這種新型量化技術(shù)稱為 各向異性向量量化 (Anisotropic Vector Quantization) 。它能夠用較低內(nèi)積增加的量化誤差來(lái)?yè)Q取較高內(nèi)積的超高精度，這是創(chuàng)新的關(guān)鍵所在，也是其性能提升的源泉。

在上圖中，橢圓表示等損失的輪廓：在各向異性向量量化中，平行于原始數(shù)據(jù)點(diǎn) x 的誤差會(huì)受到更大的懲罰

各向異性向量量化使 ScaNN 可以更好地估計(jì)前 k 個(gè) MIPS 結(jié)果中的內(nèi)積，從而實(shí)現(xiàn)更高的精度。在來(lái)自 ann-benchmarks.com 的 glove-100-angular 基準(zhǔn)測(cè)試中，ScaNN 的性能優(yōu)于其他 11 個(gè)經(jīng)過(guò)精心調(diào)整的向量相似性搜索庫(kù)，在給定的精度下，每秒處理的查詢數(shù)量大約是第二快的庫(kù)的兩倍。*

Recall@k 是最近鄰搜索精度的常用指標(biāo)，用于衡量算法返回的 k 個(gè)近鄰中存在的真正的最近 k 個(gè)近鄰的比例：ScaNN（上面的紫色線）在速度-精度權(quán)衡的各個(gè)方面始終展現(xiàn)出優(yōu)異的性能

我們已經(jīng)將 ScaNN 在 GitHub 上開源。ScaNN 可以通過(guò) Pip 直接安裝，并具有用于 TensorFlow 和 Numpy 輸入的接口。有關(guān)安裝和配置 ScaNN 的更多說(shuō)明，請(qǐng)參見 GitHub 代碼庫(kù)。

通過(guò)修改向量量化目標(biāo)以與 MIPS 的目標(biāo)保持一致，我們?cè)谧罱徦阉骰鶞?zhǔn)上獲得了 SOTA 性能，這是基于嵌入向量的搜索性能的關(guān)鍵指標(biāo)。

各向異性向量量化是一項(xiàng)重要技術(shù)，我們今天主要討論的是通過(guò)優(yōu)化算法實(shí)現(xiàn)性能提升的一個(gè)例子，其最終目標(biāo)是實(shí)現(xiàn)提高搜索精度，而不是壓縮失真等中間目標(biāo)。

這篇文章是整個(gè) ScaNN 團(tuán)隊(duì)共同努力的成果：David Simcha、Erik Lindgren、Felix Chern、Nathan Cordeiro、Ruiqi Guo、Sanjiv Kumar 以及 Zonglin Li。我們還要感謝 Dan Holtmann-Rice、Dave Dopson 和 Felix Yu。