Zenn
📊

Transformer Memory as a Differentiable Search Index (NeurIPS 2022)

2023/01/03に公開
機械学習
情報検索
tech

論文紹介: Transformer Memory as a Differentiable Search Index (NeurIPS 2022)

この記事は情報検索・検索技術 Advent Calendar 2022 の 16 日目の記事です．

この記事では，NeurIPS 2022 に採択された T5 を用いた検索手法に関する Google Research の論文を紹介します．紹介する論文の情報は以下の通りです．

  • タイトル: Transformer Memory as a Differentiable Search Index
  • 著者: Yi Tay, Vinh Q. Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, Donald Metzler
    • 所属: Google Research
  • 採択会議: NeurIPS 2022 (2022 年 12 月)
    • arXiv における初出: 2022 年 2 月
  • 論文リンク

この記事の想定読者は以下の通りです．

  • 想定読者: 情報検索分野をある程度知っている人向け
    • Required: 検索システム ― 実務者のための開発改善ガイドブックの第一部の内容を把握している
      • クエリ，文書，転置インデックス，TF-IDF/BM25 などの用語を聞いたことがある
    • Preferred: 最近の自然言語処理でよく用いられている，大規模言語モデルについて把握している
      • 大規模言語モデル (事前学習済み言語モデル/ニューラル言語モデル)，BERT/BART/T5，系列変換モデル (seq2seq) などの用語を聞いたことがある

この論文がやったこと

まず，この論文の問題設定について説明します．この論文はアドホック文書検索を問題としています．つまり，入力としてはクエリとしてキーワードや質問文などが与えられ，出力としてはクエリに適合する文書のランキングを返します．これは Web 検索などと同様の一般的な文書検索の問題設定です．たとえば下の図のように，Web 検索においては「つくばでおすすめの観光地は？」というクエリが与えられると，じゃらんや Yahoo!トラベルなどのつくば市内のおすすめの観光スポットを紹介する Web ページ (文書) が上位にくるようなランキングが返されると考えられます．

次にこの論文が提案した手法の概要を説明します．この論文では，系列変換モデルを用いた検索モデルを提案しています．系列変換モデルとは，ある系列データ (ここでは文字列からなる文) を入力すると，その系列データに依存した別の系列データを出力する機械学習モデルのことです．seq2seq や encoder-decoder とも呼ばれており，近年の代表的な手法として BART[1] や T5[2] が挙げられます．系列変換モデルが用いられる典型的なタスクとしては，ある言語の文を入力として，別の言語の対応する文を出力とする機械翻訳や，ある文書を入力として，その文書を要約した文を出力する文書要約が挙げられます．

本論文では系列変換モデルであるである T5 を用いて，クエリから文書 ID などの文書に一意に対応する文字列を生成し，その文字列の生成確率でランキングを作成します．たとえば下図のように，「つくばでおすすめの観光地は？」というクエリが与えられると，doc42 などの文書 ID を生成し，その生成確率の降順でソートしたランキングを作成します．

著者らはこの手法のことを微分可能検索インデックス (Differentiable Search Index) と呼んでいます．ここでの「インデックス」は転置インデックスなどのインデックスと似た意味で著者らは用いていると考えてよいです．転置インデックスはクエリを入力として，クエリに含まれる単語を含むような文書の ID と単語頻度などの統計情報を返すものです．一方で，今回提案されている微分可能検索インデックスもクエリを入力して文書 ID とその ID の生成確率を返すものといえます．よって微分可能検索インデックスは転置インデックスと同様に，クエリを入力として文書 ID を返す，という点で共通しています．微分可能検索インデックスにおいては T5 が用いられており，転置インデックスとは違って微分可能である (ここでは訓練データを用いて学習が可能であるくらいの意味) ということで著者らは微分可能検索インデックスと呼んでいるようです．

ではここからはこの論文が登場してきた背景について，その先行研究を紹介しながら説明していきます．

背景: 生成的エンティティ検索モデル GENRE

この論文の提案手法のように，「クエリを系列変換モデルに入力して文書に一意に対応する文字列を生成し，その文字列の生成確率でランキングをする」という方法は，この論文の後続論文において生成的検索 (Generative Retrieval) と呼ばれており，この論文が 2022 年 2 月に arXiv で公開されて以来複数の論文が出版されています．

では生成的検索という手法はこの論文が初めて提案したものなのでしょうか？実はそうではありません．生成的検索を初めて提案したのは Autoregressive Entity Retrieval (ICLR 2021)[3] という論文です．

この論文はエンティティ検索に取り組んでいます．エンティティとは，なんらか参照可能な「モノ」や「概念」のことで，典型的なエンティティとしては人物（例: バラク・オバマ)，組織 (例: 筑波大学)，地名 (例: つくば市) などが挙げられます．よってエンティティ検索とは，クエリを入力としてクエリに適合するエンティティのランキングを返すタスクとなります．例えば「つくばでおすすめの観光地は？」というクエリに対しては「筑波山」や「つくばエキスポセンター」などの地名が上位にくるランキングが返されます．

この問題に対して，著者らは GENRE (Generate Entity Names auroREgressively) という手法を提案しています．GENRE では下の図のように，クエリを系列変換モデルである BART の入力として，エンティティに対応する文字列としてエンティティ名を生成し，その生成確率でランキングを作成しています．本論文においてエンティティ名は全てのエンティティで一意であるという仮定が置かれているため，エンティティ名とエンティティ名の生成確率をもとにランキングをすれば，対応するエンティティのランキングが分かります．

系列変換モデルの学習は翻訳などの通常の系列変換モデルが用いられるタスクの学習と同様に行われます．つまり，クエリと適合エンティティ名の組からなるペアのデータを大量に用意し，この学習データを用いてクエリを入力としてエンティティ名を生成するモデルを教師ありで学習します．

ここで述べておきたい点として，GENRE はエンティティに付随するエンティティの説明文 (entity description) やエンティティの属性を用いずに学習を行うという特徴があります．例えば「筑波大学」というエンティティには，「筑波大学は茨城県にある大学である」という説明文や，筑波大学は「組織名」であるという属性情報が付随しています．通常のエンティティ検索においては，これらの情報も用いて検索モデルを学習しますが，GENRE においてはこれらを用いません．これにより新語などの新しいエンティティに対する検索も，そのエンティティ名の生成確率を計算するればランク付け可能という特徴があります．

実験についても少し触れておきます．データセットとしては，近年のエンティティ関連タスクとして頻繁に用いられている KILT データセット[4] を利用しています．このデータセットにおける実験において提案手法である GENRE は，単語の厳密一致による検索手法である TF-IDF や，近年性能が高いと言われている密(ベクトル)検索手法の Dense Passage Retrieval [5] などと比較して良い性能であることが確認されています．

提案: 文書検索への適用

今回紹介する論文がやったことを一言で述べるとするならば「GENRE を文書検索に適用した」ということになります．つまり GENRE がエンティティを検索対象にしていたことに対して，提案手法である Differentiable Search Index (以下 DSI) はより一般の文書を検索対象としています．

ここでの「文書」とは何を指すのでしょうか？しかしながら，本論文における「文書」がどのようなドメイン (例: Web) のテキストデータを指すかは論文中では触れられていません．よって，手法として文書のドメインに関する制約はないものとこの論文の著者は考えているのではないかと推測しています．ただし実際に実験で用いられているデータセットは Natural Questions[6] という質問応答のデータセットであり，クエリは質問文，検索対象の文書は Wikipedia の文書 (をパラグラフなどで分割したもの) となっています．そのため，他のドメインの文書にもそのまま適用できるかどうかは未知数であるとも言えます．

文書検索のための訓練方法: 文書 ID を大規模言語モデルに教える

では，GENRE が検索対象としていたエンティティと本論文の DSI が対象としている (一般の) 文書では何が異なるのでしょうか？言い換えると，何が解決する技術的な課題となるのでしょうか？答えは，文書検索においてはクエリに対して何を生成すべきかが自明ではないという点です．エンティティ検索においては，エンティティ名が全てのエンティティで一意であり，エンティティ名を生成すれば対応するエンティティを元にランキングが作成できます．しかしながら文書検索においては，エンティティにおけるエンティティ名とは違って，一意に対応する意味のある文字列が存在するとは限りません．

一意な文字列を得る単純な方法としては，文書ごとに一意な文書 ID を付与し，この ID を生成することが考えられます．しかしながらエンティティ名とは違って，大規模言語モデルは付与された文書 ID が何を表すのか，何を意味するのかを知りません．そのため，単純に GENRE と同様にクエリから文書 ID を生成するように系列変換モデルを学習するだけでは性能の良い検索モデルを得ることはできないことが予想されます．

そこで本論文では，通常のクエリを入力として文書 ID を生成するタスクに加えて，文書を入力として文書 ID を生成するタスクも同時に学習することを提案しています．この後者のタスクによって，大規模言語モデルに文書 ID がどのような文書であるのかを学習させていると見ることができます．論文においては前者のクエリから文書 ID を生成するタスクを Rertieval と呼び，後者の文書から文書 ID を生成するタスクを Indexing と呼んで区別しています．

文書 ID の割り振り: 似た文書が似た ID を持つようにする

また本論文では文書 ID の表現についても検討を行っています．単純にはランダムに割り当てられた一意な整数や文字列として表現できます．例えば本論文では以下の 2 つの単純な方法があるとしています．

  • Atomic: ランダムな 1 つの整数を割り振る．生成時にはそれぞれの ID を 1 つのトークンとして扱う．
  • Naive String: Atomic の ID を文字列として扱う．生成時にはサブワードに分割されうる．

しかしこのように割り当てられた ID は特に意味を持ちません．そこで本論文では似た文書が似た ID を持つような ID の割り振り方として，semantically structured identifier を提案しています．この論文において ID が似ているとは，2 つの ID が同じ接頭辞 (prefix) を持つことと定義しています．たとえば，ID:123 と ID:124 の組は ID:123 と ID:223 の組よりも ID が似ています．これは，前者が同じ接頭辞として 12 を持っていることに対して，後者は同じ接頭辞を持っていないためです．

このように似た文書の接頭辞を似たようにする理由としては，文書 ID を生成する際に有用であるという点が挙げられます．検索においては似たような文書は似た情報を持っていることが多いと考えられます．そのためあるクエリに対して似た 2 つの文書は，どちらも適合もしくはどちらも不適合のいずれかになりやすいことが予想されます．よって，T5 や BART のような生成時に前から順に (自己回帰的に) 文字列を生成する系列変換モデルにおいて，似た文書が同じ接頭辞を持つような ID を割り振ると，その生成の途中の段階で適合文書を含むような接頭辞が上位となりやすく，最終的なランキングにも適合文書が含まれやすいと考えられます．

例えば，ID として 123 と 124 という 2 つの ID を考えます．まず提案手法を使わない場合，つまり似た文書が同じ接頭辞を持たないように ID を割り振った場合を見てみます．このとき文書 123 が適合文書であっても，文書 123 と文書 124 は似ておらず，かつ適合文書は膨大な文書集合全体のほんの一部であることを考えると，文書 124 は不適合である可能性が高いです．そのため，2 つの文書のうち一方のみが不適合文書であるという状況において，ID の生成の過程において 1→2 と ID を生成する確率はそこまで高くならないと考えられます．一方で提案手法のように似た文書が同じ接頭辞を持つように ID を割り振ると，文書 123 が適合文書である場合，文書 124 も適合文書である可能性が高いと考えられます．そのため，ID の生成の過程において 1→2 と ID を生成する確率は提案手法を用いない場合よりも高くなると考えられます．

では実際に似た文書が同じ接頭辞を持つように ID を割り振る方法を説明します．本論文では文書ベクトルをもとに再帰的にクラスタリングを行うことで，トライ木の一種である Decimal Tree を導出します (下図)．具体的には，文書集合を k-means 法で C 個のクラスタに分割します．分割したクラスタに整数を割り振った後，さらに各クラスタを再度 C 個のクラスタに k-means 法で分割します．これを，クラスタ内の文書数が c 個以下になるまで再帰的に実行します．下図は，C=3, c=4 としてこの手順を実施した図とみなせます．実際の実験では C=10, c=100 が用いられています．

実験

本論文では質問応答データセットを用いて，提案手法とベースライン手法を比較するためにいくつか実験を行い，提案手法の有効性を検証しています．

実験設定と比較手法

本論文の実験設定 (データセット，評価指標，使用する系列変換モデル，ベースライン手法) は以下の通りです

  • データセット: Natural Questions[6:1] (質問応答データセット)
    • クエリ: 質問文，文書: Wikipedia (をパラグラフなどで分割したテキスト)
    • データセットサイズが小さい場合も比較するために，NQ10K/NQ100K/NQ320K を用意して比較
  • 評価指標: Hits@N (N=1,10)
  • 系列変換モデル: T5 を利用
    • モデルサイズ: Base (250M), Large (800M), XL (3B), XL (11B)
      • カッコ内はパラメータ数
  • 比較手法
    • 単語の厳密一致: BM25
    • 密(ベクトル)検索: T5 ベースの dual-encoder モデル[7]

実験 1: 教師あり設定の比較

まずメインの実験として，上記の質問応答データセットで教師あり学習した実験を行っています．結果は以下の表の通りです．最上段が BM25，次の 4 行が密(ベクトル)検索手法，その下は全て提案手法である DSI であり，文書 ID の割り振り方法 (Atomic, Naive, Semantic) ごとに結果が算出されています．最後の 4 行 (Semantic) が提案した ID の割り振り方で，似た ID が同じ接頭辞を持つように ID を割り振る方法です．この表の太字の部分から，全ての評価指標の全てのサイズのデータセットにおいて，提案手法のいずれかがベースライン手法よりも上回っていることが分かります．

この表は少し大きいので一部を切り出して見てみます．例えば T5-base という比較的小さいサイズのモデルを用いた場合を比較してみると，下図のようにデータセットのサイズによって異なる結果が出ており，必ずしも提案手法が良いと言える結果が出るわけではないようです．


一方で，T5-XXL という非常に大きいモデルを用いた場合の結果を見てみると，ID の工夫のあり/なしはあれど，提案手法が密(ベクトル)検索手法よりも良い結果となっていることが多いことが分かります．ここから，モデルサイズが大きいほうが提案手法の恩恵を受けやすい傾向があると言えそうです．

実験 2: ゼロショット設定の比較

また別の実験としてゼロショット設定の実験をしています．ここでのゼロショット設定とは，教師データであるクエリと適合文書の組を用いたモデルの学習を行えない設定のことです．特に提案手法におけるゼロショット設定とは，従来の学習タスクである「クエリ→文書 ID」の生成タスクと提案手法独自の学習タスクである「文書→文書 ID」の生成タスクのうち，前者の従来の学習タスクをせずに後者のみを行う設定となっています．

この設定における実験結果は以下の通りです．今回も教師あり設定と同様に，ID の工夫のあり/なしはあれど，提案手法のいずれかが密(ベクトル)検索手法よりも良い結果となっていることが分かります．

まとめと感想

本論文のまとめは以下の通りです．

この論文を読んだ感想や疑問点としては以下の通りです．

  • T5 で全部やっちゃえというかなり狂った⼤胆な発想
    • ニューラル⾔語モデルの⼒をどこまで信じればこんな発想ができるのだろうか？
    • いかにも BigTech っぽい，計算リソースを⼤量に使って性能を出した論⽂という印象
      • T5-XXL の学習に 128-256 個の TPUv4 を利⽤して (少なくとも) 丸 1 ⽇らしい
  • 推論の efficiency についての実験結果が⼀切ない
    • おそらく推論の efficiency に関してはかなり悪いのだろうと予想できる
    • これは実用においてはかなり厳しく，今後の改善が期待される
  • 質問応答以外のデータセットでの実験がない
    • MS MARCO Document Ranking などの⻑い⽂書でも可能なのか？
    • Robust04 のような(より)⼩さいデータセットでも同様か？

その後の発展

GENRE や DSI のような，「クエリを系列変換モデルに入力して文書に一意に対応する文字列を生成し，その文字列の生成確率でランキングをする」する生成的検索は 2022 年に入ってから多くの論文が出版されています．調べた限りで出てきた論文をまとめたものが以下の表です．NeurIPS 2022 にも 3 本の論文が採択され，そのうち 1 本が outstanding paper (いわゆる best paper) を取るなど盛り上がりを見せています．

手法名 データセット 初出 会議・ジャーナル
GENRE[3:1] KILT 2020/10 ICLR 2021 (spotlight)
DSI (本論文) NQ 2022/02 NeurIPS 2022
DynamicRetriever[8] MS MARCO 2022/03 -
SEAL[9] KILT 2022/04 NeurIPS 2022
GERE[10] FEVER 2022/04 SIGIR 2022 (short)
NCI[11] NQ, TriviaQA 2022/06 NeurIPS 2022 (outstanding)
DSI-QG[12] MS MARCO, XOR QA 2022/06 -
CorpusBrain[13] KILT 2022/08 CIKM 2022
Ultron[14] MS MARCO, NQ 2022/08 -
CGR[15] KILT 2022/10 -
DSI++[16] MS MARCO, NQ 2022/10 -
  • データセットについて
    • エンティティ関連タスクのデータセット: KILT[4:1]
    • 質問応答データセット: NQ (Natural Questions)[6:2]，Trivia QA[17]，XOR QA[18]
    • 検索データセット: MS MARCO[19]

おまけ: 著者と関連文献について

本論文の著者は全員 Google Research の所属です．第一著者である Yi Tay 博士は Transformer 系の大規模言語モデルに関する研究を精力的に行っている方で，Sparse Sinkhorn Attention[20] や効率的な Transformer に関するサーベイ論文[21]で有名な印象です．情報検索の論文は少ないですが，ランキング学習における勾配ブースティング木とニューラルネットを比較する論文[22]を書いており，これは面白い論文でした．

著者の中で SIGIR などの情報検索の研究コミュニティで論文を書いている方(書いていた方)としては，最終著者である Donald Metzler 博士が最も有名な印象です．例えば bi-gram を使った古典的な検索モデルである Sequential Dependence Model[23]，Listwise なランキング学習の手法である Coordinate Ascent[24]，評価指標の ERR[25]，検索エンジンの Indri[26] などの論文の著者です．もともとはいわゆる IR 暁の四戦士の一人である Croft 先生の元で博士を取られた方で，Croft 先生の本[27]の第 2 著者でもあります．

また今回紹介した論文に関連するものとして，本論文の著者である Donald Metzler 博士や Yi Tay 博士などが書いた Rethinking Search という論文 (正確には Research Proposal)[28] があります．この論文は昨年 (2021 年) に以下の紹介ツイートにより，情報検索界隈で少し話題になりました．この論文で提唱されている Model-based IR を実現する方法の一つが今回紹介した DSI であると見ることもできます．

https://twitter.com/rejasupotaro/status/1394107056559312898?s=20

謝辞

Advent Calendar を企画してくださった @hurutoriya さんありがとうございました．また，この文書は IR Reading 2022 秋において発表させていただいたスライドを元に作成しており，このような会を開催・運営して下さった ACM SIGIR 東京支部の役員の皆様に感謝いたします．また，IR Reading 2022 秋の発表に対して質問・コメントして下さった佃先生，飯田さんに感謝いたします．

参考リンク

脚注

  1. Lewis et al.: BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension. ACL 2020. https://aclanthology.org/2020.acl-main.703/ ↩︎

  2. Raffel et al.: Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR, 2020. https://jmlr.org/papers/v21/20-074.html ↩︎

  3. De Cao et al.: Autoregressive Entity Retrieval. ICLR 2021. https://openreview.net/forum?id=5k8F6UU39V ↩︎ ↩︎

  4. Petroni et al.: KILT: a Benchmark for Knowledge Intensive Language Tasks. NAACL 2021. https://aclanthology.org/2021.naacl-main.200/ ↩︎ ↩︎

  5. Karpukhin et al.: Dense Passage Retrieval for Open-Domain Question Answering. EMNLP 2020. https://aclanthology.org/2020.emnlp-main.550/ ↩︎

  6. Kwiatkowski et al.: Natural Questions: a Benchmark for Question Answering Research. TACL, 2019. https://doi.org/10.1162/tacl_a_00276 ↩︎ ↩︎ ↩︎

  7. Ni et al.: Sentence-T5: Scalable Sentence Encoders from Pre-trained Text-to-Text Models. Findings of ACL 2022. https://aclanthology.org/2022.findings-acl.146/ ↩︎

  8. Zhou et al.: DynamicRetriever: A Pre-training Model-based IR System with Neither Sparse nor Dense Index. 2022. https://arxiv.org/abs/2203.00537 ↩︎

  9. Bevilacqua et al.: Autoregressive Search Engines: Generating Substrings as Document Identifiers. NeurIPS 2022. https://openreview.net/forum?id=Z4kZxAjg8Y ↩︎

  10. Chen et al.: GERE: Generative Evidence Retrieval for Fact Verification. SIGIR 2022. https://dl.acm.org/doi/10.1145/3477495.3531827 ↩︎

  11. Wang et al.: A Neural Corpus Indexer for Document Retrieval. NeurIPS 2022. https://openreview.net/forum?id=fSfcEYQP_qc ↩︎

  12. Zhuang et al.: Bridging the Gap Between Indexing and Retrieval for Differentiable Search Index with Query Generation. 2022. https://arxiv.org/abs/2206.10128 ↩︎

  13. Chen et al.: CorpusBrain: Pre-train a Generative Retrieval Model for Knowledge-Intensive Language Tasks. CIKM 2022. https://dl.acm.org/doi/abs/10.1145/3511808.3557271 ↩︎

  14. Zhou et al.: Ultron: An Ultimate Retriever on Corpus with a Model-based Indexer. 2022. https://arxiv.org/abs/2208.09257 ↩︎

  15. Lee et al.: Contextualized Generative Retrieval. 2022. https://arxiv.org/abs/2210.02068 ↩︎

  16. Mehta et al.: DSI++: Updating Transformer Memory with New Documents. 2022. https://arxiv.org/abs/2212.09744 ↩︎

  17. Joshi et al.: TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. ACL 2017. https://aclanthology.org/P17-1147/ ↩︎

  18. Asai et al.: XOR QA: Cross-lingual Open-Retrieval Question Answering. NAACL-HLT 2021. https://aclanthology.org/2021.naacl-main.46/ ↩︎

  19. Nguyen et al.: MS MARCO: A Human Generated MAchine Reading COmprehension Dataset. CoCo Workshop at NeurIPS 2016. http://ceur-ws.org/Vol-1773/CoCoNIPS_2016_paper9.pdf ↩︎

  20. Tay et al.: Sparse Sinkhorn Attention. ICML 2020. https://proceedings.mlr.press/v119/tay20a.html ↩︎

  21. Tay et al.: Efficient Transformers: A Survey. 2022. https://arxiv.org/abs/2009.06732 ↩︎

  22. Qin et al.: Are Neural Rankers still Outperformed by Gradient Boosted Decision Trees? ICLR 2021. https://openreview.net/forum?id=Ut1vF_q_vC ↩︎

  23. Metzler et al.: A Markov Random Field Model for Term Dependencies. SIGIR 2005. https://dl.acm.org/doi/abs/10.1145/1076034.1076115 ↩︎

  24. Metzler et al.: Linear Feature-based Models for Information Retrieval. Information Retrieval, 2007. https://link.springer.com/article/10.1007/s10791-006-9019-z ↩︎

  25. Chapelle et al.: Expected Reciprocal Rank for Graded Relevance. CIKM 2009. https://dl.acm.org/doi/abs/10.1145/1645953.1646033 ↩︎

  26. Strohman et al.: Indri: A Language Model-based Search Engine for Complex Queries. ICIA 2005. http://ciir.cs.umass.edu/pubfiles/ir-407.pdf ↩︎

  27. Croft et al.: Search Engines: Information Retrieval in Practice. 2009. http://www.search-engines-book.com/ ↩︎

  28. Donald Metzler, Yi Tay, Dara Bahri, Marc Najork: Rethinking Search: Making Domain Experts out of Dilettantes. SIGIR Forum, 2021. https://dl.acm.org/doi/abs/10.1145/3476415.3476428 https://arxiv.org/abs/2105.02274 ↩︎

Discussion