ヤフー株式会社は、2023年10月1日にLINEヤフー株式会社になりました。LINEヤフー株式会社の新しいブログはこちらです。LINEヤフー Tech Blog

テクノロジー

NGT - 検索ロジックの詳細と実装方法 -

NGT

Yahoo! JAPAN研究所の岩崎です。前回は探索起点の選択方法について説明しました。今回は初回の記事(NGT -グラフ型インデックスによる近似近傍検索の基本-)で説明した検索ロジックを詳細に説明するだけでなく、実装時の注意点を説明しますので、かなり突っ込んだ内容となっています。

グラフインデックスの検索ロジック

では、早速、検索ロジックを簡略化したコードを以下に示します。探索起点ノードは既に取得済みである前提としていますので、入力データの一つとなっています。コード中のwhileループではクエリオブジェクトから最も近い探索候補ノードを一つずつ取り出しては処理します。forループでは、取り出した探索候補ノードにエッジで接続されたすべての近傍ノードについて処理します。順次、各近傍ノードからクエリまでの距離を算出し、検索結果なのか、探索候補ノードなのかを判定します。さらに、検索結果内でクエリオブジェクトから最も遠いオブジェクトの距離を検索半径とします。探索候補ノードがなくなったり、最も近い探索候補ノードすら検索範囲に入らなくなったら、探索を終了します。ステップ単位の詳細については、初回の記事も見ながら、コード中のコメントを読んでいただければ理解できると思います。

// 入力
//   query : クエリオブジェクト
//   k : 検索数
//   seeds : 探索起点ノード
//   nodeNeighbors : グラフ(各ノードの近傍ノードリスト)
//   objects : 検索対象オブジェクトリスト
// 出力
//   result : 検索結果
priority_queue<Node,...> result;                                     // 検索結果(1)
priority_queue<Node> exploredCandidates = seeds;                     // 探索候補ノード(2)に探索起点をセット
unordered_set<NodeID,...> processedNodes;                            // 処理済みノード(3)
float radius = std::numeric_limits<float>::max();                    // 検索範囲に最大値(無限)をセット
while (!exploredCandidates.empty()) {                                // 探索候補ノードがなくなるまでループ
    auto targetNode = exploredCandidates.top();                      // クエリに最も近い探索候補ノードを取得
    exploredCandidates.pop();                                        // 取得したノードを削除
    if (targetNode.distance > radius) break;                         // 取得ノードが検索範囲外であれば終了
    neighbors = nodeNeighbors.at(targetNode.id);                     // エッジで接続された近傍ノードの取得
    for (auto neighbor = neighbors.begin(); neighbor != neighbors.end(); ++neighbor) {
        // 次のノード(neighbor+1)のオブジェクトをプリフェッチ(*)
        if (processedNodes.at((*neighbor).id)) continue;             // ノードが処理済みの場合はスキップ
        processedNodes.insert((*neighbor).id);                       // ノードを処理済みとする
        double distance = compare(query, objects.at((*neighbor).id));// クエリからの距離計算
        if (distance < radius) {                                     // 検索半径内か?
            Node node((*neighbor).id, distance);
            exploredCandidates.push(node);                           // ノードを探索候補に追加
            result.push(node);                                       // ノードを検索結果に追加
            if (result.size() > k) result.pop();                     // 検索結果の検索数を超えたら削減
            if (result.size() >= k) radius = result.top().distance;  // 検索半径の更新
        }
    }
    return result;
}

なお、検索精度を上げるためにグラフインデックスによる検索では、通常、検索範囲を拡大しますが、コードの理解のし易さを優先して検索範囲拡大のコードは省略しています。検索範囲の拡大の方法には以下の2種類があります。

  • 検索結果数(k)を一定割合増やすことで、検索範囲を拡大する。
  • 検索範囲(radius)を一定割合拡大する。

前者は実装が簡単ですが、NGTは後者を実装しています。検索範囲とは別に探索範囲を設定し、探索範囲を検索範囲よりも一定割合拡大して、検索範囲内のオブジェクトは検索結果に追加し、探索範囲内のオブジェクトは探索候補ノードとする処理を行います。

実装上の注意点

以上が検索ロジックの説明でしたが、ここからはこのロジックの実装の話となります。検索ロジックは上記のように複雑なものではないのですが、安易に実装すると検索時間が増大する傾向があるので、実装時の注意点をいくつか紹介します。

距離計算の高速化

高次元のベクトルデータの場合には検索処理全体の時間に占める距離計算時間の割合が多くなります。したがって距離計算の時間を削減することが検索時間の削減に効果的なので、最初に手をつけるべき部分が距離計算の高速化となります。
まず、距離計算はSIMDで計算することで高速化が期待できます。ただし、最近のコンパイラは賢いので、それほど高速にならない場合もあるので期待は禁物です。

さらに、通常、個々の次元ごとに計算を行うので、次元数分のループを回しがちですが、数次元をまとめて1回のループで計算することでも高速化ができます。つまり、命令実行時のパイプラインを乱す分岐数を減らせるからです。ただし、数次元を1回のループで計算すると、ループの最後に残った端数の次元の処理が発生してしまいます。そこで、オブジェクトの対して端数の次元が出ないように(1回のループで計算する次元数の倍数になるように)最後に0の値の次元を加えます。こうすることで、端数の処理がなくなり検索時間を削減できます。

データ構造

上記コードの中で検索結果(1)、探索候補ノード(2)、処理済みノード(3)のデータ構造は処理時間に影響するので注意が必要です。

  • 検索結果 result(1)
    随時最新の検索結果を保持します。検索結果内の最遠のノードの距離により検索半径が更新されるので、検索結果から常に最大値が得られなければなりません。したがって、優先度付きキュー(この例ではC++のpriority_queue)を用いて実装するのが一般的です。

  • 探索候補ノード exploredCandidates(2)
    クエリまでの距離から検索範囲内に存在すると判定されたノード、つまり、探索候補のノードを保持します。探索候補ノードが随時追加されて、かつ、クエリに最も近い探索候補ノードを随時取得できなければなりません。したがって、こちらも優先度付きキューを利用するのが一般的です。検索結果のデータ構造にも言えますが、より高速なデータ構造を使うことによって高速化される可能性がありますし、高速化のためにこのデータ構造を独自に実装している手法も見受けられます。

  • 処理済みノード processedNodes(3)
    探索時に、同じノードを繰返し処理しないように処理済みのノードを保持します。3つのデータ構造の中でも最も頻繁にアクセスされるデータ構造なので特に注意が必要です。この例では連想コンテナであるC++のunordered_setを利用していますが、データ数が少なければ初期化に時間がかからないのでbit vectorを使ってフラグとして実装(全オブジェクトに対するフラグを用意)するのが最も高速になります。NGTではunordered_setとvectorを組み合わせた独自のデータ構造を使ってオブジェクト数に依存せずに常に高速な処理を可能としています。

オブジェクトのプリフェッチ

高速な検索を実現するためにオブジェクトデータ(ベクトルデータ)はメモリ上に展開しておきます。そして、距離計算時にメモリ上のオブジェクトにアクセスします。上記コードではcompare()の呼び出しでオブジェクトにアクセスすることになります。通常のメモリアクセスではコンパイラやCPUのメモリキャッシュの最適化機能により、効果的にメモリキャッシュが使われて、プログラミング上で意識する必要性はあまりないと思います。しかし、グラフ探索時にはメモリ空間上で分散したノードをアクセスすることになり、メモリキャッシュの効果が得られにくいです。分散したノードをアクセスしているといってもランダムにアクセスしているのではなく対象ノードに接続された近傍ノードを順次アクセスするので、プログラミング上はアクセス順が分かっています。したがって、後続の処理でアクセスすると想定されるノード(オブジェクト)をprefetchコマンドで使って明示的にメモリキャッシュに載せておくことで、オブジェクトのアクセス時間を短縮できます。上記コードでは煩雑になるのでプリフェッチのコードは記述していませんが、(*)の部分にプリフェッチのコードを記述します。linux系だと__mm_prefetch関数で容易にメモリのプリフェッチが可能です。

以上の点を注意すれば、実装による速度低下は抑えられると思います。ロジックは簡単なので自分で実装して、高速化に挑戦してみてはいかがでしょうか?

次回は最適化グラフの生成方法について解説します。

Yahoo! JAPAN研究所 岩崎雅二郎

こちらの記事のご感想を聞かせください。

  • 学びがある
  • わかりやすい
  • 新しい視点

ご感想ありがとうございました

このページの先頭へ