Optimum Intel と fastRAG を使用した CPU 最適化埋め込み

Hugging Face と Intel は、Optimum Intel ライブラリと fastRAG フレームワークを使用して、Intel Xeon CPU 上の埋め込みモデルを大幅に高速化する手法を実証しました。int8 への事後学習静的量子化を採用することで、ベースラインの bf16 モデルに比べてレイテンシを最大 4.5 倍、スループットを最大 4 倍向上させます。

RAG 用埋め込みモデルの高速化

埋め込みモデルは Retrieval Augmented Generation (RAG) パイプラインにおいて重要な役割を果たし、主に次の 3 つの機能を担います：オフラインの文書インデックス作成、リアルタイムのクエリエンコーディング、取得した文書の再ランク付け。セマンティック検索は BM25 のような疎検索に比べてコンテキストをよりよく捉えますが、計算コストが高くなります。

CPU バックエンドでこれらのモデルを最適化することは以下の点で重要です:

  • 文書インデックス作成: スループットを向上させ、大規模コレクションのインデックス作成を高速化。
  • クエリエンコーディング: レイテンシを削減し、リアルタイム検索の応答性を向上。
  • 再ランク付け: 時間制約のあるアプリケーション向けに、候補セットの高速処理を実現。

Optimum Intel と IPEX による技術的最適化

Optimum Intel は、Intel ハードウェア上で Hugging Face パイプラインを高速化するために設計されたオープンソースライブラリです。Intel® Advanced Vector Extensions 512 (AVX-512)、Vector Neural Network Instructions (VNNI)、および Intel® Advanced Matrix Extensions (AMX) を活用して、ディープラーニングワークロードを最適化します。

量子化ワークフロー

最適化プロセスは fp32 から int8 への精度低減に焦点を当て、以下の手順で実施します:

  1. インストール: optimum[neural-compressor] と intel-extension-for-transformers を利用。
  2. 事後学習静的量子化: キャリブレーションセット（例: qasper データセットのサブセット）を使用して重みと活性化の動的範囲を決定し、int8 への変換時の精度損失を最小化。
  3. 推論: 量子化モデルは IPEXModel を介してロードされ、Intel Extension for PyTorch (IPEX) による最適化ランタイムで実行されます。

BGE モデルのパフォーマンスベンチマーク

評価は、4 世代目 Intel Xeon 8480+ CPU 上で、BGE（Beijing Academy of Artificial Intelligence）モデルの Small（45M）、Base（110M）、Large（355M）パラメータサイズを使用して実施しました。

精度保持

int8 への量子化は、MTEB（Massive Text Embedding Benchmark）タスク全体で fp32 モデルと比較して精度低下が最小限に抑えられます:

  • 再ランク付け: エラー率 1% 未満。
  • 検索: エラー率 1.55% 未満。

レイテンシとスループット

ベースラインの PyTorch bf16 実装と比較して、int8 量子化モデルは大幅な改善を示します:

  • レイテンシ: 最大 4.5 倍の高速化。Small と Base モデルは 10ms 未満、Large モデルは 20ms 未満のレイテンシを達成。
  • スループット: 最大 4 倍の向上。すべてのモデルサイズで、256 トークンの文書に対しバッチサイズ 128 がピークスループットとなります。

fastRAG との統合

fastRAG は Intel Labs が提供する研究フレームワークで、Haystack と互換性のある最適化 RAG パイプラインを実現します。最適化された埋め込みモデルは、次の 2 つのモジュールを通じて fastRAG に統合されます:

  • QuantizedBiEncoderRetriever: 密なインデックスから文書をインデックス化および取得するために使用。
  • QuantizedBiEncoderRanker: 取得した文書を再ランク付けし、関連性を向上させるために使用。

これらのモジュールにより、開発者は Haystack ベースのパイプライン内で標準の埋め込みモデルを Intel 最適化版に簡単に置き換えることができ、インデックス作成フェーズと検索フェーズの両方の効率を向上させることができます。

Sources