商用DRAM-PIMに基づく近似最近傍探索エンジン、DRIM-ANN

Concepts de base

DRIM-ANNは、従来のCPUやGPUの限界を超えて、大規模データセットに対する高速かつエネルギー効率の高い近似最近傍探索（ANNS）を実現する、商用DRAM-PIMアーキテクチャを活用した革新的なANNSエンジンである。

Résumé

DRIM-ANNの概要

本稿では、商用DRAM-PIMを用いた初の近似最近傍探索（ANNS）フレームワークであるDRIM-ANNについて述べている。ANNSは、情報検索や検索拡張生成（RAG）などの重要なアプリケーションにおいて、大規模データセットにおける効率的な意味的類似性検索を可能にする、基礎的な構成要素となっている。しかし、ANNSは、高次元データの大容量化により、 massive storage を必要とすることが多く、計算量に対してI/O量の多いアルゴリズムであることが知られている。このことが、CPUではI/Oのボトルネックに、GPUではメモリ容量の限界につながっている。高帯域幅、大容量メモリ、そしてデータの近傍またはデータ内で効率的な計算を実行する能力を提供する、DRAMベースのProcessing-in-Memory（DRAM-PIM）アーキテクチャは、ANNSにとって有望なソリューションとなる。

本稿では、UPMEMのDRAM-PIMを用いたANNS向けに最適化されたANNSエンジンであるDRIM-ANNを提案する。まず、クラスタベースのANNSの計算パターンを特定し、IVF-PQとその派生形であるOPQ [16]やDPQ [25]をサポートするANNSエンジンを設計する。次に、UPMEMの高いメモリ帯域幅と比較的低い計算能力を活用するために、一般的なANNSエンジンで必要とされる計算量の多いユークリッド距離計算を、損失のないルックアップテーブル（LUT）に変換することで、計算能力の制限を補う。この変換されたANNSエンジンに基づいて、ターゲットとなるDRAM-PIMアーキテクチャ向けにANNS構成を最適化する、PIM対応のアルゴリズムチューニングフレームワークを提案する。具体的には、アーキテクチャを意識したチューニングをガイドするために、効果的なANNSパフォーマンスモデルを開発する。さらに、2,000個を超える並列DPUを搭載したUPMEMにANNSを展開する際に生じる、重大な負荷分散の課題にも取り組む。この課題に対処するために、静的なデータレイアウトの最適化と動的なランタイムリクエストスケジューリングを組み合わせた負荷分散戦略を開発し、それによって負荷分散の問題を軽減する。代表的なデータセットを用いた実験結果から、DRIM-ANNは、32スレッドCPUと比較して平均2.92倍の性能向上を達成することが示された。

DRIM-ANNの構成要素

DRIM-ANNは、以下の3つの主要な最適化技術から構成されている。

1. 乗算器レスANNS変換

DRAM-PIMは、計算能力は低いものの、非常に優れた帯域幅を提供する。そこで、乗算をなくすために帯域幅をある程度犠牲にする。具体的には、計算量の多い乗算を、損失のないLUTに置き換える。基本的な考え方は、L2距離計算の乗算のオペランドは常に同じであるため、すべての可能な2乗結果の量は、オペランドのすべての可能な値の量と同じであり、8ビットまたは16ビットのオペランドの場合はそれぞれ128または64Kとなる。そこで、オフラインですべての可能な値を計算し、LUTに保存する。このLUTは、オンライン検索のためにすべてのDPUにブロードキャストされる。ビット幅の大きいオペランドの場合、LUT全体が許容できない場合は、オフラインでは小さい値の2乗結果のみを格納するLUTを構築する。乗算からLUT検索への変換により、計算量の多い演算がメモリアクセスに変換されるため、高帯域幅のPIMプラットフォームで高速化することができる。

2. ANNパフォーマンスモデリング

乗算器レス変換を適用することで、DRAM-PIMから計算量の多い乗算が排除される。クラスタベースのANNSの性能と精度を両立させるパラメータはいくつかある。あるパラメータが性能と精度に与える影響は、他のパラメータによって補償することができる。例えば、候補となる近傍リストを縮小すると、TSフェーズの性能は向上するが、精度は低下する。これは、コードブックのエントリ数やアクセスするクラスタ数を増やすことで補償することができる。

3. ANN設計空間探索

特定のデータセットとPIMプラットフォームでは、N、Q、D、BxとBWx、PE、Fxが決まっている。したがって、設計空間探索（DSE）のタスクは、精度制約の下で最適な性能を達成する（K、P、C、M、CB）の組み合わせを見つけることである。

DRIM-ANNの評価

SIFT100MとDEEP100Mの代表的なデータセットを用いた実験結果から、DRIM-ANNは、32スレッドCPUと比較して平均2.92倍の性能向上を達成することが示された。また、DRIM-ANNはCPUベースラインと比較して1.63倍から2.42倍のエネルギー効率を達成し、その平均値は1.97倍である。

結論

DRIM-ANNは、商用DRAM-PIMアーキテクチャを活用した、効率的でスケーラブルなANNSソリューションを提供する。乗算器レスANNS変換、PIM対応アルゴリズムチューニング、効果的な負荷分散戦略により、DRIM-ANNはCPUベースラインと比較して大幅な性能向上とエネルギー効率の向上を実現する。

Stats

DRIM-ANNは、32スレッドCPUと比較して平均2.92倍の性能向上を達成。
DRIM-ANNはCPUベースラインと比較して1.63倍から2.42倍のエネルギー効率を達成し、その平均値は1.97倍。
乗算器レスANNS変換により、エンドツーエンドの性能で平均1.40倍、LCで平均1.93倍の高速化を実現。
提案する性能モデルと実際のパフォーマンスの差は3.32倍から6.48倍で、平均値は5.23倍。

Citations

Idées clés tirées de

DRIM-ANN: An Approximate Nearest Neighbor Search Engine based on Commercial DRAM-PIMs

by Mingkai Chen... à arxiv.org 10-22-2024

https://arxiv.org/pdf/2410.15621.pdf

Questions plus approfondies

DRIM-ANNは、他のPIMアーキテクチャ（例えば、SamsungのHBM-PIM）にどのように適応できるだろうか？

DRIM-ANNは、UPMEMのDRAM-PIMをベースに設計されていますが、そのコアとなる設計思想は、SamsungのHBM-PIMのような他のPIMアーキテクチャにも適用できます。
HBM-PIMへの適用における重要なポイント:

高帯域幅の活用: DRIM-ANNは、LUT（ルックアップテーブル）を用いることで乗算演算をメモリ参照に置き換え、PIMの高いメモリ帯域幅を活用しています。HBM-PIMも高帯域幅メモリであるため、この設計は有効です。
計算能力の考慮: HBM-PIMは、UPMEMのDRAM-PIMと比較して計算能力が高い傾向がありますが、依然としてCPUやGPUと比較すると制限があります。DRIM-ANNのアルゴリズム調整や負荷分散の戦略は、HBM-PIMの計算能力を最大限に引き出すために重要です。
アーキテクチャ固有の最適化: HBM-PIMのアーキテクチャに合わせて、データ配置やデータ転送の最適化が必要になります。例えば、HBM-PIMのメモリ階層や処理ユニットの構成を考慮したデータ配置や、処理ユニット間のデータ転送を最小限に抑えるアルゴリズム設計が求められます。
具体的な適応例:

データ配置の最適化: HBM-PIMのメモリ階層構造に合わせて、頻繁にアクセスされるデータをより高速なメモリに配置するなどの最適化を行います。
処理ユニットへのタスク割り当て: HBM-PIMの処理ユニットの特性を考慮し、計算負荷を均等に分散するようにタスクを割り当てます。
データ転送の最小化: HBM-PIMのメモリ帯域幅は高いため、データ転送のオーバーヘッドを最小限に抑えることが重要です。DRIM-ANNのLUTを用いた設計は、この点で有利に働きます。
まとめ:
DRIM-ANNは、高帯域幅メモリと計算ユニットを組み合わせたPIMアーキテクチャに適した設計思想を持っています。HBM-PIMのような他のPIMアーキテクチャにも、アーキテクチャ固有の最適化を加えることで適用できる可能性があります。

DRIM-ANNの性能は、データセットの次元やデータ分布によってどのように変化するだろうか？

DRIM-ANNの性能は、データセットの次元やデータ分布に影響を受けます。
データセットの次元:

高次元データへの影響: DRIM-ANNは、LUTを用いて距離計算の負荷を軽減していますが、データセットの次元数が大きくなると、LUTのサイズも増大し、メモリ容量の制限を受ける可能性があります。
計算量の増加: 次元数の増加は、距離計算やベクトル演算の計算量増加に直結します。DRIM-ANNの並列処理能力を活用し、計算負荷を分散する必要があります。
データ分布:

偏ったデータ分布の影響: データ分布が偏っている場合、特定のクラスタにクエリが集中し、負荷分散がうまく機能しない可能性があります。DRIM-ANNの動的なクエリスケジューリングやクラスタ複製などの負荷分散技術が重要となります。
均一なデータ分布の影響: データ分布が均一な場合、DRIM-ANNの負荷分散技術は効果的に機能し、高いパフォーマンスを発揮する可能性があります。
性能への影響を最小限に抑えるための対策:

次元削減: 高次元データに対しては、次元削減技術を適用することで、DRIM-ANNのメモリ容量の制限を緩和し、計算量を削減できます。
データ分布の分析: データ分布を事前に分析し、偏りがある場合には、クラスタリング手法の調整やデータ拡張などの対策を検討します。
動的な負荷分散: DRIM-ANNの動的なクエリスケジューリングやクラスタ複製などの負荷分散技術を、データセットの特性に合わせて最適化します。
まとめ:
DRIM-ANNの性能は、データセットの次元やデータ分布に影響を受けます。高次元データや偏ったデータ分布に対しては、適切な対策を講じることで、性能への影響を最小限に抑えることが重要です。

将来のPIMアーキテクチャにおける計算能力とメモリ帯域幅の向上は、DRIM-ANNの性能にどのような影響を与えるだろうか？

将来のPIMアーキテクチャにおける計算能力とメモリ帯域幅の向上は、DRIM-ANNの性能を飛躍的に向上させる可能性があります。
計算能力の向上による影響:

ボトルネックの解消: 現在のDRIM-ANNは、PIMの計算能力がボトルネックとなる場合がありますが、計算能力が向上することで、より高速な検索処理が可能になります。特に、高次元データや複雑な距離計算を用いる場合に効果を発揮します。
より高度なアルゴリズムの適用: 計算能力の向上は、より高度なアルゴリズムの適用を可能にします。例えば、より精度の高い近似最近傍探索アルゴリズムや、動的なインデックス構築などが可能になります。
メモリ帯域幅の向上による影響:

データ転送の高速化: メモリ帯域幅の向上は、データ転送の高速化に直結します。DRIM-ANNは、LUTを用いることでメモリ参照を多用するため、メモリ帯域幅の向上は、検索処理全体のパフォーマンス向上に大きく貢献します。
大規模データセットへの対応: メモリ帯域幅の向上は、より大規模なデータセットの処理を可能にします。これにより、DRIM-ANNは、大規模な画像データベースや自然言語処理など、より広範なアプリケーションに適用できるようになります。
更なる性能向上のための課題:

計算能力とメモリ帯域幅のバランス: 計算能力とメモリ帯域幅のバランスを最適化することが重要です。メモリ帯域幅に対して計算能力が不足すると、データ転送がボトルネックとなり、性能が低下する可能性があります。
ソフトウェアスタックの最適化: PIMアーキテクチャの性能を最大限に引き出すためには、ソフトウェアスタックの最適化が不可欠です。コンパイラ、ランタイムライブラリ、アプリケーションレベルでの最適化が必要です。
まとめ:
将来のPIMアーキテクチャにおける計算能力とメモリ帯域幅の向上は、DRIM-ANNの性能を飛躍的に向上させる可能性があります。これらの進化を活用することで、DRIM-ANNは、様々な分野で大規模データの高速な近似最近傍探索を実現する重要な技術となるでしょう。