實作兩塔檢索功能,以便大規模產生候選項

Last reviewed 2025-01-16 UTC

本文提供參考架構,說明如何使用 Vertex AI 實作端對端雙塔候選項目生成工作流程。雙塔模型架構是強大的擷取技術,適用於個人化用途,因為這項技術會學習兩個不同實體 (例如網頁查詢和候選項目) 之間的語意相似度。

本文適用於資料科學家和機器學習工程師等技術人員,他們正在開發大規模推薦應用程式,且需要低延遲服務。如要進一步瞭解雙塔模型的建模技術、問題框架和資料準備作業,請參閱「使用 TensorFlow Recommenders 和 Vector Search 擴充深度檢索功能」。

架構

下圖顯示的架構可用於訓練雙塔模型,並分別部署每個塔,以執行不同的部署和服務工作:

訓練雙塔模型並分別部署每個塔的架構。

圖中的架構包含下列元件:

  • 訓練資料:訓練檔案儲存在 Cloud Storage 中。
  • 雙塔訓練:使用 Vertex AI Training 服務離線訓練合併的雙塔模型;每個塔會分別儲存,並用於不同工作。
  • 已註冊的查詢和候選項目塔:訓練塔後,每個塔會分別上傳至 Vertex AI Model Registry。
  • 已部署的查詢塔:已註冊的查詢塔會部署至 Vertex AI 線上端點。
  • 批次預測嵌入:註冊的候選項目塔會在批次預測工作中預先計算所有可用候選項目的嵌入表示法。
  • 嵌入 JSON:預測的嵌入會儲存到 Cloud Storage 中的 JSON 檔案。
  • ANN 索引:Vertex AI Vector Search 用於建立服務索引,並設定為近似最鄰近 (ANN) 搜尋。
  • 已部署的索引:ANN 索引會部署至 Vertex AI Vector Search 索引端點。

使用的產品

這項參考架構使用下列 Google Cloud 產品:

  • Vertex AI Training:全代管訓練服務,可讓您大規模訓練模型。
  • Vector Search:向量相似度比對服務,可讓您儲存、建立索引及搜尋語意相似或相關的資料。
  • Vertex AI Model Registry:中央存放區,可用於管理機器學習模型的生命週期。
  • Cloud Storage:適用於多種資料類型的物件儲存庫,成本低廉且沒有限制。 資料在 Google Cloud 內外都能存取,且會複製到多個位置,以便提供備援機制。 Google Cloud

用途

為滿足低延遲服務需求,大規模建議工具通常會部署至生產環境,做為兩階段系統,有時也會部署為多階段系統。第一階段的目標是生成候選項目,從大量候選項目中篩選出數百個相關項目,以供後續的篩選和排序作業使用。如要最佳化這項擷取作業,請考量下列兩項核心目標:

  1. 在模型訓練期間,瞭解要解決的問題或工作的最佳表示法,並將此表示法編譯為<query, candidate>嵌入。
  2. 在模型服務期間,快速擷取相關項目,以符合延遲時間規定。

下圖顯示兩階段建議系統的概念元件:

兩階段建議工具的概念元件。

在圖表中,候選人生成會篩選數百萬個候選人項目。排序功能會篩選出數百個候選項目,然後傳回數十個建議項目。

本文中的參考架構會訓練以雙塔為基礎的擷取模型。在架構中,每個塔都是神經網路,可處理查詢或候選項目特徵,然後產生這些特徵的嵌入表示法。每個塔架都會分別部署,因為每個塔架都會用於生產中的不同工作:

  • 候選項目塔:候選項目塔用於預先計算所有候選項目的嵌入。預先計算的嵌入項目會部署至 Vertex AI Vector Search 索引端點,該端點經過最佳化,可實現低延遲檢索。
  • 已部署的塔:在線上放送期間,已部署的查詢塔會將原始使用者查詢轉換為嵌入表示法。然後,系統會使用嵌入表示法,在已部署的索引中尋找類似的項目嵌入。

雙塔架構非常適合用於許多擷取工作,因為雙塔架構會擷取查詢和候選實體的語意關係,並將其對應至共用的嵌入空間。將實體對應至共用嵌入空間後,語意相似的實體會更緊密地聚集在一起。因此,如果您計算特定查詢的向量嵌入,就能在嵌入空間中搜尋最接近 (最相似) 的候選項目。這類架構的主要優點是能夠將查詢和候選項目表示法的推論作業分離。這種解耦方式主要有兩大優點:

  • 您可以放送新 (新鮮) 項目,不必重新訓練新項目詞彙。只要將任何一組項目特徵提供給候選項目塔,您就能計算任何一組候選項目的項目嵌入,即使這些項目在訓練期間未出現也沒問題。執行這項運算有助於解決冷啟動問題。
    • 候選項目塔可支援任意一組候選項目,包括尚未與推薦系統互動的項目。雙塔架構會處理每個<query, candidate>配對的豐富內容和中繼資料特徵,因此可以提供這項支援。這類處理方式可讓系統根據已知項目描述不明項目。
  • 您可以預先計算所有候選項目嵌入,藉此最佳化擷取推論。這些預先計算的嵌入內容可以編入索引,並部署至經過最佳化的服務基礎架構,以利低延遲擷取。
    • 雙塔共同學習可讓您根據查詢描述項目,反之亦然。如果您有一半的配對 (例如查詢),需要尋找其他對應項目,可以預先計算一半的方程式。預先計算可讓您盡快做出其餘決策。

設計須知

本節提供指引,協助您在 Google Cloud 中開發候選人生成架構,滿足安全性和效能需求。本節的指引僅列出部分範例。根據您的具體需求,您可能會選擇考量其他設計因素和取捨。

安全性

Vertex AI Vector Search 支援公開和虛擬私有雲 (VPC) 端點部署作業。如要使用 VPC 網路,請按照「設定 VPC 網路對等互連連線」一文的說明操作。如果 Vector Search 索引部署在 VPC 邊界內,使用者必須從同一個 VPC 網路存取相關聯的資源。舉例來說,如果您是透過 Vertex AI Workbench 進行開發,就必須在與已部署索引端點相同的虛擬私有雲網路中,建立 Workbench 執行個體。同樣地,任何預期會建立端點或將索引部署至端點的管道,都應在同一個虛擬私有雲網路中執行。

效能最佳化

本節說明使用這個參考架構設計 Google Cloud 拓撲時,應考量的因素,以符合工作負載的效能需求。

剖析訓練工作

為提升資料輸入管道和整體訓練圖的效能,建議您使用 Cloud Profiler 分析訓練效能。Profiler 是開放原始碼 TensorBoard Profiler 的代管實作。

在訓練工作中傳遞 –profiler 引數,即可啟用 TensorFlow 回呼,為每個訓練週期設定特定批次數量的剖析檔。設定檔會擷取主機 CPU 和裝置 GPU 或 TPU 硬體的追蹤記錄。追蹤記錄會提供訓練工作的資源耗用量資訊。為避免記憶體不足錯誤,建議您先將設定檔時間長度設為 2 到 10 個訓練步驟,然後視需要增加。

如要瞭解如何搭配使用 Profiler 與 Vertex AI Training 和 Vertex AI TensorBoard,請參閱「分析模型訓練成效」。如需偵錯最佳做法,請參閱「最佳化 GPU 效能」。如要瞭解如何提升效能,請參閱「使用 Profiler 提升 TensorFlow 效能」。

充分運用加速器

附加 NVIDIA GPU 或 Cloud TPU 等訓練加速器時,請務必充分運用這些資源。充分運用訓練加速器是成本管理的最佳做法,因為加速器是架構中最昂貴的元件。充分運用訓練加速器也是提高工作效率的最佳做法,因為沒有閒置時間就能減少整體資源耗用量。

為充分運用加速器,您通常會執行幾次迭代,找出瓶頸、最佳化瓶頸,然後重複這些步驟,直到加速器裝置利用率可接受為止。由於這個用途的許多資料集都太大,無法放入記憶體,因此通常會在儲存空間、主機 VM 和加速器之間發現瓶頸。

下圖顯示機器學習訓練輸入管道的概念階段:

機器學習訓練輸入管道的概念階段。

在圖表中,資料會從儲存空間讀取並預先處理。資料經過前處理後,就會傳送到裝置。如要最佳化效能,請先判斷整體效能是否受限於主機 CPU 或加速器裝置 (GPU 或 TPU)。裝置負責加速訓練迴圈,主機則負責將訓練資料提供給裝置,並接收裝置傳回的結果。以下各節說明如何改善輸入管道效能和裝置效能,解決效能瓶頸。

提升輸入管道效能
  • 從儲存空間讀取資料:如要提升資料讀取效能,請嘗試快取、prefetching循序存取模式和並行 I/O。
  • 預先處理資料:如要改善資料預先處理作業,請為資料擷取和轉換作業設定平行處理,並在資料輸入管道中調整 interleave 轉換
  • 將資料傳送至裝置:如要縮短整體工作時間,請從主機平行傳輸資料至多部裝置。
提升裝置效能
  • 增加迷你批次大小。迷你批次是指每個裝置在訓練迴圈的單次疊代中使用的訓練樣本數。增加微批次大小可提高作業間的平行處理能力,並改善資料重複使用情形。不過,迷你批次必須能夠與其餘訓練程式一起放入記憶體。如果大幅增加迷你批次大小,可能會發生記憶體不足錯誤和模型發散。
  • 向量化使用者定義函式。一般來說,資料轉換可以表示為使用者定義函式,說明如何轉換輸入資料集的每個元素。如要向量化這個函式,請一次對一批輸入內容套用轉換作業,而不是一次轉換一個元素。任何使用者定義函式都有與排程和執行相關的額外負擔。轉換一批輸入內容時,每批內容只會產生一次額外負擔,而不是每個資料集元素一次。
先擴充再擴展

設定訓練工作的運算資源時,建議您先擴充再擴展。也就是說,使用多部效能較低的裝置前,請先選擇效能較高的裝置。建議您按照下列方式擴大規模:

  1. 單一工作人員 + 單一裝置
  2. 單一工作人員 + 更強大的裝置
  3. 單一工作人員 + 多部裝置
  4. 分散式訓練

如要評估 ANN 搜尋的優點,可以測量特定查詢的延遲時間和召回率。為協助調整索引,Vertex AI Vector Search 提供建立暴力搜尋索引的功能。暴力搜尋索引會執行詳盡搜尋,但延遲時間較長,目的是找出指定查詢向量的真正近鄰。暴力搜尋索引不適用於實際運作,但計算索引調整期間的召回率時,可做為良好的基準。

如要評估延遲時間的召回率,請將預先計算的候選項目嵌入內容部署至一個設定為 ANN 搜尋的索引,以及另一個設定為暴力搜尋的索引。暴力搜尋索引會傳回絕對鄰近的鄰居,但通常會比 ANN 搜尋花費更長的時間。您可能願意犧牲一些檢索召回率,以換取檢索延遲的改善,但應評估這種取捨。影響喚回率和延遲時間的其他特徵包括:

  • 模型化參數:許多模型化決策都會影響嵌入空間,而這最終會成為放送索引。比較從淺層和深層擷取模型建構的索引中擷取的候選項目。
  • 尺寸:尺寸是最終由模型決定的另一個層面。ANN 索引的維度必須與查詢和候選塔向量的維度相符。
  • 擁擠和篩選標記:標記可提供強大的功能,針對不同的製作用途調整結果。建議您瞭解標記如何影響擷取的候選人和成效。
  • ANN 計數:增加這個值可提高召回率,但延遲時間也會相應增加。
  • 要搜尋的分葉節點百分比:要搜尋的分葉節點百分比是評估回想率與延遲時間取捨關係時,最重要的選項。提高這個值會增加喚回度,並可能成比例增加延遲時間。

後續步驟

如需更多參考架構、圖表和最佳做法,請瀏覽 Cloud 架構中心

貢獻者

作者:

其他貢獻者:Kaz Sato | 員工開發人員服務代表