Vポイントマーケティング|TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

“Retrieval-augmented in-context learning”を実現する、DSP(Demonstrate-Search-Predict)の論文を読んでまとめてみました。

こんにちは、CCCMKホールディングスTECH LAB三浦です。

最近ふとしたきっかけで読んでみたコミックがとても面白くて、良い出会いをしたなぁとしみじみと感じています。コミックだけでなく、映画とか音楽もこれまで知らなかったけど触れてみたらとてもお気に入りの作品になった、ということがあったりします。こういう出会いは本当にいいものだな、と思います。

前回DSPyという、LLMを使った一連の処理パイプラインをModuleを組み合わせて構築することが出来、さらにModuleのパラメータを機械学習のように学習データで調整することが出来るフレームワークについてまとめました。

techblog.cccmkhd.co.jp

プロンプトエンジニアリングをすることなくLLMパイプラインを最適化出来るアプローチがとても新鮮で興味を持ち、それ以来論文などを読んだりしてそのアプローチの仕組みなどを調べています。

実はDSPyにはDemonstrate-Search-Predict (DSP)という前のバージョンが存在します。DSPはDemonstrate, Search, Predictという3つの推論ステージを通じて複雑なタスクをLLMに解かせることを可能にするフレームワークで、3つの推論ステージの実装はあらかじめ組み込まれたModule(DSPではTransformationと呼ばれています)を組み合わせてプロンプトエンジニアリングではなくプログラミングすることで実装することが可能です。

DSPが提案されている論文がこちらです。

  • Title: Demonstrate-Search-Predict: Composing retrieval and language models for knowledge-intensive NLP
  • Authors: Omar Khattab, Keshav Santhanam, Xiang Lisa Li, David Hall, Percy Liang, Christopher Potts, Matei Zaharia
  • Submitted: 28 Dec 2022 (v1), last revised 23 Jan 2023
  • arXivURL: https://arxiv.org/abs/2212.14024

今回はこの論文を読んでDSPで作った一連のLLMプログラムがどのように実行されるのかについて理解したことをまとめてみたいと思います。特にDemonstrateのステージで行われていることは、DSPyの学習フェーズで行われていることに近いと思います。

Demonstrate-Search-Predict (DSP)とは

Demonstrate-Search-Predict (DSP)はin-context learningRetrieval-augmented generation(RAG)を組み合わせたLLMパイプライン: "Retrieval-augmented in-context learning"をモジュールを組み合わせて構築することを可能にするフレームワークです。In-context learningはプロンプトの中に与えられたタスクを解くために例示を含めるテクニックで、RAGは与えられた質問に関連する情報を検索してプロンプトに組み込むことでデータを増強し、様々なドメインに関する質問に回答出来るようにするテクニックです。DSPはこの2つのテクニックを組み合わせたLLMパイプラインをあらかじめ用意されたモジュールを組み合わせた"LLMプログラム"で表現することが可能です。

DSPによって可能になること

in-context learningやRAGはどちらもよく活用されているテクニックですが、それらをシンプルに使うだけでは解くことが難しい複雑な問題もあります。たとえば、「三浦が働いている会社の住所は?」といった質問が与えられたとします。私たちはこういった問題を解くとき、いくつかのステップを踏んで問題に取り組むのではないでしょうか?たとえば最初に「三浦が働いている会社は?」という質問を考えて情報を検索して解決し、そのあとで「その会社の住所は?」という次の質問を考えてまた情報を検索し解決することで最終的な回答にたどり着くと思います。

これは結構難しい問題で、特に与えられた質問からどのように細かい質問を生成するのか、とか、前の質問の結果からどうやって次の質問を生成したらいいのかなどは具体的な例示、つまりin-context learningのテクニックが必要になるかもしれません。人の手によるプロンプトエンジニアリングでこれを実現しようとするとかなり煩雑な作業になってしまいますが、DSPではまずどうやってタスクを解決するのかをTransformationというモジュールを組み合わせた"LLMプログラム"を実装して表現した後、解決したい質問と同じタスクに属する別の質問と正解データで構成される学習データを与えることでin-context learningに必要になる例示(Demos)を生成し、回答精度を向上させることを可能にします。

DSPの処理の流れ

DSPの処理はDemonstrate→Search→Predictの順に推論ステージを遷移して質問に対する回答を生成します。具体的にどのように処理が実行されるのかは論文に掲載されているこちらの図で表現されています。

DEMONSTRATE–SEARCH–PREDICT: Composing retrieval and language models, Figure 2

この図で表現されていることとして、まずDSPのプログラムではdsp.Example型の変数が入力され(図中のx)、ステージを経過するにつれどんどん新しいデータが追加されていきます(図中の"Demos"や"Hop1","Psg1"などです)。前のステージで追加されたこれらのデータは、次のステージで必要に応じて参照されます。たとえばDemonstrateステージで追加された例示"Demos"が、Searchステージで質問を別の質問に書き換えるときにin-context learningに使われたりします。

処理はDemonstrate→Search→Predictの順に行われるのですが、実際にプログラムを実装するときは最初にSearch, Predictのステージの実装をしてから最後にDemonstrateステージの実装をする流れになると思います。なのでここではSearch, Predictのステージについてまとめた後、最後にDemonstrateのステージについてまとめていきたいと思います。

このステージは質問に関連する情報を検索して情報をリッチにして次のPredictのステージに渡す役割を担います。シンプルに質問に関連する情報を検索して追加情報として加える、という処理ももちろん可能ですが、DSPではもっと複雑な、たとえば質問を別の質問に置き換えて情報を検索する、さらにそれを複数回繰り返す、といった処理をこのステージで行うことも可能です。

実際先の図中では"multi-hop question answering"というテクニックが記載されていて、Searchステージで質問の生成と関連情報の取得、といった処理が2回行われています。1回質問の生成と関連情報の取得を行うと、dsp.Example型のxhop1psg1という新しいデータが追加され、次のhop2psg2の生成の際に参照されます。また質問の生成の際には前のステージで生成されxに格納されている例示demosを参照させることが可能です。

Predict

これまでのステージで追加されたデータを参照して質問に対する最終的な回答を生成するステージです。このステージではたとえばChain of Thought(CoT)やSelf-Consistencyという回答生成を複数回実行した後に多数決で最終的な回答を決める、といったテクニックを実装することが出来ます。

Demonstrate

Demonstrateはタスクを解くために取るべき具体的な例示であるDemosを生成するステージで、in-context learningを実現するための重要なステージです。DSPでは回答してほしい質問と同時に学習データ(図中のTrain)をプログラムに入力します。Demonstrateでは学習データに対しLLMプログラムを実行し、LLMプログラムによる回答を生成します。その過程で先に説明したようにSearchとPredictの処理が実行され、それらの中間処理で生成された質問や関連情報(hop1psg1など)が追加されていきます。

たとえば先の図では、学習データとして質問"When was the discoverer of Palomar 4 born?"と回答"1889"が与えられています。ここで組んでいるLLMプログラムではSearchのステージで2回、質問の書き換えと検索処理が実行されます。その結果、hop1(生成された質問1), psg1(関連情報1), hop2(生成された質問2), psg2(関連情報2)という属性をキーを持つdsp.Example型の値がdemosに追加されます。そしてPredictのステージで学習データの質問に対する最終的な回答が生成されpredもまた新しい情報としてdemosに格納されたdsp.Example型の値に追加されます。

Demonstrateのステージではこのように最初は質問と回答しか含まれていなかった学習データに対し、LLMプログラムを通じて生成される中間情報が追加され以降のステージで使用されるin-context learning用のデータが生成されるのですが、それだけでなく学習データのフィルタリングを行うことも出来ます。たとえば先の図では"Q: When was the discoverer of Palomar 4 born? A: 1889"と"Q: In which city did Akeem Ellis play in 2017? A: Ellesmere Port"の2つの学習データが与えられていて、1つ目の学習データに対してはLLMプログラムの処理によって正解データと同じ回答が生成出来ていますが、2つ目の方は一致する回答が生成出来ていません。この場合2つ目の学習データは例示として含めるのに相応しくないと考えられ、以降のステージで使わないようにすることが出来ます。

DSPとDSPy

DSPはDSPyの前のバージョンで、DSPyで何をやっているのかを理解するうえでDSPを理解することは重要だと思います。前回のブログでもまとめたのですが、特にDSPyのcomplileを実行した時のLLMパイプライン最適化の処理の中ではDSPのDemonstrateと同じような処理が実行され、自動的にin-context learningに必要な例示データが生成されていると考えられるからです。

まとめ

ということで、今回はDemonstrate-Search-Predict (DSP)というRetrieval-augmented in-context learningを実現するためのフレームワークについて、具体的にどのような仕組みなのかを論文を読んでまとめてみました。自分の中ではかなり新鮮な切り口のアプローチだったので、なかなか理解するのに時間がかかってしまいました。ただこういうアプローチを見ていると、プロンプトエンジニアリングをすることなく安定してLLMの性能を引き出すテクニックが今後どんどん誕生してくるのでは、と感じました。