大規模な候補生成のための 2 つのタワーの検索を実装する

Last reviewed 2025-01-16 UTC

このドキュメントでは、Gemini Enterprise Agent Platform を使用してエンドツーエンドの Two-Tower 候補生成ワークフローを実装する方法を示すリファレンス アーキテクチャについて説明します。Two-Tower モデリング フレームワークは、ウェブ検索と候補アイテムなど、2 つの異なるエンティティ間のセマンティック類似性を学習するため、パーソナライズのユースケースに有効な検索手法です。

このドキュメントは、低レイテンシのサービング要件で大規模なレコメンデーション アプリケーションを開発しているデータ サイエンティストや ML エンジニアなどの技術担当者を対象としています。2 つのタワーモデルの構築のためのモデリング手法、問題のフレーミング、データ準備の詳細については、TensorFlow Recommender とベクトル検索を使用したディープ リトリーブのスケーリングをご覧ください。

アーキテクチャ

次の図は、Two-Tower モデルをトレーニングし、さまざまなデプロイとサービングのタスク用に各タワーを個別にデプロイするアーキテクチャを示しています。

Two-Tower モデルをトレーニングし、各タワーを個別にデプロイするアーキテクチャ。

この図のアーキテクチャには、次のコンポーネントが含まれています。

  • トレーニング データ: トレーニング ファイルは Cloud Storage に保存されます。
  • Two-Tower トレーニング: Gemini Enterprise Agent Platform Managed Training サービスを使用して、複合 Two-Tower モデルをオフラインでトレーニングします。各タワーは別々に保存され、異なるタスクに使用されます。
  • 登録されたクエリタワーと候補タワー: タワーのトレーニング後、各タワーは Gemini Enterprise Agent Platform の Model Registry に個別にアップロードされます。
  • デプロイされたクエリタワー: 登録されたクエリタワーが Agent Platform のオンライン エンドポイントにデプロイされます。
  • バッチ予測エンベディング: 登録された候補タワーは、バッチ予測ジョブで使用され、使用可能なすべての候補アイテムのエンベディング表現を事前計算します。
  • エンベディング JSON: 予測されたエンベディングが Cloud Storage の JSON ファイルに保存されます。
  • ANN インデックス: Agent Platform Vector Search を使用して、近似最近傍探索(ANN)用に構成されたサービング インデックスを作成します。
  • デプロイされた検索インデックス: ANN インデックスが Agent Platform Vector Search 検索インデックス エンドポイントにデプロイされます。

使用するプロダクト

このリファレンス アーキテクチャでは、次の Google Cloud プロダクトを使用します。

  • Gemini Enterprise Agent Platform Managed Training: 大規模なモデル トレーニングの運用を可能にするマネージド トレーニング サービス。
  • ベクトル検索: 意味的に類似または関連するデータを保存、インデックス化、検索できるベクトル類似度マッチング サービス。
  • Gemini Enterprise Agent Platform の Model Registry: ML モデルのライフサイクルを管理できる中央リポジトリ。
  • Cloud Storage: 低コストで無制限のオブジェクト ストア。さまざまなデータ型に対応しています。データには Google Cloudの内部および外部からアクセスでき、冗長性を確保するために複数のロケーションに複製されます。

ユースケース

低レイテンシでのサービング要件を満たすために、大規模なレコメンダーは多くの場合、2 段階システムとして本番環境にデプロイされます。場合によっては、多段階システムとしてデプロイされることもあります。最初の段階である候補生成の目標は、候補アイテムの膨大なコレクションをふるいにかけ、後工程のフィルタリングおよびランキング タスクのために関連性の高いアイテムのサブセット(数百個程度)を抽出することです。このリトリーブ タスクを最適化するため、次の 2 つの基本的な目標を検討します。

  1. モデルのトレーニング中に、解決する問題やタスクの最適な表現を学習し、この表現を <query, candidate> エンベディングにコンパイルします。
  2. モデル提供中に、レイテンシ要件を満たす速さで関連性の高いアイテムを抽出する。

次の図は、2 段階のレコメンダーの概念コンポーネントを示しています。

2 段階の Recommender の概念コンポーネント。

図では、候補生成で数百万の候補アイテムがフィルタリングされています。ランキングでは、数百個の候補アイテムをフィルタして、数十個のおすすめアイテムを返します。

このドキュメントのリファレンス アーキテクチャでは、2 つのタワーベースの検索モデルをトレーニングします。このアーキテクチャでは、各タワーはクエリまたは候補アイテムの特徴量を処理し、それらの特徴量のエンベディング表現を生成するニューラル ネットワークです。各タワーは本番環境で異なるタスクに使用されるため、別々にデプロイされます。

  • 候補タワー: 候補タワーは、すべての候補アイテムのエンベディングを事前計算するために使用されます。事前計算されたエンベディングは、低レイテンシの取得に最適化された Vector Search インデックス エンドポイントにデプロイされます。
  • デプロイされたタワー: オンライン サービング中、デプロイされたクエリタワーは、ユーザーの生クエリをエンベディング表現に変換します。エンベディング表現は、デプロイされたインデックスで類似アイテムのエンベディングを検索するために使用されます。

Two-Tower アーキテクチャは、クエリと候補の各エンティティのセマンティック関係を取得して、共通のエンベディング空間にそれらをマッピングできるため、多くのリトリーブ タスクに最適です。エンティティが共有エンベディング空間にマッピングされると、意味的に類似したエンティティが近くにクラスタ化されます。したがって、あるクエリのベクトル エンベディングを計算する場合に、エンベディング空間内で最も近い(最も類似した)候補アイテムを検索できます。このようなアーキテクチャの主な利点は、クエリと候補の表現の推論を分離できることです。この分離の主なメリットは次の 2 つです。

  • 新しいアイテムの語彙を再トレーニングしなくても、新しい(新鮮な)アイテムを提供できます。候補アイテム タワーにアイテム特徴量の任意のセットを渡すことで、トレーニング中に見られなかった候補を含む、候補の任意のセットのアイテム エンベディングを計算できます。この計算を行うことで、コールド スタートの問題に対処できます。
    • 候補タワーは、まだレコメンデーション システムとやり取りしていないアイテムを含め、任意の候補アイテムのセットをサポートできます。このサポートは、2 つのタワー アーキテクチャが各 <query, candidate> ペアに関するリッチ コンテンツとメタデータ機能を処理するために可能です。このような処理により、システムは不明なアイテムを既知のアイテムの観点から説明できます。
  • すべての候補アイテム エンベディングを事前に計算することで、検索推論を最適化できます。これらの事前計算されたエンベディングは、低レイテンシの取得に最適化されたサービング インフラストラクチャにインデックス登録してデプロイできます。
    • タワーの共同学習により、アイテムをクエリの観点から説明したり、その逆も可能になります。クエリなどのペアの半分があり、対応するもう一方のアイテムを探す必要がある場合は、方程式の半分を事前に計算できます。事前計算により、残りの決定をできるだけ迅速に行うことができます。

設計上の考慮事項

このセクションでは、セキュリティとパフォーマンスのニーズを満たす候補生成アーキテクチャを Google Cloud で開発する際に役立つガイダンスを提供します。このセクションのガイダンスはすべてを網羅しているわけではありません。特定の要件によっては、追加の設計要素やトレードオフを考慮することをおすすめします。

セキュリティ

Agent Platform ベクトル検索は、パブリック エンドポイントと Virtual Private Cloud(VPC)エンドポイントの両方のデプロイをサポートしています。VPC ネットワークを使用する場合は、VPC ネットワーク ピアリング接続を設定するの手順に沿って操作してください。ベクトル検索インデックスが VPC 境界内にデプロイされている場合、ユーザーは同じ VPC ネットワーク内から関連リソースにアクセスする必要があります。たとえば、Gemini Enterprise Agent Platform Workbench から開発している場合は、デプロイされたインデックス エンドポイントと同じ VPC ネットワーク内に Workbench インスタンスを作成する必要があります。同様に、エンドポイントの作成やエンドポイントへのインデックスのデプロイが想定されるパイプラインは、同じ VPC ネットワーク内で実行する必要があります。

パフォーマンスの最適化

このセクションでは、このリファレンス アーキテクチャを使用して、ワークロードのパフォーマンス要件を満たす Google Cloud のトポロジを設計する際に考慮すべき要素について説明します。

トレーニング ジョブをプロファイリングする

データ入力パイプラインとトレーニング グラフ全体を最適化するには、Cloud Profiler を使用してトレーニング パフォーマンスをプロファイリングすることをおすすめします。プロファイラは、オープンソースの TensorBoard プロファイラのマネージド実装です。

トレーニング ジョブで –profiler 引数を渡すと、TensorFlow コールバックが各エポックのバッチ数をプロファイリングできるようになります。プロファイルは、ホスト CPU とデバイス GPU または TPU ハードウェアからトレースをキャプチャします。トレースには、トレーニング ジョブのリソース消費に関する情報が提供されます。メモリ不足のエラーを回避するには、プロファイルの期間を 2 ~ 10 個のトレーニング ステップで開始し、必要に応じて増やすことをおすすめします。

Agent Platform Managed Training と Vertex AI TensorBoard で Profiler を使用する方法については、モデルのトレーニング パフォーマンスのプロファイルを作成するをご覧ください。デバッグのベスト プラクティスについては、GPU パフォーマンスの最適化をご覧ください。パフォーマンスを最適化する方法については、プロファイラを使用して TensorFlow のパフォーマンスを最適化するをご覧ください。

アクセラレータを最大限に活用する

NVIDIA GPU や Cloud TPU などのトレーニング アクセラレータを接続する場合は、それらを最大限に活用することが重要です。アクセラレータはアーキテクチャの中で最も高価なコンポーネントであるため、トレーニング アクセラレータを最大限に活用することは、費用管理のベスト プラクティスです。トレーニング アクセラレータを最大限に活用することも、ジョブの効率を高めるためのベスト プラクティスです。アイドル時間がないため、リソースの全体的な消費量が少なくなります。

アクセラレータを最大限に活用するには、通常、ボトルネックの特定、ボトルネックの最適化、アクセラレータ デバイスの使用率が許容範囲になるまでこれらの手順を繰り返します。このユースケースのデータセットの多くはメモリに収まらないほど大きいため、通常、ボトルネックはストレージ、ホスト VM、アクセラレータの間にあります。

次の図は、ML トレーニング入力パイプラインの概念ステージを示しています。

ML トレーニング入力パイプラインの概念的なステージ。

図では、データがストレージから読み取られ、前処理されています。データは前処理された後、デバイスに送信されます。パフォーマンスを最適化するには、まず、全体的なパフォーマンスがホスト CPU によって制限されているか、アクセラレータ デバイス(GPU または TPU)によって制限されているかを判断します。デバイスはトレーニング ループの高速化を担当し、ホストはデバイスへのトレーニング データのフィードとデバイスからの結果の受信を担当します。以降のセクションでは、入力パイプラインのパフォーマンスとデバイスのパフォーマンスを改善してボトルネックを解消する方法について説明します。

入力パイプラインのパフォーマンスを改善する
  • ストレージからのデータの読み取り: データ読み取りを改善するには、キャッシュ保存、prefetchingシーケンシャル アクセス パターン、並列 I/O を試してください。
  • データの前処理: データの前処理を改善するには、データ抽出と変換の並列処理を構成し、データ入力パイプラインで interleave 変換を調整します。
  • デバイスへのデータの送信: ジョブ全体の時間を短縮するため、ホストから複数のデバイスにデータを並行して転送します。
デバイスのパフォーマンスを改善する
  • ミニバッチ サイズを大きくする。ミニバッチは、トレーニング ループの 1 回の反復で各デバイスが使用するトレーニング サンプルの数です。ミニバッチサイズを大きくすると、オペレーション間の並列処理が増加し、データの再利用が改善されます。ただし、ミニバッチは、残りのトレーニング プログラムとともにメモリに収まる必要があります。ミニバッチサイズを大きくしすぎると、メモリ不足エラーやモデルの分散が発生する可能性があります。
  • ユーザー定義関数をベクトル化する。通常、データ変換は、入力データセットの各要素を変換する方法を記述するユーザー定義関数として表現できます。この関数をベクトル化するには、一度に 1 つの要素を変換するのではなく、入力のバッチに対して一度に変換オペレーションを適用します。ユーザー定義関数には、スケジューリングと実行に関連するオーバーヘッドがあります。入力のバッチを変換する場合、オーバーヘッドはデータセット要素ごとに 1 回ではなく、バッチごとに 1 回発生します。
スケールアウトの前にスケールアップする

トレーニング ジョブのコンピューティング リソースを構成する場合は、スケールアウトする前にスケールアップすることをおすすめします。つまり、複数の低電力デバイスを使用する前に、より強力な大型デバイスを選択する必要があります。次の方法でスケールインすることをおすすめします。

  1. 単一ワーカー + 単一デバイス
  2. 単一ワーカー + より強力なデバイス
  3. 単一ワーカー + 複数のデバイス
  4. 分散トレーニング

ANN 検索のメリットを評価するには、特定のクエリのレイテンシと再現率を測定します。インデックスのチューニングに役立つように、Agent Platform ベクトル検索にはブルート フォース インデックスを作成する機能が用意されています。ブルート フォース インデックスは、レイテンシが高くなる代わりに、網羅的な検索を実行して、特定のクエリベクトルの真の最近傍を見つけます。ブルート フォース インデックスの使用は本番環境での使用を想定していませんが、インデックス チューニング中に再現率を計算する際の適切なベースラインとなります。

レイテンシに対する再現率を評価するには、事前計算された候補エンベディングを、ANN 検索用に構成されたインデックスと、ブルート フォース検索用に構成されたインデックスにデプロイします。ブルート フォース インデックスは絶対最近傍を返しますが、通常は ANN 検索よりも時間がかかります。検索レイテンシを短縮するために検索再現率を多少犠牲にしてもよいかもしれませんが、このトレードオフは評価する必要があります。再現率とレイテンシに影響するその他の特性は次のとおりです。

  • モデリング パラメータ: モデリングに関する多くの決定事項は、エンベディング空間に影響を与え、最終的にサービング インデックスになります。浅い検索モデルとディープ リトリーブ モデルの両方から構築されたインデックスに対して取得された候補を比較します。
  • ディメンション: ディメンションは、最終的にモデルによって決定される別の側面です。ANN インデックスのディメンションは、クエリと候補タワー ベクトルのディメンションと一致している必要があります。
  • クラウディング タグとフィルタリング タグ: タグは、さまざまな本番環境のユースケースに合わせて結果を調整する強力な機能を提供します。タグが取得された候補にどのように影響し、パフォーマンスにどのような影響を与えるかを理解することをおすすめします。
  • ANN カウント: この値を大きくすると再現率が向上し、それに比例してレイテンシも増加する可能性があります。
  • 検索するリーフノードの割合: 検索するリーフノードの割合は、レイテンシと再現率のトレードオフを評価するうえで最も重要なオプションです。この値を大きくすると再現率が向上し、それに比例してレイテンシも増加する可能性があります。

次のステップ

Cloud アーキテクチャ センターで、リファレンス アーキテクチャ、図、ベスト プラクティスを確認する。

寄稿者

著者:

その他の寄稿者: Kaz Sato | スタッフ デベロッパー アドボケイト