[論文レビュー] Data Selection for Language Models via Importance Resampling

Sang Michael Xie, Shibani Santurkar+2 more|arXiv (Cornell University)|Feb 6, 2023

Natural Language Processing Techniques被引用数 16

ひとこと要約

本論文は、ハッシュ化されたn-gramの特徴空間で重要度重みを推定し、それに基づいて再サンプリングすることで前訓練データを選択し、ターゲット分布に合わせるスケーラブルなフレームワークであるDSIRを紹介します。乱数選択やヒューリスティックベースラインよりダウンストリーム性能を向上させます。

ABSTRACT

Selecting a suitable pretraining dataset is crucial for both general-domain (e.g., GPT-3) and domain-specific (e.g., Codex) language models (LMs). We formalize this problem as selecting a subset of a large raw unlabeled dataset to match a desired target distribution given unlabeled target samples. Due to the scale and dimensionality of the raw text data, existing methods use simple heuristics or require human experts to manually curate data. Instead, we extend the classic importance resampling approach used in low-dimensions for LM data selection. We propose Data Selection with Importance Resampling (DSIR), an efficient and scalable framework that estimates importance weights in a reduced feature space for tractability and selects data with importance resampling according to these weights. We instantiate the DSIR framework with hashed n-gram features for efficiency, enabling the selection of 100M documents from the full Pile dataset in 4.5 hours. To measure whether hashed n-gram features preserve the aspects of the data that are relevant to the target, we define KL reduction, a data metric that measures the proximity between the selected pretraining data and the target on some feature space. Across 8 data selection methods (including expert selection), KL reduction on hashed n-gram features highly correlates with average downstream accuracy (r=0.82). When selecting data for continued pretraining on a specific domain, DSIR performs comparably to expert curation across 8 target distributions. When pretraining general-domain models (target is Wikipedia and books), DSIR improves over random selection and heuristic filtering baselines by 2-2.5% on the GLUE benchmark. Code is available at https://github.com/p-lambda/dsir.

研究の動機と目的

• データ選択を、ターゲット分布に一致させるために大規模なラベルなしデータセットのサブセットを選ぶこととして形式化する
• 計算可能な特徴空間で重要度重みを推定するスケーラブルな方法を開発する
• ハッシュ化されたn-gram特徴がデータ選択にとってターゲット関連情報を保持することを示す
• DSIRが領域特化の継続前訓練と一般領域のLM前訓練で有効であることを示す
• 特徴空間の忠実度とダウンストリーム性能を結ぶ指標としてKL削減を導入する

提案手法

• 生データとターゲットデータを h によって特徴空間に写し、q_feat と p_feat を作成する
• 各生データ z_i = h(x_i) に対して重要度重み w_i = p_feat(z_i) / q_feat(z_i) を推定する
• 正規化された重みに従い、生データから置換なしで k 個の例を再サンプリングする
• ハッシュ化されたn-gram特徴（m 個バケツ、例: m=10000）と bag-of-ngrams モデルパラメータ gamma を用いてDSIRを実装する
• ハッシュ化されたバケツ頻度を数えて特徴分布を推定し gamma_hat を得る
• 1 CPUノードでThe Pileから100M文書を4.5時間で選択することでスケーラビリティを示し、KL削減をターゲット類似性の代理指標として測定する

実験結果

リサーチクエスチョン

• RQ1 縮小された特徴空間での重要度再サンプリングは、ターゲット分布に一致する前訓練データを効果的に選択できるか？
• RQ2 ハッシュ化されたn-gram特徴は、ダウンストリームターゲットに関連するデータを選択するのに必要な情報を保持するか？
• RQ3 領域特化の継続前訓練と一般領域の前訓練において、DSIRは乱択やヒューリスティックフィルタリングと比べてダウンストリーム性能を改善するか？
• RQ4 KL削減はデータ選択手法間でダウンストリーム性能の信頼できる低コスト代理指標か？
• RQ5 あるターゲット用に前訓練データを選択し、別のターゲットで評価した場合、ターゲットドメインの整合性は転移にどう影響するか？

主な発見

• DSIRは領域特化の継続前訓練において乱択やヒューリスティックなベースラインより約2%向上し（8データセットのF1の平均）、手動選定と同等である。
• 一般領域の前訓練では、DSIRはベースラインより平均で2–2.5%GLUE性能を向上させる。
• ハッシュ化されたn-gram特徴はThe Pileから100M文書を4.5時間で選択可能とし、ほとんどの時間は重要度重みの計算に費やされる。
• KL削減はダウンストリーム性能と強い相関を示す（Pearson r = 0.82）。
• 識別型DSIRの派生は生成型DSIRを下回る； unigram+bigram（ hashed）特徴は unigram のみよりも優れている。
• DSIRは多くの設定で manual curation を代替できる自動データ選択を可能にする；品質フィルターと計算予算が結果を制御する。
• DSIRでデータを選択した場合、同一ドメイン内の前訓練は異ドメイン間より転移が良くなる（平均約1.7％の改善）。