[论文解读] 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

一句话总结

本文提出 DSIR，一种可扩展框架，通过在降维特征空间（哈希 n-gram）中估计重要性权重并据此重采样以匹配目标分布，从而在下游性能上优于随机或启发式基线。

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 在领域特定的连续预训练和通用领域语言模型预训练中的有效性
• 引入 KL reduction 作为将特征空间保真度与下游性能联系起来的度量

提出的方法

• 通过 h 将原始数据和目标数据映射到特征空间，创建 q_feat 和 p_feat
• 对每个原始样本 z_i = h(x_i) 估计重要性权重 w_i = p_feat(z_i) / q_feat(z_i)
• 根据归一化权重无放回地从原始数据中重采样 k 个样本
• 使用哈希 n-gram 特征实例化 DSIR（m 桶，例如 m=10000）和袋-的 n-gram 模型参数 gamma
• 通过计数哈希桶频率来估计特征分布，得到 gamma_hat
• 通过在 1 个 CPU 节点用 4.5 小时从 The Pile 中选择 1 亿文档来展示可扩展性；以 KL reduction 作为对目标相似性的代理进行衡量

实验结果

研究问题

• RQ1 在降维特征空间进行重要性重采样是否能有效选择与目标分布匹配的预训练数据？
• RQ2 哈希 n-gram 特征是否能保留选择与下游目标相关的数据所需的信息？
• RQ3 相较于随机选择和启发式筛选，DSIR 在领域特定的继续预训练和通用领域预训练中是否提升下游性能？
• RQ4 KL reduction 是否是跨数据选择方法对下游性能的可靠廉价代理？
• RQ5 当为一个目标选择数据、在另一个目标进行评估时，目标领域对迁移有何影响？

主要发现

• DSIR 在领域特定的连续预训练中，相对于随机选择和启发式基线平均提升约 2%（8 个数据集的 F1），并且与人工筛选相当。
• 在通用领域预训练中，DSIR 相对于基线在 GLUE 上的表现平均提高 2–2.5%。
• 哈希 n-gram 特征使得能够在 4.5 小时内从 The Pile 中选择 1 亿文档，绝大多数时间花在计算重要性权重上。
• 对哈希 n-gram 的 KL reduction 与下游性能高度相关（Pearson r = 0.82）。
• 判别式 DSIR 变体的表现不及生成式 DSIR； unigram+bigram（哈希）特征优于仅 unigram 的版本。
• DSIR 使自动数据选择能够在许多设置中替代人工筛选；质量过滤器和计算预算控制结果。
• 同域预训练比跨域预训练在使用 DSIR 选择数据时迁移效果更好（平均提升约 1.7%）。