はじめに

機械学習エンジニアの本田志温です。最近担当した類似アイテム推薦の案件で、BigQueryを使って最大内積検索（MIPS; maximum inner-product search）¹ を実装したので、その方法と高速化のテクニックを紹介します。

類似アイテム推薦は「多数のアイテム候補から、クエリとなるアイテムに最も類似したK件を抽出する」というタスクなので、MIPSないし近傍探索の枠組みで解くことが一般的です。

一定の規模を持つサービスでMIPSを実装しようとすると、アイテム数×特徴量次元の行列が何かと厄介です²。第一に、MIPSを素朴な行列積で実装すると、時間・空間計算量がアイテム数の2乗でかかってきます。典型的には空間計算量の方がボトルネックになりやすく、RAMの制約に収めるための工夫が必要になるでしょう。第二に、アイテム数が膨大な場合、特徴量マートから全アイテムの特徴量を転送してくるのにかかる時間も無視できません。

もし特徴量マートを BigQuery 上に作成しているのであれば、MIPSをBigQuery上で実行できます。BigQueryならメモリ制約のことをあまり気にせずにクエリを実行できて、データ転送も不要です。私が担当した案件でも、このメリットを享受するためにMIPSをBigQuery上で実行することにしました。

ただし、この方法が有効なのはバッチ推論の場合に限ります。オンライン推論にはBigQueryは向かないので、近似近傍探索など別の計算量削減方法を検討すべきです。Google Cloudであれば、 Vertex AI Matching Engine というMIPSのマネージドサービスがあります。

以降では、類似アイテム推薦のバッチ推論という事例を想定して、BigQueryによるMIPSの実装方法を紹介します。

ダミーテーブルの準備

まずは実験用のデータを用意します。アイテム数を1万とし、それぞれにL2正規化済みの特徴量（4次元）を乱数で付与します。また、後で使うためにアイテムのカテゴリ（全部で10種類）をランダムに付与します。実際にはアイテム数・特徴量次元・カテゴリ数のいずれももっと大きくなると思いますが、手軽に実行できるようにあえて小さくしています。

create temp table items as
  select 
    generate_uuid() as item_id
    , mod(cast(rand() * 10 as int64), 10) as category
  from unnest(generate_array(1, 10000))
;
create temp table rands as
  select
    item_id
    , category
    , (row_number() over (partition by item_id)) as vector_index
    , elem / sqrt(sum(elem * elem) over (partition by item_id)) as elem
  from (
    select
      item_id
      , category
      ,rand() * 2 - 1 as elem
    from items
    cross join unnest(generate_array(1, 4))
  )
;

create table <TABLE_ID> as
select
  item_id
  , any_value(category) as category
  , array_agg(elem order by vector_index) as vector
from rands
group by item_id

上記クエリを実行すると、このようなテーブルができるはずです。

ベクトルの内積を計算する方法

ベクトル同士の内積は以下のUDF (user-defined function)で計算できます。

create temp function dot_product(vec1 array<float64>, vec2 array<float64>)
returns float64
as (
 (select sum(elem1 * elem2)
  from unnest(vec1) elem1 with offset pos1 
  join unnest(vec2) elem2 with offset pos2 
  on pos1 = pos2)
);

MIPS（全探索）

上記のUDFを使うと、MIPSは簡単に書けます。各アイテムについて（自分以外の）全アイテムとの内積を計算し、その値の降順でK（=10）件を抽出するだけです。

create temp table item_pairs_with_similarity as
select 
  q.item_id as item_query
  , r.item_id as item_result
  , dot_product(q.vector, r.vector) as similarity
from 
  <TABLE_ID> q
cross join
  <TABLE_ID> r
where
  q.item_id != r.item_id
;

select 
  item_query
  , item_result
  , similarity
  , rank
from (
  select
    item_query
    , item_result
    , similarity
    , rank() over (partition by item_query order by similarity desc) as rank
  from
    item_pairs_with_similarity
)
where 
  rank <= 10
order by
  item_query
  , rank

以上でやりたかったことが実現できました。「処理されたバイト数」は7.82 GBでした。

メタデータによる候補削減

全探索よりも効率的な方法は考えられないでしょうか？

実際の類似アイテム推薦の現場では、アイテムのメタデータを使って候補を絞り込むことができるケースがほとんどです。例えば、何かの商品であれば商品カテゴリ、ホテルであればエリア、ニュース記事であれば投稿日などがメタデータとして利用できるでしょう。この方針で候補を絞り込めば、推薦精度をほとんど犠牲にすることなく計算量を削減できます。

そこで、ダミーデータの作成時に用意したカテゴリのカラムを使って事前に候補を削減する方法を試してみます。BigQueryでは処理したデータ量によって課金額が決まるので、候補削減はコストの節約にも繋がります。

実はこの候補削減はとても簡単で、先程のwhere句に1行足すだけで実現できます。

create temp table item_pairs_with_similarity as
select 
  q.item_id as item_query
  , r.item_id as item_result
  , dot_product(q.vector, r.vector) as similarity
from 
  <TABLE_ID> q
cross join 
  <TABLE_ID> r
where 
  q.item_id != r.item_id
  # 同一カテゴリに絞り込む
  and q.category = r.category
;

この工夫によって「処理されたバイト数」は802.05 MBになりました。カテゴリ数が10だったので、ちょうど10分の1程度になっています。

ちなみに、ダミーデータなので抽出アイテムの類似度が全探索のときよりも小さくなっていますが、実際には類似度が高いアイテムは同一カテゴリ内に集中しているはずなので、大きな変化は起きないと考えられます。

ベクトルの離散化による候補削減

アイテムのメタデータを使わずに候補削減をすることもできます。特徴ベクトルの各要素を離散化し、その値が一致する数を利用して絞り込むというアプローチです。

以下のクエリでは、特徴ベクトルの各要素を-1または0の2通りに離散化し、候補を4次元の要素すべてが一致するものに絞り込んでいます。

create temp function array_equal_count(vec1 array<int64>, vec2 array<int64>)
returns int64
as (
 (select count(1)
  from unnest(vec1) elem1 with offset pos1 
  join unnest(vec2) elem2 with offset pos2 
  on pos1 = pos2 and elem1 = elem2)
);

create temp table item_quantized_vector as (
select
  item_id
  , category
  , vector
  , array(select cast(floor(elem) as integer) from unnest(vector) elem) as quantized_vector
from <TABLE_ID>
);

create temp table item_pairs_with_similarity as (
select 
  q.item_id as item_query
  , r.item_id as item_result
  , dot_product(q.vector, r.vector) as similarity
from 
  item_quantized_vector q
cross join
  item_quantized_vector r
where
  q.item_id != r.item_id
  # 離散ベクトルの全要素が一致する候補に絞り込む
  and array_equal_count(q.quantized_vector, r.quantized_vector) = 4
);

この工夫によって「処理されたバイト数」は505.24 MBになりました。2の4乗は16なので、ちょうど16分の1程度になっています。

ランダムサンプリングによる候補削減

メタデータによる候補削減はカテゴリが一様に分布している場合には極めて有効ですが、アイテムが一部のカテゴリに集中している場合だと削減幅が限定的になってしまいます。実際の現場では、例えばホテルが大都市や主要観光地に集中していたりと、カテゴリに偏りがあることの方が多いでしょう。ベクトルの離散化による候補削減も同様の問題をはらみます。

こういったケースでも計算量を減らしたいという場合は、推薦精度を多少犠牲にしてランダムサンプリングで候補削減することも可能です。

「メタデータによる候補削減」で各アイテムの候補数は1000程度でした。これをさらに減らして100程度にしてみましょう。

# カテゴリごとのアイテム数をカウント
create temp table features as
select
  *
  , count(distinct item_id) over (partition by category) as category_size
from
  <TABLE_ID>
;

create temp table item_pairs_with_similarity as
select 
  q.item_id as item_query
  , r.item_id as item_result
  , dot_product(q.vector, r.vector) as similarity
from 
  features q
cross join 
  features r
where 
  q.item_id != r.item_id
  and q.category = r.category
  # どのカテゴリでも候補数が100程度になるようにする
  and rand() < 100 / q.category_size
;

この工夫によって「処理されたバイト数」は81.56 MBになりました。ランダムサンプリングで候補削減したので、「メタデータによる候補削減」のときよりも抽出アイテムの類似度が小さくなっていることに注意してください。

一緒に働きませんか？

弊社では、様々な職種のエンジニアを募集しています。興味のある方は、以下の採用ページをご覧ください。

近傍探索の中でも、類似度として内積を採用したもののこと。 ↩︎
本記事では扱いませんが、巨大なデータを扱う分散処理では、DIMSUMと呼ばれる近似方法を使うことで全てのベクトル間の距離を効率的に計算できます。例えば、 SparkではDIMSUMを使ってコサイン類似度の行列を計算する機能が提供されています。 ↩︎