【DL輪読会】In-Context Pretraining: Language Modeling Beyond Document Boundaries

In-Context Pretraining: Language Modeling Beyond Document Boundaries Toru Fujino 1

書誌情報 ● ICLR 2024 (spotlight) ● 概要 ○ Large Language Model（LLM）の学習サンプルを作成する際に、文書をランダムに連結さ せるのではなく、関連した文書同士を連結させる手法の提案 提案手法 従来手法 ○ 長い文脈の理解が要求されるタスクにおいて既存手法よりも精度が向上 ● スライド中の図・表は論文から引用

LLMの課題 ● LLMは様々なタスクで良い結果を出しているが、複雑な文脈の理解が要求さ れるタスクで苦労することがある ○ プロンプトのinstructionに正確に従わない ■ ○ 重要な情報がプロンプトのどの位置にあるかで精度が大きく変わってしまう ■ ○ [McKenzie+ 2023] [Liu+ 2023] プロンプトの文章の微妙な変化が大きな精度の低下につながる ■ [Zhao+ 2021]

LLMの学習サンプルの構築方法 ● 学習コーパス内の文章をランダムな順序で連結した上でnext token predictionを行う ● 同一サンプル内において、前の文章は後ろの文章に何のシグナルも与えな い ○ 不要な計算のオーバーヘッドが発生 ● 類似した文章を連結させることで、意味のある文脈を持ったサンプルを作成 できないか

提案手法：In-Context Pretraining ● 類似文書を連結して学習サンプルを作成する ○ LLMが学習時に長い文脈に触れることができるようになる

In-Context Pretraining（におけるサンプル作成）の流れ 1. 類似した文書の検索・グラフの構築 ● 各文書に対し、その文書と類似した文書を抽出する ● 文書をノード、類似度をノード間のエッジの重みとしてグラフを構築する 2. 類似度の合計が最大となるグラフ上のパスの探索 ● すべてのノードを一度だけ通り、かつ通ったエッジの重みの合計値が最 大となるパスを見つける ○ ○ ノードを二度以上通ると、学習で出現する文書に偏りが生じてしまう 巡回セールスマン問題

類似した文書の検索・グラフの構築 ● 学習コーパス内の各文書をEmbeddingに変換 ○ 変換にはContriever model [Izacard+, 2022] を使用 ● すべての文書間のコサイン類似度を計算 ● 近似kNNにより、各文書に対する類似文書の集合を形成 ○ ○ product quantised inverted file（IVFPQ）FAISS indexという手法を使用 k=3の場合 ■ ■ ■ 文書1の類似文書={文書2、文書4、文書6} 文書2の類似文書={文書1、文書5、文書8} … ● 文書をノード、類似度をノード間のエッジの重みとしてグラフを構築 ○ kNNで検出された文書間でのみエッジをはるので、完全グラフとは限らない

類似度の合計が最大となるグラフ上のパスの探索 ● すべてのノードを一度だけ通り、かつ通ったリンクの重みの合計値が最大と なるパスを見つける ○ ○ 巡回セールスマン問題 NP hardなので解けない ● Greedyなアルゴリズムが近似的な解を与えることが知られている

Greedyなアルゴリズムによるパスの探索 ● 最も次数の低いノードからスタートし、重みが大きいエッジを通っていく ○ 通れるエッジがなくなったら、未到達ノードのうち最も次数の低いものから再度スタート ● 訪れた順に文書を並べ、context windowサイズで分割する ○ 割り切れない場合は文書を途中で区切る

アルゴリズム全体

学習データセットのクリーニング ● 文書間の類似度を計算する過程において、大部分が重複している文書が多 く見つかった ● 提案手法の学習はすべて重複を削除した上で実施 ○ 重複の削除の有効性についても検証（後述）

先行研究 ● kNNを用いて（重複を許して）サンプルを作成する方法 ○ ○ [Levine+ 2022]、[Guu+ 2020] 同じデータが複数回登場する可能性がある（巡回セールスマンではない）ため、一定 iteration回数におけるサンプルの多様性がランダムサンプリングに比べ低くなる ● 類似文書を同じバッチに含める方法 ○ ○ [Zhong+ 2022] サンプル間の文書の並びはランダムなので、類似度の低い文書が前後に並ぶ可能性があ る ● 関連する文書でのaugmentation ○ ○ ○ [Yasunaga+2022] 文書内にあるhyperlink/引用先の文章を挿入する手法 hyperlinkや引用のある文章は限られる

実験 ● LLaMAの実験設定を踏襲 ○ ○ 提案しているのはデータの前処理についてであり、残りの部分は既存のモデル・手法をその まま流用することが可能 学習データ：English Cmmoncrawl ■ 学習リソースの制限のため、ランダムに抽出した 235M文書を利用 ● 文脈の理解が要求される複数のタスクで効果を検証 ○ ○ ○ ○ 言語モデル In-Context learningでの文書分類 読解（Reading comprehension） Question Answering ● モデルサイズは0.3 ~ 7bで検証

比較手法 ● Standard ○ ○ コーパス内の文書をランダムに連結 LLMで通常用いられる手法 ● kNN ○ ○ コーパス内の各文書に対しその文書に類似する文書をkNNで抽出し、元の文書と連結する 1 epochする間に同じ文書が複数回登場する可能性がある ● In-Context Pretraining Language Models（ICLM） ○ 今回の提案手法

言語モデル ● 3つのデータセットにおいて、文書をランダムに並べて連結した上で評価 ● 実験したすべてのデータセットでPerplexityが低下 ○ kNNはStandardよりも悪い結果となった

In-Context Learningでの文書分類 ● few shotで例を提示し、文書分類させるタスク ○ 文脈を読み解く能力が要求される ● 実験した3つすべてのデータセットにおいて精度が向上

読解（Reading Comprehension） ● 与えられたパラグラフを元に、質問に答えるタスク ● 実験したすべてのデータセットで精度が大きく向上 ○ 特にHoppotQA ■ HotpotQAでは複数のドキュメントをまたいだ理解が要求される

Question Answering ● 2種類の設定で実験 ○ Closed-book方式：質問文が与えられ、LLMはその質問文だけをもとに回答する ■ ○ LLMのパラメータ内に蓄えられた知識のみで回答する必要がある Open-book方式：質問文に加えて、関連する外部の文章が追加の文脈として与えられる ■ より長い文脈を読み解く必要がある ● Openでは精度が向上するが、Closedでは既存と同じかそれ以下 ○ 論文の主張「ICLMで学習させたLLMは既存手法よりも記憶できる知識の量が少ないから （... because our model memorizes less.）」

学習時のLossの推移 ● 通常のランダムなサンプリングと比較すると、学習のLossの下がり幅が大き くなっている ○ 類似した文章が続いているので、学習時のtokenの予測が容易になっている可能性

学習データセットのクリーニング ● 重複文書を学習コーパスから取り除く（deduplication）ことにより、Perplexity が大きく低下する ○ In-Context Learningでは類似文書を隣り合うように配置するので、重複した文書が隣り合っ ているとnext token predictionがほぼコピーの学習になってしまう可能性がある（のでtest時 のpplが高くなっている可能性）

検索にかかる時間 ● 近似kNNを使って類似文書の構築（= グラフの構築）にかかる時間は、 32GPUで6時間 ○ GPUの詳細は不明 ● グラフの探索にかかる時間は20CPUで12時間 ○ CPUの詳細も不明 ● 一度探索が終われば（データセットが変わらない限り）再計算は必要ないの で、そこまでボトルネックにはならない？

まとめ・感想 ● LLMの学習データの構築方法の提案 ○ ○ 類似度の高い文章を連結させて学習データを作成すると、文脈の理解が要求されるタスクで 精度が大きく向上する あくまでも前処理の工夫なので、モデルアーキテクチャや最適化、プロンプティング等の他軸 の手法との組み合わせが容易 ● 感想 ○ 手法はシンプルだが精度は確実に上がっていそう ■ ○ （計算資源の問題は置いといて）試せるなら試してみたい データセットのクリーニングも結構効いてそう？ ■ Perplexity以外の指標での比較実験はされていなかった

参考文献 ● ● ● ● ● ● ● ● [Levine+, 2022] The inductive bias of in-context learning: Rethinking pretraining example design. ICLR. 2022. [McKenzie+, 2023] Inverse scaling: When bigger isn't better. 2023. [Liu+, 2023] Lost in the middle: How language models use long contexts. 2023. [Zhao+, 2021] Calibrate before use: Improving few-shot performance of language models. ICML. 2021. [Guu+, 2020] Retrieval augmented language model pre-training. ICML. 2020. [Zhong+, 2022] Training language models with memory augmentation. EMNLP. 2022. [Yasunaga+, 2022] LinkBERT: Pretraining language models with document links. ACL. 2022. [Izacard+, 2022] Unsupervised dense information retrieval with contrastive learning. TMLR. 2022.