Improving Language Models by Retrieving from Trillions of Tokens

[shnakazawa]

Borgeaud, Sebastian, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, et al. 2021. “Improving Language Models by Retrieving from Trillions of Tokens.” arXiv [cs.CL]. arXiv. http://arxiv.org/abs/2112.04426.

• 何兆ものトークンを持つデータベースから検索しながら、任意のテキストをモデル化する "Retrieval-Enhanced Transformer (RETRO)" という手法を提案。
• 記憶を外部のデータソースに頼ることで、25倍規模のパラメータ数を持つモデルと同等の性能を得ることができた。
• 「セミパラメトリック」なアプローチの有用性を示した。

Abstract

We enhance auto-regressive language models by conditioning on document chunks retrieved from a large corpus, based on local similarity with preceding tokens. With a 2 trillion token database, our Retrieval-Enhanced Transformer (Retro) obtains comparable performance to GPT-3 and Jurassic-1 on the Pile, despite using 25x fewer parameters. After fine-tuning, Retro performance translates to downstream knowledge-intensive tasks such as question answering. Retro combines a frozen Bert retriever, a differentiable encoder and a chunked cross-attention mechanism to predict tokens based on an order of magnitude more data than what is typically consumed during training. We typically train Retro from scratch, yet can also rapidly Retrofit pre-trained transformers with retrieval and still achieve good performance. Our work opens up new avenues for improving language models through explicit memory at unprecedented scale.

(DeepL翻訳)

我々は、大規模コーパスから取得した文書チャンクを、先行トークンとの局所的な類似性に基づいて条件付けすることにより、自己回帰型言語モデルを強化する。2兆個のトークンデータベースを用いた我々の検索強化型変換器（Retro）は、25倍少ないパラメータで、Pile上のGPT-3やJurassic-1と同等の性能を得ることができる。Retroの性能は、微調整の後、質問応答のような下流の知識集約的なタスクに反映される。Retroは、凍結バートレトリバー、微分可能エンコーダー、チャンク型クロスアテンションメカニズムを組み合わせ、学習時に消費されるデータよりも一桁多いデータを基にトークンを予測します。私たちは通常、Retroをゼロから学習しますが、事前に学習した変換器を検索に迅速にRetrofitすることも可能であり、それでも良好な性能を達成することができます。私たちの研究は、前例のない規模の明示的記憶によって言語モデルを改善する新しい道を開くものです。

コード

• labmlai/annotated_deep_learning_paper_implementations
• lucidrains/RETRO-pytorch

解決した課題/先行研究との比較

• 近年の大規模自然言語処理モデルの高性能化は学習データの増加・計算能力の向上・モデルサイズの増加によって達成されている。
  • 実際、BERT (0.3B) → GPT-2 (1.5B) → T5 (11B) → GPT-3 (175B) → Gopher (280B) と進むにつれどんどん性能向上。
• モデルサイズの増加は「学習・推論の処理能力の増加」と「学習データの記憶力」という２つの利点があると考えられている。
• 本論文は、2つの利点がそれぞれどの程度効いているのかの分離を目指し、特に後者の「学習データの記憶力」という側面に着目。
• 記憶力に相当するものとして外部のデータベースを用いることで、モデル自体の計算量を大幅に増やすことなく、言語モデルを拡張する方法を提案した。
  • 過去にも検索を組み合わせる手法は行われてきた (Guu et al., 2020; Khandelwal et al., 2020; Lewis et al., 2020; Yogatama et al., 2021) が、モデルサイズやデータベースが小規模なものであったため、何兆ものトークンからなるデータベースを用いた初の報告。
    • 「過去の仕事を大規模にやってみました」の一種？
      • Chunked Cross Attentionは本論文で初出のように思われる。検索テキストの取り込み方に新規性！

技術・手法のポイント

• インプット文字列をチャンクに分割。現在のチャンクの予測のために、前のチャンクと似たテキストをデータベースから検索。
• 外部データベースからの検索には事前学習済みBERTを使用。近傍探索を行い、インプットに似たテキストを抽出。
• 検索されたテキストをChuncked Cross-Attentionモジュールを用いてRETROに取り込む。

評価指標

• 用いたデータセット

  • C4
  • Wikitext103
  • Curation Corpus
  • Lambada
  • Pile
  • マニュアルで選んだWikipediaの記事 (基準：データセットを集めた後に編集された記事)

• 比較した指標

  • Bits-per-byte
  • Perplexity
  • Lambadaデータセットはaccuracy on the last word
  • Q&A (The Natural Questions. Kwiatkowski et al., 2019) のAccuracy

• 比較対象のモデル

  • Transformer (パラメータサイズが 172M, 425M, 1.5B, 7.5Bのものを用意。Baselineと呼称)
  • RETRO (検索なし)
  • RETRO (検索あり)

いずれにおいてもBaselineからの改良がみられ、Fine-tuningを行うことでQ&A taskでもstate-of-the-artとのcompetitive performanceを示した。検索なしでもbaselineと同程度の性能が出る。

残された課題・議論

• 検索のための外部データベースを工夫すれば、差別や暴力的表現などを前もって除くことができるかもしれないと議論。 (Bender et al. 2021; Weidinger et al., 2021あたりも参照)
• 大規模モデルより軽量化したとはいえど、外部データベースのサイズが 1T トークンくらい必要そうなので、一般人が使うのは難しそう。

重要な引用

• 本論文以前の検索を組み合わせる手法。データベースのサイズ的に、一般人が現実的に使えるのはこれらか。
  • Guu, Kelvin, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. 13--18 Jul 2020. “Retrieval Augmented Language Model Pre-Training.” In Proceedings of the 37th International Conference on Machine Learning, edited by Hal Daumé Iii and Aarti Singh, 119:3929–38. Proceedings of Machine Learning Research. PMLR.
  • Khandelwal, Urvashi, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2020. “Generalization through Memorization: Nearest Neighbor Language Models.” https://openreview.net/pdf?id=HklBjCEKvH.
  • Lewis, Patrick, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, et al. 2020. “Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.” In Proceedings of the 34th International Conference on Neural Information Processing Systems, 9459–74. NIPS’20 793. Red Hook, NY, USA: Curran Associates Inc.
  • Yogatama, Dani, Cyprien de Masson d’Autume, and Lingpeng Kong. 2021. “Adaptive Semiparametric Language Models.” Transactions of the Association for Computational Linguistics 9: 362–73.
• Gopher. DeepMindが開発した2,800億個のパラメータを持つ言語モデル。モデルのサイズを大きくして性能が良くなる分野と大きく変わらない分野を議論
  • Rae, Jack W., Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, et al. 2021. “Scaling Language Models: Methods, Analysis & Insights from Training Gopher.” arXiv [cs.CL]. arXiv. http://arxiv.org/abs/2112.11446.
• 大規模言語モデルが持つ潜在的弊害について議論
  • Bender, Emily M., Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell. 2021. “On the Dangers of Stochastic Parrots: Can Language Models Be Too Big? 🦜.” In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, 610–23. FAccT ’21. New York, NY, USA: Association for Computing Machinery.
  • Weidinger, Laura, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, et al. 2021. “Ethical and Social Risks of Harm from Language Models.” arXiv [cs.CL]. arXiv. http://arxiv.org/abs/2112.04359.