ドキュメント画像RAGのための検索機能をJDocQAを使って比較(ColQwen2、Jina CLIP v2、OCR+テキスト埋め込み)
スライドやPDF資料などのドキュメントは、テキストデータだけでなく、図表やレイアウトに情報が含まれています。
このようなドキュメント群にRAGを行う場合、ページに含まれる画像をcaptioningしたり、レイアウトを考慮してチャンキングするなど、複雑なパイプラインを組むことで精度が上がるユースケースもあると思います。
このような構築・保守コストのかかる複雑な前処理パイプラインを避けつつ視覚的な情報も取り入れる手法として、6月にはDocument Screenshot Embedding(DSE)、7月はColPaliという検索モデルなどがリリースされ話題になっていたと思います。
これらを使えば、画像を直接ベクトルにして検索し、結果の画像を直接VLMに入力することで前処理が非常に簡略化されます。
今回は、JDocQAという日本語の画像質問応答データセットを使って、検索手法やモデルの比較を行いました。
後段のVLMの部分ではなく、検索の部分のみ評価しています。
データセット JDocQA
「JDocQA: 図表を含む日本語文書質問応答データセットによる 大規模言語モデルチューニング」という論文で構築されたデータセットで、官公庁や地方自治体の公開するドキュメントに対し、人手で質問を作成した画像質問応答のデータセットです。
これはVLMなどの画像質問応答タスクの学習・評価を目的に構築されたデータセットですが、今回はこれを画像質問に対する「検索」精度評価データセットとして扱います。
つまり、質問への回答が含まれるページを正解とし、それ以外のページを不正解として、質問に対して正解のページを検索できるかを評価します。
データセットには学習時のハルシネ傾向抑制のために回答が存在しない質問も含まれています。これを除外し、testセットの質問1,176クエリのうち、回答が可能な質問1,031クエリを評価対象にします。
検索対象のドキュメントは、testセットのクエリに紐づけられたページを含むドキュメント551ファイルの全ページを使用し、合計25,073ページを対象としました。
trainセットまで含めたデータセット全体には5,000以上のドキュメントがあるのですが、私の環境で処理するには規模が大きすぎるためtestセットに関連するファイルのみに絞っています。
これらを1ページずつpdf2imageを使ってjpgに変換し、ColQwen2やOCR+テキスト埋め込みモデルでベクトル化していきます。
↓イメージ図
クエリも画像側と同じモデルでベクトル化してドキュメントとの類似度を計算し、nDCG@10などの検索評価指標で評価します。
手法1 ColQwen2
冒頭で紹介したColPaliはVLMのPaliGemmaをベースに、クエリとドキュメント画像をペアとしたデータセットでColBERTスタイルのtoken embeddingモデルとして訓練したものです。
同じチームにより、ベースモデルをQwen2-VL-2Bに変更したColQwen2も出ており、日本語ではこちらの方が有力な選択肢になりそうです。
論文によると、テキストのみを使う場合(上段)と、画像のCaptioningを追加する場合(中段)に比べ、ColPali(下段)は複数のデータセットで高い検索精度を示しています。
手法2 Jina CLIP v2
画像の埋め込みモデルについては Jina CLIP v2 も今月2024/11にリリースされました。v2でマルチリンガル対応し、解像度も224x224から512x512に向上しています。
このモデルはそもそもドキュメント画像向けのモデルかはよく分かりません。v2のテクニカルレポートは11/24時点で出ておらず、どのようなタスクで学習・評価しているかは不明です。
手法3 OCR + テキスト埋め込みモデル
OCRとしては、日本語の縦書きにも対応したOCRであるNDLOCR v.2.1と、日本語対応している多言語OCRライブラリであるSuryaを試しました。Suryaはおそらく縦書き非対応です。
最近はPDFの前処理パイプラインを自動化するライブラリとしてDoclingなど色々と出てきていますが、今回はシンプルにOCRしたテキストをText Embeddingしています。
後段をVLMにしてOCR結果は検索に使うだけであれば、レイアウト等を考慮した前処理はあまり重要でなさそうな気がします。
テキスト埋め込みモデルはBGE-m3を使い、denseとcolbertベクトルを比較しました。
評価結果
検索評価の結果
評価セットの1,031クエリ、25,000ドキュメントに対して上記の手法を比較した結果が以下の表になります。
結果として、ColQwen2が最も良い結果となりました。
また、ColQwen2は検索のための学習には日本語データを使っていない(マルチリンガル性能はベースモデルのQwen2-VL頼り)なので、チューニングすることで向上の余地もあるかもしれません。この点は後述します。
次点で Surya + BGE-m3(colbert) となりました。このOCRモデルは縦書きを認識できないためその分精度は下がるはずですが、データセットの縦書きの割合は確認できておらず、影響の度合いは不明です。
同じSuryaのテキストで比較すると、テキスト埋め込みとしてはdenseよりもcolbertの方が5ポイント以上良くなりました。
BGE-m3のcolbertベクトルはprojectionがなく、そのままの1024次元xトークン数を出力するためストレージとメモリには注意が必要です。
ちなみに、ColQwen2にはprojectionがあり1トークンあたり128次元となっています。
NDLOCRは縦書きに対応しているものの、Suryaよりもテキスト行検出に失敗したページが多くなったせいか精度は低くなりました。
以下が、1ページ当たりのOCR結果の文字数をヒストグラムにしたものですが、50文字未満のページ数がSurya(オレンジ)と比べてNDLOCR(青)ではかなり多くなってしまっています。
例えば以下のような画像はSuryaではほぼ全体が文字起こしできましたが、NDLOCRでは認識に失敗し、"\n"のみが出力されました。
この点は私のNDLOCR利用時の設定か何かが正しくないかもしれず、実際はもう少し良いかもしれません。
Jina CLIP v2はRecall@10が0.14と非常に低い精度となりました。日本語ドキュメント画像の検索に適したモデルではなさそうです。
処理時間
処理はすべてRTX3090x1を225W電力制限して行いました。
隣にさしているGPUで別の処理を行っていたこともあったので、この処理時間は参考程度と考えてください。
OCRは25,000ページで6時間程度かかりました。
OCRはモデル自体は小さいのですが、テキスト行として検出されたすべての矩形に対してモデルを動かすために、トータルでかなり時間がかかったのだと思います。
Jina CLIP v2とColQwen2はモデルサイズはOCRよりも大きいですが、画像1枚に対して1回forwardするだけなので処理時間は短くなりました。
ColQwen2はJina CLIP v2より12倍時間がかかっていますが、パラメタサイズの差(2.2B 対 0.9B)に加えて、Jina CLIPでは画像サイズが縮小されていることが影響していると思います。
ColQwen2には幅720pxの画像を入力していますが、Jina CLIP v2では短い辺が512pxになるようにリサイズされるようです。
それを考慮してもJina CLIP v2に比べてColQwen2は時間がかかりすぎているように思うので、この測定にはやや納得できていないです。
失敗例の観察
OCRの失敗例
OCRではテキスト行検出に失敗して文字起こしが行われないと、その部分の情報が完全に落ちてしまうのが難点です。
例えば、「八王子神社は「はちおっつぁん」と呼ばれ住民に親しまれていますが、事故が起きたような言い伝えはありますか。」というクエリでは以下の画像が正解となります。
縦書き部分「八王子神社」の認識が必須ですが、Suryaでは画像下部の横書き部分しか文字起こしできず、検索に失敗します。
一方、ColQwen2ではこの画像をtop1で検索出来ています。縦書きも「八王子神社」ぐらいの部分は読めているということでしょう。
ColQwen2の失敗例
少ししか確認していませんが、ColQwen2は、難しめの漢字に弱いかもしれません。
クエリ「この前小瀬鵜飼をはじめてみましたが、春夏秋冬問わず行っていますか。参加したい場合はどのようにしたらいいですか。」に対し、正解は以下の画像です。
小瀬鵜飼という文字が複数個所に出てきており、OCRを使った手法ではうまくできましたが、ColQwen2ではtop1000以内にこの画像をランク付けできませんでした。
日本語の漢字画像の学習が足りておらず、「小瀬鵜飼」という文字に対応するトークン列と、その文字の画像の対応関係が理解できていないのではと想像しています。
このような点は、日本語OCR用のデータを使ってチューニングを行うことで改善できるかもしれません。