日本語 RAG タスクで e5-large 並みの性能の ColBERT
先日公開された、ColBERT の日本語pretrainモデル、JaColBERTの性能が良いらしい。早速、普段評価に利用している、AIクイズ王のQ&A RAGタスクで評価してみた。
結果としては、multilingual-e5-large よりわずかに低い、という結果となった。学習データも少なく、モデルサイズも12層のBert(= multilingual-e5-smallと同じサイズ)とほとんど一緒なのにね、すごい。
ColBERT の実装と論文を読む
というわけで、ColBERT に興味が湧いたので、論文と実装を読んでみる。
- 読んだ論文
- 実装
ColBERTはいわゆるSentenceTransformer等を用いて文からEmbeddings(文ベクトル)を出力して類似度比較で検索、ではなくトークンベースの類似度検索だ。文の最終隠れ層は、その文章のコンテキスト文脈の情報をトークン単位でも持っているので、文全体ではなくトークン単位で使って類似度を算出する。
類似度計算は、queryとdocumentのトークンの出力のコサイン類似度をとって、最大のものを足し合わせるという MaxSim という計算手法。MaxSimの計算自体は単純である。
query と document は別にエンコードする必要があるのだけど、使っているモデルは Bert (12層) + 独自HEAD (linear 128次元)。Bert 12層の隠れ層出力は768次元なので、それを128次元へと線形層を通して変換している。
実実装では、queryとdocumentは、query の場合は prefix に [CLS][unused0] を、document の場合は [CLS][unused1] といった形で CLS の後に判断のための独自トークンをつけて区別をしているだけで、エンコーダ自体は同じである。
これらでエンコードした結果を、MaxSimをとることで、一番数値が大きいdocumentがqueryに似ていると判定する。なおColBERT実装では、documentのtokenのうち、記号やpadding tokenなどは無視(mask)して計算していた。
検索時のパフォーマンス問題の解消
普通の文ベクトルを探す手法なら、ANNを始めとした近似最近傍探索を使えば、数億以上の文章からも高速な検索が可能だ。しかしながらColBERTは文ベクトルではなく、トークン同士の類似度を用いたMaxSim計算なので、この手法はママ使えない。
という問題に対応するために高速に近傍探索ができるindexを作る方法が載っているのが論文ColBERTv2。あれこれとベクトルを情報圧縮し、Kmeansでcentroidを求め、そこから探索しているようだ。実装ではfaissをimportしていたので、そのままfaissのindexを使っているのかな、と思いきや、faissはKmeanでクラスタ中心点を計算しているのに使っているだけであった。一度インデクスを作ってしまえば、その後の検索はサクサク。
なお、前述のAIクイズ王のQ&A RAGタスクの550万Passageをこの方法でindex作成するのに、Ryzen 3900 + RTX 3090の環境で5時間前後であった。faissはgpu利用版のfaiss-gpuでないと、かなり遅いので注意だ。
省データで学習可能
JaColBERT のレポートによると、bert-base-japanese-v3を元に、1000万のトリプレットデータとNVidia L4 GPU * 8 の省リソースで10時間学習させて作ったとのこと。データ量が少なく、かつ学習時間が短いとなると、夢が広がる結果である。
ColBERT の問題点・利用が大変
ColBERT の実装は読んでみて分かったが、処理の複雑性とそもそもの実装コードの複雑性も相まって、サクッと使うのが難しい。それを解消するために、ゼロコンフィグでもすぐに使えるアプローチが RAGatouille である。今回の評価にもRAGatouilleを利用した。
RAGatouille はインデクスの作成・検索はもちろん、Trainerで学習させることもできる。また、LangChain の retriver にすることもできるなど、今風の対応も入っている。
ColBERT は本番環境の検索サーバで使えそうか？
ColBERT の気になりどころとしては、そもそも本番で運用できそうかというところがある。今(2024/02頭)時点では、自前で検索用APIサーバ実装して運用が前提になり、手軽に動かすのがなかなか大変そうである。
しかしながらRAGatouilleのドキュメントによると、近日中に検索エンジンのVespaが対応するっぽいことが書かれているので、そうなるとだいぶ運用が楽になりそうである。またIndexのデータ追加もまだ試験的なようで、その辺も問題なくできるようになれば、本番環境でも十分検討できそうである。
何より、自前ドメインのデータを省コストで学習させることができるなら、自前データに対して高品質なRAG用のデータ取得検索として使えそうなので、そのようなユースケースなら積極的に検討したい。
おわりに
というわけで ColBERT すごい！というお話でした。情報検索を研究している人からすると ColBERT は既知の話かとは思うが、自分は知らなかったので調べてみて新鮮だった。また、JaColBERT が無ければそもそも興味がわかなかったと思うので、作者の Benjamin Clavié 氏に感謝。なお氏は RAGatouille の作者でもある、ありがたい。
