SIGMOD會議是數(shù)據(jù)庫領(lǐng)域最具影響力的頂級國際學(xué)術(shù)會議之一，與VLDB和ICDE并稱為數(shù)據(jù)庫領(lǐng)域的三大頂級會議。PKU-DAIR實驗室發(fā)表于SIGMOD 2024的研究成果《CAFE: Towards Compact, Adaptive, and Fast Embedding for Large-scale Recommendation Models》榮獲了SIGMOD24 Honorable Mention for Best Artifact，該獎項每年僅授予至多三篇文章，旨在表彰那些在可復(fù)現(xiàn)性、靈活性和可移植性方面表現(xiàn)卓越的研究工作。

導(dǎo)讀

近年來，深度學(xué)習(xí)推薦模型（DLRM）中嵌入表的內(nèi)存需求不斷增長，給模型訓(xùn)練和部署帶來了巨大的挑戰(zhàn)?，F(xiàn)有的嵌入壓縮解決方案無法同時滿足三個關(guān)鍵設(shè)計要求：高內(nèi)存效率、低延遲和動態(tài)數(shù)據(jù)分布的適應(yīng)性。本文提出了CAFE，一種緊湊、自適應(yīng)和低延遲的嵌入壓縮框架，可以同時滿足上述要求。CAFE的設(shè)計理念是動態(tài)地為重要的特征（稱為熱特征）分配更多的內(nèi)存資源，為不重要的特征分配更少的內(nèi)存。在CAFE中，我們提出了一種快速且輕量級的草圖數(shù)據(jù)結(jié)構(gòu)，名為HotSketch，用于捕獲特征重要性并實時識別熱特征。對于每個熱特征，我們?yōu)槠浞峙湮ㄒ坏那度耄?/span>對于非熱門特征，我們使用哈希嵌入技術(shù)允許多個特征共享一個嵌入。在該設(shè)計理念下，我們進一步提出了多級哈希嵌入框架來優(yōu)化非熱門特征的嵌入表。我們從理論上分析了HotSketch的準(zhǔn)確性，并分析了模型的收斂性。實驗表明，CAFE顯著優(yōu)于現(xiàn)有的嵌入壓縮方法，在10000倍的壓縮比下，在Criteo Kaggle數(shù)據(jù)集和CriteoTB數(shù)據(jù)集上的測試AUC分別提高了3.92%和3.68%。

一、問題背景

近年來，嵌入技術(shù)廣泛應(yīng)用于數(shù)據(jù)庫領(lǐng)域，其中深度學(xué)習(xí)推薦模型（DLRM）是嵌入技術(shù)最重要的應(yīng)用之一。典型的DLRM將分類特征向量化為可學(xué)習(xí)的嵌入，然后將嵌入與其他數(shù)值特征一起喂給下游神經(jīng)網(wǎng)絡(luò)。最近，嵌入表的內(nèi)存需求隨著DLRM分類特征數(shù)量的指數(shù)級增長而增長，給各種應(yīng)用帶來了巨大的存儲挑戰(zhàn)。在本文中，我們重點關(guān)注壓縮超大規(guī)模DLRM的嵌入表。DLRM有離線訓(xùn)練和在線訓(xùn)練兩種范式。離線訓(xùn)練時，預(yù)先收集訓(xùn)練數(shù)據(jù)，并在整個訓(xùn)練過程后部署模型以供使用。在線訓(xùn)練中，訓(xùn)練數(shù)據(jù)是實時生成的，模型同時更新參數(shù)并服務(wù)請求。本文主要針對在線訓(xùn)練的場景，因為在線訓(xùn)練的難度更大，且在線訓(xùn)練的壓縮方法可以直接應(yīng)用于離線訓(xùn)練。在線訓(xùn)練中的嵌入壓縮有三個重要的設(shè)計要求：

• 高內(nèi)存效率。為了在內(nèi)存限制內(nèi)保持模型質(zhì)量，一種方案是使用分布式實例，但它們會帶來大量的通信開銷。此外，嵌入表的訓(xùn)練和部署通常發(fā)生在存儲容量較小的邊緣或終端設(shè)備上，使得內(nèi)存問題更加嚴(yán)重。另一方面，由于模型質(zhì)量直接影響利潤，即使AUC（ROC 曲線下面積）0.001的微小變化也會帶來相當(dāng)大的對利潤的影響。因此，我們需要保持模型質(zhì)量的內(nèi)存高效壓縮方法。
• 低延遲。延遲是服務(wù)質(zhì)量的關(guān)鍵指標(biāo)，因此嵌入的壓縮方法必須足夠快，以避免引入顯著的延遲。
• 對數(shù)據(jù)分布變化的適應(yīng)性。在線訓(xùn)練中，數(shù)據(jù)分布并不像離線訓(xùn)練那樣固定，一般會隨著時間動態(tài)變化。我們需要適合在線訓(xùn)練的自適應(yīng)壓縮方法。

現(xiàn)有的行壓縮方法可以分為基于哈希和自適應(yīng)方法兩類?；诠５?/span>方法利用簡單的哈希函數(shù)將特征映射到具有沖突的嵌入中[1]，盡管它們簡單易用，但預(yù)先確定的哈希函數(shù)會扭曲特征的語義信息，導(dǎo)致模型精度大幅下降；集成特征頻率信息[2]可以提高離線訓(xùn)練中基于哈希方法的模型質(zhì)量，但不能應(yīng)用于在線訓(xùn)練。自適應(yīng)方法在訓(xùn)練過程中識別重要特征以適應(yīng)在線訓(xùn)練[3]，但現(xiàn)有方法的壓縮率受到重要性分?jǐn)?shù)存儲的限制，導(dǎo)致內(nèi)存效率較低，且額外的檢查、采樣操作會增加整體延遲?？傊?，現(xiàn)有方法無法同時滿足所有的三個關(guān)鍵要求。

二、CAFE設(shè)計方案

1. 總體流程

我們提出的嵌入框架CAFE能夠同時滿足高內(nèi)存效率、低延遲、適應(yīng)數(shù)據(jù)分布動態(tài)變化三個要求。CAFE的核心思想是動態(tài)區(qū)分重要特征（稱為熱特征）和不重要特征（稱為非熱特征），并為熱特征分配更多的資源。具體來說，根據(jù)我們的理論推導(dǎo)，我們使用梯度的L2范數(shù)定義特征的重要性得分。我們觀察到，在大多數(shù)訓(xùn)練數(shù)據(jù)中，特征重要性遵循高度傾斜的分布，其中大多數(shù)特征的重要性得分較小，只有一小部分熱門特征非常重要。如圖1所示，Criteo Kaggle數(shù)據(jù)集和CriteoTB數(shù)據(jù)集中的特征重要性分布分別近似參數(shù)1.05和1.1的Zipf分布。因此，CAFE為熱點特征的嵌入分配更多的內(nèi)存，為非熱點特征的嵌入分配更少的內(nèi)存，從而在嵌入表的內(nèi)存使用量不變的情況下顯著提高模型質(zhì)量。 圖2是CAFE框架的示意圖，在CAFE 中，我們提出了一種新穎的HotSketch算法，用于實時記錄top-k熱特征的重要性。在每次訓(xùn)練迭代中，我們首先在HotSketch里查詢輸入樣本中的每個特征，重要性分?jǐn)?shù)超過閾值的視為熱門特征；然后我們分別查找熱門特征和非熱門特征的嵌入，我們給熱門特征分配唯一的嵌入，而讓多個非熱門特征共享一個嵌入；隨著訓(xùn)練數(shù)據(jù)分布的變化，熱門特征和非熱門特征之間可以互相遷移；在神經(jīng)網(wǎng)絡(luò)反向階段，我們使用特征的梯度范數(shù)來更新特征重要性。

圖1. 特征重要性分布

圖2. CAFE示意圖

2. HotSketch算法

HotSketch算法旨在訓(xùn)練中識別高重要性分?jǐn)?shù)的熱門特征，這本質(zhì)上是在流數(shù)據(jù)中查找top-k頻繁項（特征）的問題。目前，Space-Saving[4]是解決該問題最受認(rèn)可的算法，它在排序的雙向鏈表中維護頻繁項，并使用哈希表來索引。然而，該哈希表不僅使內(nèi)存使用量增加了一倍，而且指針操作引起的大量內(nèi)存訪問導(dǎo)致時間效率低下。我們提出的HotSketch刪除了哈希表，同時仍然保持o(1)時間復(fù)雜度，并很好地繼承了Space-Saving的理論結(jié)果。

數(shù)據(jù)結(jié)構(gòu)：如圖3所示，HotSketch由多個桶組成，我們使用哈希函數(shù)將特征映射到桶中。 每個桶包含多個槽，其中每個槽都存儲一個特征及其重要性分?jǐn)?shù)。

插入：對于每個傳入特征，我們首先計算哈希函數(shù)來定位桶，并遇到三種可能的情況：(1) 該特征記錄在桶中，我們將新分?jǐn)?shù)添加到其重要性得分上；(2)桶中沒有記錄該特征且存在空槽，那么我們將該特征和分?jǐn)?shù)插入到空槽中；(3)桶已滿且未記錄該特征，那么我們找到得分最小的特征替換為插入的特征并添加分?jǐn)?shù)。

討論：HotSketch具有以下優(yōu)點：（1）HotSketch插入速度快，具有o(1) 時間復(fù)雜度，并且只有一次內(nèi)存訪問；(2) HotSketch內(nèi)存效率高，它不存儲指針，并且短暫冷啟動后不會出現(xiàn)空槽；(3)HotSketch對硬件友好，可以通過多線程和SIMD進行加速，從而實現(xiàn)數(shù)據(jù)并行性。

圖3. HotSketch示意圖

3. 遷移策略

在DLRM的在線訓(xùn)練過程中，特征重要性的分布隨著數(shù)據(jù)分布的變化而波動，這意味著在整個訓(xùn)練過程中熱點特征并不是恒定的。HotSketch已經(jīng)記錄了特征重要性，自然可以通過嵌入表之間的嵌入遷移來支持動態(tài)熱特征。我們設(shè)置了一個閾值來區(qū)分熱點特征，幾乎每個首次達(dá)到HotSketch的特征都被視為非熱門特征。當(dāng)非熱門特征的重要性得分超過閾值時，它會轉(zhuǎn)變?yōu)闊衢T特征，并且其嵌入將從共享表遷移出來，確保特征的表示在整個訓(xùn)練過程中保持平滑；相反，HotSketch中特征的重要性會隨時間進行衰減，當(dāng)熱門特征的重要性得分低于閾值時，它就變成非熱門特征，考慮到新遷移的非熱特征不再重要，因此我們簡單地忽略其原始的獨占嵌入，而使用共享嵌入。通過在HotSketch中設(shè)置合適的閾值，我們能夠保持適度的遷移頻率從而使模型適應(yīng)分布的變化，而不會對收斂和延遲產(chǎn)生負(fù)面影響。

4. 多級哈希嵌入

非熱特征的數(shù)量遠(yuǎn)超熱門特征，鑒于這些非熱特征的重要性得分也符合高度傾斜的Zipf分布，因此我們可以根據(jù)其重要性進一步分隔并分配不同的哈希嵌入表。如圖4所示，我們將非熱門特征劃分為不同重要性級別，并為它們分配來自多個表的不同數(shù)量的嵌入。簡單起見，我們關(guān)注2級哈希嵌入，其中非熱特征分為中等特征和冷特征。我們擴展了HotSketch的功能，以估計中等特征的重要性分?jǐn)?shù)。由于中等特征比冷特征更重要，我們賦予它們兩個不同哈希嵌入表中的兩個嵌入向量，而冷特征僅查找單個嵌入向量。

圖4展示了一個示例。(1)輸入特征f1、f2、f3中，f1的得分大于熱閾值，f2的得分大于中閾值，f3的得分低于閾值，因此它們分別被分為熱特征、中特征和冷特征；(2) 熱特征和冷特征查找嵌入向量的方式不變；（3）中特征從兩個哈希嵌入表中查找兩個嵌入向量，并通過池化過程獲得最終的嵌入。為了確保訓(xùn)練過程保持平滑，當(dāng)一個特征在中類和冷類之間遷移時，它總是從第一個嵌入表中檢索相同的嵌入向量。對于池化操作，在實踐中，我們發(fā)現(xiàn)嵌入的簡單求和效果很好，因為特征的嵌入向量總是在相同的方向上更新。多級哈希嵌入的設(shè)計基于這樣的觀察：唯一嵌入是一種沒有信息丟失的綜合表示，而對于哈希嵌入，涉及的嵌入數(shù)量越多，沖突越少，信息量越多。

圖4. 多級哈希嵌入示意圖

三、實現(xiàn)

我們將CAFE實現(xiàn)為基于PyTorch的插件嵌入層模塊，可以直接替換任何基于PyTorch的推薦模型中的原始Embedding模塊。

CAFE后端：我們用C++實現(xiàn)HotSketch算法以減少總體延遲，并使用PyTorch實現(xiàn)CAFE的其余部分。我們HotSketch中桶的數(shù)量設(shè)置為預(yù)先確定的熱門特征數(shù)量，每個桶有4個槽；我們對所有特征字段使用一個HotSketch結(jié)構(gòu)，而不是每個字段一個結(jié)構(gòu)，因為跨字段的熱門特征分布不明確，最好直接用重要性分?jǐn)?shù)來處理。

容錯性：我們將所有HotSketch的狀態(tài)注冊為CAFE的PyTorch模塊中的buffer，使其可以與模型參數(shù)一起保存和加載。這種簡單的設(shè)計不需要額外的修改，使帶有CAFE的DLRM能夠使用檢查點進行訓(xùn)練和推理。

內(nèi)存管理：考慮到HotSketch不是計算密集型任務(wù)，我們將整個HotSketch結(jié)構(gòu)放在CPU上。CAFE的嵌入模塊基于PyTorch運算符構(gòu)建，可以在任何支持PyTorch的加速器（包括 CPU、GPU）上運行。

四、實驗結(jié)果

我們用三個模型DLRM[5]、WDL[6]、DCN[7]在四個數(shù)據(jù)集Criteo，CriteoTB，Avazu，KDD12上做了實驗。我們將CAFE與Hash Embedding[1]、Q-R Trick[8]、AdaEmbed[3]、MDE[9]進行了比較。

1、模型質(zhì)量比較

我們用測試AUC和訓(xùn)練loss分別表示離線和在線場景下的模型質(zhì)量。如圖5所示，相比Hash Embedding、Q-R Trick、MDE方法，CAFE能夠在同樣的壓縮率下達(dá)到最好的模型質(zhì)量，且在訓(xùn)練過程中一直保持更低的訓(xùn)練loss。AdaEmbed是一個自適應(yīng)方法，在低壓縮率場景下有很好的效果，而CAFE不僅能夠達(dá)到更高的壓縮率，還能在低壓縮率下取得和AdaEmbed相當(dāng)或者更好的效果。

具體而言，所有方法中只有CAFE和Hash能達(dá)到極限的10000x壓縮率，而其他方法只能支持較小的壓縮率，例如AdaEmbed和MDE只能支持不到100x的壓縮率。DLRM模型里，在Criteo數(shù)據(jù)集上，與Hash和Q-R Trick相比，CAFE的測試AUC平均提高了1.79%和0.55%；在CriteoTB數(shù)據(jù)集上，改進分別為1.304%和0.427%； 在KDD12數(shù)據(jù)集上，改進分別為1.86%和3.80%。 考慮訓(xùn)練loss，與Hash和Q-R Trick相比，CAFE在Criteo數(shù)據(jù)集上降低了2.31%、0.72%，在CriteoTB數(shù)據(jù)集上降低了1.35%、0.59%，在Avazu數(shù)據(jù)集上降低了3.34%、0.76%。與AdaEmbed相比，CAFE在Criteo數(shù)據(jù)集上達(dá)到了幾乎相同的測試AUC和訓(xùn)練損失，在CriteoTB數(shù)據(jù)集上提升了0.04%的測試AUC和減少0.12%的訓(xùn)練損失，在KDD12數(shù)據(jù)集上提升了0.82%的測試 AUC，在Avazu數(shù)據(jù)集上的訓(xùn)練損失減少了0.83%。

圖5. 模型質(zhì)量對比

我們進一步查看了多級哈希嵌入的效果，如圖6所示，其中CAFE-ML表示CAFE結(jié)合多級哈希嵌入。在不同的壓縮比下，CAFE-ML始終比CAFE表現(xiàn)更好，測試AUC提高了0.08%，訓(xùn)練損失降低了0.25%。CAFE-ML在較小的壓縮比下表現(xiàn)尤其出色，在100倍壓縮比下幾乎沒有造成性能下降。這是因為CAFE-ML在小壓縮比下為多級哈希嵌入表分配了更多的內(nèi)存，使得中等特征的表示更加精確。

圖6. 比較多級哈希嵌入

2、吞吐量對比

我們在圖7中測試了每種方法的吞吐量，實驗是在壓縮率為10x的CriteoTB數(shù)據(jù)集上進行的。由于不同方法的數(shù)據(jù)加載時間和DNN計算時間是相同的，因此差異僅在于嵌入層的操作。與未壓縮的嵌入操作相比，Hash僅需要額外的模運算，因此是訓(xùn)練和推理中最快的方法。Q-R Trick只額外引入了哈希過程和嵌入向量的聚合，而MDE引入了矩陣乘法、但需要更少的內(nèi)存訪問來獲取嵌入?yún)?shù)，因此這兩個方法也有較高吞吐量。為了適應(yīng)在線場景，AdaEmbed和CAFE引入了嵌入的動態(tài)調(diào)整，即重新分配或遷移，從而導(dǎo)致吞吐量降低。AdaEmbed會定期采樣數(shù)千個數(shù)據(jù)來確定是否重新分配，帶來較大時間開銷；相比之下，CAFE在HotSketch中的計算開銷可以忽略不計。與AdaEmbed相比，CAFE的訓(xùn)練吞吐量提高了33.97%，推理吞吐量提高了1.20%。

圖7. 吞吐量對比

3、HotSketch評測

桶中槽數(shù)的影響（圖8a-b）：我們記錄使用不同槽數(shù)的HotSketch的召回率和吞吐量。實驗使用Criteo數(shù)據(jù)集上1000x壓縮率下的熱門特征數(shù)量作為k。在圖8a中，召回隨著存儲變大而增加。根據(jù)我們的理論推導(dǎo)（見論文推論3.5），給定參數(shù)1.05到1.1的Zipf分布，每個桶的最佳槽數(shù)位于11到21。因此，槽數(shù)為8和16比4和32表現(xiàn)出更好的召回率。圖 8b中所示的序列化插入（寫入）和查詢（讀?。┑耐掏铝考s為1e7，比DLRM本身的吞吐大得多?？紤]到我們在實踐中可以并行化操作，HotSketch只占訓(xùn)練和推理的一小部分。隨著槽數(shù)量的增加，吞吐量會下降，因為需要更多的時間在桶內(nèi)進行比較。權(quán)衡召回率和吞吐量，我們在實現(xiàn)中采用每個桶4個槽，此時的模型質(zhì)量已經(jīng)足夠好。

尋找實時top-k特征（圖8c-d）：我們進行實驗來評估HotSketch在線訓(xùn)練中尋找兩類實時熱門特征的性能：最新的top-k特征，以及前一個時間窗口中的top-k特征。這兩類top-k特征在在線訓(xùn)練過程中隨著數(shù)據(jù)分布的變化而變化，因此可以有效體現(xiàn)HotSketch適應(yīng)動態(tài)工作負(fù)載的能力。實驗在 Criteo 上進行，使用6天的在線訓(xùn)練數(shù)據(jù)，滑動窗口大小為0.5天。圖8c-d顯示了HotSketch在不同壓縮比下在線訓(xùn)練時的實時召回率，始終滿足>90%，這意味著它可以很好地跟上不斷變化的數(shù)據(jù)分布。

圖8. HotSketch評測

五、總結(jié)

在本文中，我們提出了CAFE，一種緊湊、自適應(yīng)且快速的嵌入壓縮方法，它滿足三個基本設(shè)計要求：高內(nèi)存效率、低延遲和對動態(tài)數(shù)據(jù)分布的適應(yīng)性。我們引入了一種輕量數(shù)據(jù)結(jié)構(gòu)HotSketch記錄特征的重要性分?jǐn)?shù)，它產(chǎn)生的時間開銷可以忽略不計，并且其內(nèi)存消耗明顯低于原始嵌入表。通過將獨占嵌入分配給一小組重要特征，并將共享嵌入分配給其他不太重要的特征，我們在有限的內(nèi)存約束內(nèi)實現(xiàn)了卓越的模型質(zhì)量。為了適應(yīng)在線訓(xùn)練過程中的動態(tài)數(shù)據(jù)分布，我們采用了基于HotSketch的嵌入遷移策略。我們通過多級哈希嵌入進一步優(yōu)化CAFE，創(chuàng)建更細(xì)粒度的重要性組。實驗結(jié)果表明，CAFE優(yōu)于現(xiàn)有方法，在Criteo和CriteoTB數(shù)據(jù)集上10000x壓縮比下，測試AUC分別提高了3.92%、3.68%，訓(xùn)練損失降低了4.61%、3.24%，在離線訓(xùn)練和在線訓(xùn)練中均表現(xiàn)出優(yōu)異的性能。

實驗室簡介

北京大學(xué)數(shù)據(jù)與智能實驗室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR實驗室）由北京大學(xué)計算機學(xué)院崔斌教授領(lǐng)導(dǎo)，長期從事數(shù)據(jù)庫系統(tǒng)、大數(shù)據(jù)管理與分析、人工智能等領(lǐng)域的前沿研究，在理論和技術(shù)創(chuàng)新以及系統(tǒng)研發(fā)上取得多項成果，已在國際頂級學(xué)術(shù)會議和期刊發(fā)表學(xué)術(shù)論文200余篇，發(fā)布多個開源項目。課題組同學(xué)曾數(shù)十次獲得包括CCF優(yōu)博、ACM中國優(yōu)博、北大優(yōu)博、微軟學(xué)者、蘋果獎學(xué)金、谷歌獎學(xué)金等榮譽。PKU-DAIR實驗室持續(xù)與工業(yè)界展開卓有成效的合作，與騰訊、阿里巴巴、蘋果、微軟、百度、快手、中興通訊等多家知名企業(yè)開展項目合作和前沿探索，解決實際問題，進行科研成果的轉(zhuǎn)化落地。