高維空間近鄰檢索的雙層組合量化GPU加速算法

鄧?yán)眍＃?涵，陳 靚，全成斌，趙有健

1(清華大學(xué) 計算機科學(xué)與技術(shù)系，北京 100084) 2(中國石油勘探開發(fā)研究院，北京 100083)

1 引 言

高維空間中海量數(shù)據(jù)的高效檢索問題，可以說是計算機視覺、信息檢索等應(yīng)用所面臨的基本問題之一.主流的多媒體檢索框架，都需要將音頻、圖像或文本等多媒體數(shù)據(jù)抽取一定特征，轉(zhuǎn)化為高維向量，通過在高維空間中檢索近似向量來查詢所需的多媒體內(nèi)容.

伴隨互聯(lián)網(wǎng)技術(shù)和多媒體技術(shù)的不斷發(fā)展，網(wǎng)絡(luò)空間內(nèi)的各類數(shù)據(jù)呈現(xiàn)出爆炸增長的趨勢.以社交網(wǎng)站Instagram為例，目前Instagram僅每日新增的分享圖片數(shù)就接近1億張，在線視頻直播服務(wù)的興起則更是產(chǎn)生了海量的音視頻數(shù)據(jù)信息，對現(xiàn)有的圖像、視頻和文本內(nèi)容分析技術(shù)的處理能力提出了更高的要求.隨著數(shù)據(jù)規(guī)模的持續(xù)增長，結(jié)合數(shù)據(jù)本身存在的維度高、分布不確定等特點，如何對超過十億規(guī)模、甚至更多的高維數(shù)據(jù)進(jìn)行快速、有效地檢索逐漸成為研究的重點和難點.

傳統(tǒng)的樹形結(jié)構(gòu)檢索、基于哈希的檢索，在經(jīng)過適當(dāng)?shù)母倪M(jìn)后對于百萬量級的高維特征庫，如SIFT1M、GIST1M等，已經(jīng)擁有較好的處理能力，但對于更高規(guī)模的數(shù)據(jù)，則往往只能借助大量服務(wù)器進(jìn)行分布式處理[1].隨著數(shù)據(jù)規(guī)模的不斷提升，基于向量量化(Vector Quantization，VQ)的檢索方法由于能夠盡可能小的劃分結(jié)果空間，在十億量級的數(shù)據(jù)庫上展現(xiàn)出了良好的性能，逐漸成為研究熱點[2-6].尤其是如何盡可能提升計算密度，在單臺計算服務(wù)器節(jié)點內(nèi)進(jìn)行高效檢索仍然有待探索.

隨著CPU的晶體管數(shù)增長速度放緩，GPU由于其SIMD模型具有計算密度高、可并行度高的特點，得到了長足發(fā)展.單張GPU的浮點運算性能目前已突破10Tflops，遠(yuǎn)遠(yuǎn)高出同等量級的CPU，在高性能計算中擁有廣泛的運用.

為了充分挖掘GPU的并行運算潛力，本文基于組合量化(Composite Quantization，CQ)思想提出了一種針對GPU優(yōu)化的大規(guī)模高維空間數(shù)據(jù)近似近鄰檢索算法，結(jié)合了組合量化算法與雙層索引樹結(jié)構(gòu)，在建立次級量化碼本時引入局部最優(yōu)化思想，提高了量化精度，改善了在數(shù)據(jù)分布不規(guī)律時系統(tǒng)的檢索性能.同時利用替代投影思想優(yōu)化數(shù)據(jù)預(yù)處理過程，大幅度減少了顯存占用，最終實現(xiàn)了高維空間海量數(shù)據(jù)的高效檢索.

在擁有十億條128維特征向量的SIFT1B測試集上，本文算法在單張GPU(NVIDIA TITAN Xp)上就取得了良好的檢索效率：在與原組合量化算法的CPU版本檢索精度基本不變(Recall@10=0.45，Recall@100=0.926)的情況下，平均單次檢索僅需0.91ms，獲得了超過50倍的加速比，性能提升顯著.

2 相關(guān)研究

文獻(xiàn)[1]總結(jié)了過去高維空間中近鄰檢索的相關(guān)算法，通常來說可以分為三大類：基于樹型劃分的方法、基于哈希的方法以及基于向量量化的檢索方法.

早期的高維數(shù)據(jù)檢索主要基于樹形結(jié)構(gòu)，例如K-D樹、R樹等經(jīng)典算法，通過劃分子空間逐層減小搜索狀態(tài)，在空間維度相對比較低時(假設(shè)查詢向量的維度K，數(shù)據(jù)量N，要求2K?N)，擁有不錯的查找效率；但一旦數(shù)據(jù)空間維度上升，由于算法本身的回溯特性，便會遭遇“維數(shù)災(zāi)難”，效率已經(jīng)被證明甚至低于暴力枚舉[7].

基于哈希的近鄰檢索算法核心思想是通過尋找合適的哈希方式將高維向量投影到低維空間，轉(zhuǎn)化為低維空間中的相似性檢索問題.其中的代表性算法如局部敏感哈希(Locality-Sensitive Hashing，LSH)[8]在百萬量級的數(shù)據(jù)集上取得了很好的效率，目前已經(jīng)在各類檢索場景中被廣泛使用[9，10].其主要缺點在于根據(jù)Johnson-Lindenstrauss定理，隨著誤差范圍的減小，算法的空間復(fù)雜度將急劇上升.十億量級的數(shù)據(jù)規(guī)模下LSH保持較高查詢精度所需要的內(nèi)存遠(yuǎn)遠(yuǎn)超出可承受的范圍.

向量量化(Vector Quantization，VQ)是一種數(shù)據(jù)壓縮方法.首先在目標(biāo)空間中確定一組向量基，通常也稱之為“碼本”(codebook)，其中的每一個基稱為一個“碼字”，由于往往通過聚類求出向量基，因此碼字通常對應(yīng)一個聚類中心(centroid).通過將待處理向量近似轉(zhuǎn)化為碼本中碼字的線性組合，實現(xiàn)數(shù)據(jù)的壓縮，相應(yīng)的組合系數(shù)也就是待處理向量量化后的編碼.

在向量量化算法(表1)中，關(guān)鍵是碼本的量化方法和數(shù)據(jù)索引存放的方法.在百萬量級上進(jìn)行檢索，量化碼本中只需要數(shù)千個碼字(centroids)即可.然而隨著數(shù)據(jù)量的不斷提升，當(dāng)面臨十億量級的數(shù)據(jù)時，為了取得較好的檢索效果往往需要上百萬個碼字，僅僅逐一比較查詢向量與碼字的距離的代價就已經(jīng)難以接受.

乘積量化(Product Quantization，PQ)算法[5]是Herve Jegou于2009年提出的一種具有里程碑意義的向量量化方法.假設(shè)原始向量屬于D維空間RD，PQ算法將原始空間劃分為M個M/D維子空間的笛卡爾積，分別在每個子空間內(nèi)獨立的進(jìn)行向量量化，得到碼本C1，C2，…，CM.若每個碼本包含K個碼字，則等價于在D維原始空間中擁有一個含有KM個碼字的碼本.對于數(shù)據(jù)集中的任一向量X，將其同樣劃分為M段，每一段在相應(yīng)子空間內(nèi)的碼本上量化至最近的碼字Ck(ik),k=1,2,…,M，故可以將其近似表達(dá)為：

x≈[C1(i1)，C2(i2)，…，CM(iM)]

(1)

對于查詢向量q，其與數(shù)據(jù)集中任一向量X的距離可以表示為：

‖q-x‖2≈‖q-[C1(i1),C2(i2)，…，CM(iM)]‖2

(2)

其中qm表示查詢向量在相應(yīng)子空間內(nèi)的分量.相應(yīng)的，在查詢時，原本需要與KM個維度為D的碼字比較大幅降低為與M·K個維度為D/M的碼字作比較，即復(fù)雜度由O(KM·D)降低為O(K·O).因此通過在預(yù)處理時建立合適的查找表可以極大改善檢索效率.

PQ算法的改進(jìn).PQ算法在大規(guī)模數(shù)據(jù)集上表現(xiàn)出了突出的性能，受到乘積量化思想的啟發(fā)，近年來相繼提出了優(yōu)化的乘積量化(Optimized Product Quantization，OPQ)[3]、局部優(yōu)化的乘積量化(Locally Optimized Product Quantization，LOPQ)[4]等，通過旋轉(zhuǎn)坐標(biāo)系在全局或局部尋求最優(yōu)劃分，提升了算法的準(zhǔn)確度.

表1 基于向量量化的算法總結(jié)Table 1 Summerize of vector quantization

由于PQ對空間的正交劃分屬于有損壓縮，同樣有許多研究試圖針對這一問題做出改進(jìn)：例如加法量化(Additive Quantization，AQ)[11]試圖通過在原始空間多次量化，避免PQ算法中由于分段量化帶來的信息損失；堆疊量化(Stacked Quantization)[6]則通過迭代量化殘差的方式減小量化誤差，但相應(yīng)地帶來了更大的計算開銷.

Wieschollek等人進(jìn)行了利用GPU加速乘積量化的嘗試[12]，通過設(shè)立兩級索引結(jié)構(gòu)，減小了搜索空間，降低顯存的訪問次數(shù)；并且利用替代投影的預(yù)處理方式極大減少了數(shù)據(jù)量，通過良好的預(yù)處理，使得10億條SIFT向量構(gòu)建的索引庫能被存放在一張GPU內(nèi).但由于乘積量化本身屬于有損壓縮，加上替代投影帶來的信息損失，使得檢索的準(zhǔn)確度受到影響.

3 本文算法

本節(jié)將針對過往算法存在的問題，提出適應(yīng)于GPU運算的基于雙層索引結(jié)構(gòu)的組合量化(Bilayer Composite Quantization，BCQ)的算法流程，并給出具體的說明.

3.1 算法流程與框架

基于向量量化的高維數(shù)據(jù)檢索，都遵循共同的基本算法框架：

1)量化.通過算法指定的目標(biāo)函數(shù)訓(xùn)練碼本，并通過指定的量化方式，將待檢索數(shù)據(jù)庫中的向量一一量化至特定的碼字上.

2)索引.將碼字以一定的結(jié)構(gòu)進(jìn)行索引，方便快速查找，通常采用的是非對稱倒排索引(IVFADC)或倒排多索引(Inverted Multi-Index)結(jié)構(gòu).

3)查詢.對于待查詢向量q，將其依照同樣的步驟量化至一個特定的碼字上，取出碼字對應(yīng)的索引列表，視為近似最近鄰的候選列表.

4)重排序.將候選列表的原始向量取出，一一計算待查詢向量q與其的原始距離，并根據(jù)距離排序，獲得最終的近似最近鄰列表.

下面我們將分步介紹BCQ算法的原理與流程.

3.2 組合量化的基本原理

組合量化(Composite Quantization，CQ)算法[13]是Wang Jingdong等人在PQ算法的基礎(chǔ)上提出的一種新量化方式.

PQ算法存在的一個基本缺陷是，子空間的劃分必須獨立且相互正交，難以刻畫更多數(shù)據(jù)分布.盡管諸如OPQ等算法通過旋轉(zhuǎn)坐標(biāo)系，能在一定程度上降低量化誤差，但不能解決對于子空間劃分的約束性過強的根本問題.

因此組合量化選擇舍棄笛卡爾積，采用更直觀的方式將向量近似量化為一組向量基的加和：

x≈C1(i1)+C2(i2)+…+CM(iM)

(3)

圖1 不同量化算法的空間劃分Fig.1 Space partition of different quantization

圖1所示為不同量化算法對空間的分割情況.由圖中可以看出，與PQ單純將空間正交劃分相比，通過將向量以加和的方式解析，CQ算法能夠在碼本大小相同的情況下，更好的擬合數(shù)據(jù)本身的分布情況，減少錯誤劃分的數(shù)據(jù)點.

根據(jù)式(3)采用的量化方式，求查詢向量q與數(shù)據(jù)集中任一向量X的距離可以表示為：

(4)

(5)

然后通過輪流最優(yōu)化迭代求解.

3.3 雙層索引與局部最優(yōu)化

然而原始的組合量化算法同樣仍面臨一些未解決的問題，例如：

1)檢索時間的線性增長.隨著數(shù)據(jù)規(guī)模增長，子空間劃分?jǐn)?shù)M和每個子空間內(nèi)的碼本所需的碼字?jǐn)?shù)K都需要相應(yīng)增長，否則帶來的后果就是倒排索下，每個編碼index對應(yīng)的向量數(shù)線性增長.而所有的量化方法，在將查詢向量量化后的重排序過程都需要線性掃描對應(yīng)編碼Index下索引的所有向量，并一一計算實際距離.在CQ算法的原始論文測試數(shù)據(jù)中也表明，在10億向量的SIFT 1B數(shù)據(jù)庫中進(jìn)行單次檢索所需的平均時間超過了在100萬向量的SIFT 1M數(shù)據(jù)庫中進(jìn)行檢索的100倍以上，需要秒級檢索時間.

2)空間分布的非均一性.盡管CQ相較于PQ與OPQ而言，擁有更靈活的空間分割方法.然而事實上在個子空間中，數(shù)據(jù)的分布性質(zhì)又可能各有不同.

為了解決這一問題，本文算法采用雙層索引樹結(jié)構(gòu)拓展了原有的組合量化算法.

在索引樹首層采用向量量化，不妨設(shè)碼本中的碼字?jǐn)?shù)量為ω；在每個碼字(Centroid)所劃分的Voronoi子空間內(nèi)，同樣的我們考慮局部子空間內(nèi)的分布最優(yōu)化，再次進(jìn)行組合量化，得到M個量化碼本，大小為K.假設(shè)單層索引量化需要將量化粒度控制在同樣粗細(xì)，則需要總空間劃分?jǐn)?shù)ω·KM.由于在檢索時需要對最終候選列表進(jìn)行重排序，因此碼本擴(kuò)大ω倍即檢索空間縮小ω倍，即量化時間由O(K·M)增長為O(ω+K·M)，但重排序時間由O(T)縮減至O(T/ω).

由于通常情況下

(6)

因此取

(7)

時，可以獲得接近ω倍的加速比.

同時，由于建立雙層索引結(jié)構(gòu)，在二級子索引時，可以分別在每個子空間內(nèi)采取局部最優(yōu)化策略.LOPQ算法正是利用這一策略取得了比OPQ算法更高的準(zhǔn)確率[4].

與傳統(tǒng)PQ算法一樣，為了確保在子空間切分后，不會遺漏有效候選向量，我們在此處還采用“多分配策略”[5]，在索引樹的多個分支中同時查找(見圖2).實踐表明多分配策略在量化檢索中能有效的提高召回率.

圖2 雙層組合量化與多分配策略Fig.2 Bilayer Composite quantization and multi assign strategy

3.4 替代投影與重排序

除了數(shù)據(jù)分布不均一外，量化查詢另一信息損失過程在于查詢向量到量化中心的距離，并不直接等于查詢向量到候選向量的距離.傳統(tǒng)的量化檢索算法在完成對向量的量化、索引樹的建立后，在最終查詢時需要將完整的向量庫讀入內(nèi)存，以便在獲取量化編碼后，能迅速獲取對應(yīng)量化編碼的倒排索引列表中的所有候選向量的原始值，然后計算查詢向量與候選向量的實際距離，并完成重排序.

假設(shè)倒排索引中每個候選列表的平均長度為L，

(8)

則重排序所需的復(fù)雜度為O(L·D·m)，其中m為多分配系數(shù)，D為向量維數(shù)；而空間復(fù)雜度為O(N·D).

以SIFT 1B為例，10億條128維浮點數(shù)向量，N=1010，D=128×4 bytes，其所需內(nèi)存為92GB，而通常GPU的顯存在10G左右，無法在單張GPU內(nèi)完成運算.因此我們還需要基于GPU的架構(gòu)優(yōu)化最終的檢索與重排序流程.

一個直觀的想法是，通過盡可能的預(yù)處理，減少所需存儲的數(shù)據(jù)維數(shù).傳統(tǒng)的量化在尋找到向量的最近碼字后，是以查詢向量到碼字的距離近似實際距離，Wieschollek提出可以用向量的投影點作為替代：

圖3 替代投影示意圖Fig.3 Substitute Project

如圖3所示，需要計算y到x距離，以下記為d(y,x).傳統(tǒng)上，我們在量化時以y到最近的碼字(不妨設(shè)為Cj)的距離d(y,Cj)代替d(y,x)，根據(jù)三角關(guān)系其誤差上界為d(Cj,x).

如果將x投影至最近的兩個碼字連接成的直線上得到x′，采用用d(y,x′)代替d(y,x)，由Δyxx′三邊關(guān)系可知其誤差上界為δ，且δ必定小于d(Cj,x).即：

d(y,x)≈d(y,x′)=b2+λ2·c2+λ(a2-b2-c2)

由于a，b，c均與查詢向量無關(guān)，可以在預(yù)處理時確定，故對于每個數(shù)據(jù)點x，無論其原始維度為多少，僅需保存相應(yīng)的碼字編號i、j，及其投影到碼字連線上的點的位置偏移λ.空間復(fù)雜度驟降至O(N·M)，M為量化子空間劃分?jǐn)?shù).此時處理SIFT 1B所需顯存已經(jīng)降低至5G，可以在單張GPU上實現(xiàn).

4 實驗及數(shù)據(jù)分析

基于上一章中給出的優(yōu)化后的雙層索引組合量化算法，本章對相應(yīng)檢索系統(tǒng)進(jìn)行了實現(xiàn).通過對算法檢索準(zhǔn)確率、運行效率的測試，以及與其他相關(guān)研究的對比，分析了算法的性能.

4.1 實驗環(huán)境

硬件環(huán)境：

CPU：Intel Xeon E5-2620V4 @ 2.2GHz，8核16線程

RAM：128G DDR4 2400

GPU：NVIDIA TITAN Xp，12G顯存，3072個CUDA核心

測試數(shù)據(jù)集(表2)：

表2 測試數(shù)據(jù)集Table 2 Test Dataset

包括標(biāo)準(zhǔn)的公開數(shù)據(jù)集SIFT 1、SIFT 1B，分別含有100萬條與10億條128維SIFT特征向量，以及包含有5000萬條384維GIST特征向量的GIST 50M數(shù)據(jù)集.

4.2 實驗結(jié)果及分析

本節(jié)實現(xiàn)了上文描述的雙層索引組合量化算法(以下簡稱BCQ)，并在數(shù)據(jù)集上測試后，與過往的CPU量化算法如OPQ、CQ以及GPU量化算法PQT進(jìn)行對比分析.

為了確保測試數(shù)據(jù)的可比較性與對比的公平性，與各算法的對比測試均保持量化編碼長度為128bits，重排序候選鏈表長度為10000，并計算10000次查詢的平均檢索時間.

1)算法性能測試

表3 SIFT 1M數(shù)據(jù)集測試結(jié)果Table 3 Performance on SIFT 1M

從表3中可以看出：對于SIFT 1M數(shù)據(jù)集，目前的主流算法均取得了較高的召回率.檢索速度也遠(yuǎn)遠(yuǎn)超出諸如視頻圖像的實時檢索(30 frame/s，即33.3ms/query)等應(yīng)用的需求.說明百萬數(shù)量級的、較低維度的特征檢索對主流的量化檢索算法已經(jīng)不構(gòu)成挑戰(zhàn).

表4 GIST 50M數(shù)據(jù)集測試結(jié)果Table 4 Performance on GIST 50M

從表4中可以看出，隨著特征維度的升高，各類量化檢索方法的準(zhǔn)確度都有所下降，且PQT算法的降幅十分明顯，其檢索精度已經(jīng)基本不具有可用性.此時BCQ算法的精度仍然較為良好，達(dá)到甚至超出了CPU算法的準(zhǔn)確率.

綜合表4表5不難發(fā)現(xiàn)，隨著數(shù)據(jù)量的迅速增長，數(shù)據(jù)點在高維空間中的分布逐漸稠密，局部子空間最優(yōu)化的量化方法由于對空間結(jié)構(gòu)的更優(yōu)分割，在檢索準(zhǔn)確度上的優(yōu)勢便越發(fā)明顯.

表5 SIFT 1B數(shù)據(jù)集測試結(jié)果Table 5 Performance on SIFT 1B

上文已經(jīng)提到，目前互聯(lián)網(wǎng)空間中每日新增的多媒體數(shù)據(jù)都是以數(shù)億量級增長，因此當(dāng)我們將目光聚焦在十億及以上量級中的高維數(shù)據(jù)檢索時，原有的CQ算法盡管準(zhǔn)確率高，但運行時間過長，在SIFT 1B中單次檢索需要242.3 ms，遠(yuǎn)遠(yuǎn)不能滿足各類多媒體檢索實時應(yīng)用(即>30 frame/s)的需求；PQT算法盡管檢索速度取得了飛躍，但召回率往往不足80%甚至更低，難以達(dá)到實際應(yīng)用的要求.

而基于GPU實現(xiàn)的BCQ算法由于充分挖掘了GPU的并行計算能力，即使在10億量級的數(shù)據(jù)庫上仍然取得了平均單次檢索時間0.91ms的極高效率，相較于CPU算法有了大幅度的提高，與OPQ算法相比也有50倍左右的加速比，足以應(yīng)對各類多媒體信息處理系統(tǒng)中的實時要求；同時相較于現(xiàn)有的GPU算法，在運算速度相近的條件下，召回率大幅提高，近似保持了CPU算法的準(zhǔn)確度，使得算法在真實應(yīng)用場景下具有更好的可用性.

2)算法穩(wěn)定性測試

相應(yīng)的，我們還測試了BCQ算法在不同編碼長度下的表現(xiàn).編碼長度越短，所需要的計算時間越短，但相應(yīng)的檢索精度也有所下降.

圖4 在SIFT 1B數(shù)據(jù)集上的召回率曲線Fig.4 Recall curve on SIFT 1B

從圖4中我們可以看出，相較于同樣是基于GPU的PQT算法，BCQ算法從128bit編碼到64bit編碼的召回率降幅明顯更??；且即使在64bit編碼下，BCQ(64bit)的檢索精度依然高于PQT(128bit)，可見算法的魯棒性十分良好.

5 結(jié) 論

本文提出了一種基于GPU優(yōu)化的高維空間組合量化高效檢索算法.針對過往算法運行速度較低的問題，在組合量化的基礎(chǔ)上，通過引入雙層索引結(jié)構(gòu)及投影量化方法，大幅提高了檢索效率，在單張GPU上對10億量級的高維數(shù)據(jù)實現(xiàn)了檢索速度超過1000fps(單次檢索0.91ms)、召回率為92.6%的高效率、高精度檢索.在保持原有CPU算法經(jīng)度的情況下獲得了超過50倍的加速比；較已有的GPU算法在運算速度相近的情況下提升了20%的召回率，達(dá)到了0.926，大大改善了原有算法的可用性，為互聯(lián)網(wǎng)空間海量多媒體數(shù)據(jù)的各類應(yīng)用打下了良好的基礎(chǔ).