こんにちは。レトリバの飯田(@meshidenn)です。カスタマーサクセス部研究チームのリーダーをしており、マネジメントや論文調査、受託のPOCを行なっています。

従来の検索アルゴリズムの問題点

従来の検索アルゴリズムの問題点といえば、"意味"を考慮できないということが挙げられます。従来の検索アルゴリズムは、単語一致をベースとして、そのスコアリングをするのが基本だからです。そのため、単語が一致しないことによる弊害がおきます。そして、「あー、意味を考慮できたらなー」という発想に至ります。

その結果、クエリも文書もベクトル表現にして計算してしまえ！ということで近年研究が盛んに行われており、BERT¹が提案されて以降、教師データがあれば、うまく行くことがわかってきています。さらに、近年、最近傍アルゴリズムが進歩し、内積計算が高速化したこともあり、全件検索でも使えることが知られています。そんな状況の中、「実は、転置インデックスのまま、文脈だけ考慮できるようにすれば、同じくらいの性能が出ます！」という論文が出ました。それが今回ご紹介する COIL: Revisit Exact Lexical Match in Information Retrieval with Contextualized Inverted Listです。この論文の何が、衝撃的かといえば、単語が一致しないことが問題だと思いきや、一致した単語の中で、文脈を考慮して順位づけできていれば、それだけで良い性能になるということを示したことです。

今回は、こちらの論文の概要と簡単な実験のご紹介をします。なお、図は特に断りのないかぎり、当該論文の引用です。

COILの概要

検索時の挙動

まず、学習済みモデルの利用シーンから見ていきたいと思います。通常の転置インデックスは下図の通り、各単語に対して、ある文書で何回その単語が出現しているか(Term Frequency(TF))の情報を格納しています。

f:id:meshitahiro:20210712183117p:plain — 通常の転置インデックス

一方、COILでは、TFの代わりにBERTから出力される、ある文書のある位置のベクトルを格納します。下の図はその様子を表しています。図をよく見ると、Riverにはdocid５が2回あり、2つのriverをそれぞれ分けてインデックスしていることがわかります。検索時は、クエリに含まれる単語について転置インデックスを参照し、各単語毎に格納したベクトルを取得します。また、クエリの各単語をBERTでベクトル化します。そして、転置インデックスから取得したベクトルとBERTを使って変換したベクトルの内積計算を行い、スコアを算出します。この時、スコアは、各文書毎に最大の内積となったペアを採用します。

なお、別途、CLSで文書全体をベクトル化しておきます。そして、クエリのCLS-tokenをBERTでベクトルにし、こちらの内積計算も行います。

f:id:meshitahiro:20210712183231p:plain — 文脈化された転置インデックス

式にすると以下の通りです。

$\begin{gather} s _ {full}(q, d) = s _ {tok} + {\boldsymbol{v} _ {CLS} ^ q} ^ T \boldsymbol{v} _ {CLS} ^ d \\ s _ {tok}(q, d) = \sum _ {q _ i \in q \cap d} \max _ {d _ j = q _ i}({\boldsymbol{v} _ i ^ q} ^ T , \boldsymbol{v} _ j ^ d ) \end{gather}$

ここで、 $q$ と $d$ はそれぞれ、クエリと文書を表し、 $q_i$ と $d_j$ はクエリと文書の単語を表します。 $q \cap d$ はそれぞれに共通に現れる単語群です。 ${ \boldsymbol{v}} _ i ^ q$ と ${\boldsymbol{v} _ j ^ d}$ は、 $q_i$ と $d_j$ をBERTでベクトル化したものです。 $s _ {tok}(q, d)$ が文脈化された転置インデックスによるスコアです。 $s _ {full}(q, d)$ がクエリと文書の最終スコアです。

学習時の挙動

さて、上記のような検索をさせるため、トレーニング時も同じようなスコアリングをした方が良いことが予想されます。そこで、先ほどの式で算出されるスコアを最終スコアとして、ペアワイズランキングロスで学習を行います。

$\begin{aligned} \mathcal{L} = -log \frac{\exp(s(q, d^+))}{\exp(s(q, d^+)) + \sum_l \exp(s(q, d_l^-))} \end{aligned}$

なお、 $d^+$ は、あるクエリ $q$ の正解となる文書、 $d^-$ は不正解となる文書で、 $l$ 個あるとしています。

下図は、先行研究とのアーキテクチャの比較です。図の(d)が提案手法COILの図です。クエリと文書で同じ単語のベクトルでのみ内積をとり、その最大のものを採用しています。また、同時に、全体を代表するものとして、CLS-tokenのベクトルの内積を取っています。これは、図中(b)と同様です。

f:id:meshitahiro:20210712183645p:plain — 先行研究およびCOILモデル

結果

以下が結果です。COILが提案手法です。COIL-tokはCLSを使用しなかった場合です。上図(a), (b), (c) はそれぞれ、表の BM25+BERT reranker, Dense Retriever, ColBERT²に対応しています。まず、MS MARCO Passage³という比較的短いデータを対象とする検索において、先行研究と同等の精度が出ていることがわかります。

f:id:meshitahiro:20210712183814p:plain — MS MARCO Passageでの結果

次に、MS MARCO Document⁴というデータにおいて、先行研究を大きく上回る結果を得ることができました。ちなみに、MS MARCO Documentの場合、文書が長くなるため、ColBERTは適用できません。

f:id:meshitahiro:20210712183955p:plain — MS MARCO Documentでの結果

最後に、次元の影響を調査しています。COILではたくさんのベクトルを転置インデックスとして保持するため、各単語のベクトルを通常のBERTで使用されるベクトルの次元768から線形変換で次元を削減しています。結果は以下の表の通りです。 $n_c$ はCLS-tokenの次元です。0次元はCLS-tokenを使用していないと意味です。 $n_t$ は、各単語の次元です。表の通り、次元を減らしても、若干の性能低下で済んでいることがわかります。一方、計算速度は、次元数を小さくするほど、速くなっています。 $n_t=8$ でCLS-tokenを使用しない場合、BM25に迫る速度で計算ができています。

f:id:meshitahiro:20210712184122p:plain — 次元が性能と計算速度に及ぼす影響

実験

さて、COILではコードが公開されているので、手元で再現実験を行いました。再現はMS MARCO Passageで行い、評価はDev Retrievalで行いました。結果は以下の通り、論文の値には満たないものの、良好な精度でした。

MRR @10: 0.320

この仕組みは、文脈化された単語ベクトルを使っています。そこで、全く学習をしない場合どうなるかを検証しました。以下の通り、あまり性能はよくありませんでした。世の中、教師なしでうまくいくほど甘くないですね。

MRR @10: 0.125

終わりに

今回は、COILという手法を紹介しました。驚いたことに、単語のミスマッチを解消しなくても、既存手法と同等の精度が出ています。また、COILを全く学習していないBERTのモデル上でどうなるか実験を行いました。残念ながら、全く学習していない場合は、あまりうまくいかない手法ということがわかりました。BERTを使った検索のモデルは近年盛んに研究されているため、また興味深い論文があれば、ご紹介したいと思います。

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. [paper]↩
ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT Omar Khattab, Matei Zaharia [paper]↩
MS MARCO passage: https://microsoft.github.io/msmarco/↩
MS MARCO document: https://microsoft.github.io/msmarco/↩