1. 项目背景与核心问题
在语言模型预训练领域，数据选择的质量直接影响模型性能。传统方法通常采用现成的通用嵌入模型（如基于检索任务训练的模型）计算样本相似度，但这种"拿来主义"存在三个关键缺陷：
- 任务失配：检索模型优化的相似度与预训练目标存在偏差
- 领域偏移：通用嵌入难以捕捉特定预训练语料的分布特性
- 评估缺失：缺乏针对预训练场景的相似度度量评估框架
我们团队在准备NeurIPS 2025投稿时发现，当使用The Pile数据集预训练1.7B参数模型时，采用现成嵌入模型选出的数据，其下游任务表现比随机选择仅提升2-3%，这促使我们重新思考相似度度量的本质要求。
2. 相似度评估框架设计
2.1 预训练损失泛化一致性
构建样本对(x_i, x_j)，通过以下步骤验证相似度度量的有效性：
- 在子集D_sub上预训练基础模型M_base
- 计算将x_j加入训练时M_base(x_i)的损失变化ΔL
- 检验ΔL与嵌入距离d(f(x_i),f(x_j))的Spearman相关性
实验显示，基于预训练任务微调的嵌入模型相关性达0.61，而通用嵌入模型仅0.28。
2.2 下游任务性能验证
设计数据选择-预训练-微调的三阶段验证流程：
python复制def evaluate_metric(embedding_model):
selected_data = diversity_sampling(train_data, embedding_model)
pretrained_model = train_lm(selected_data)
return downstream_task_score(pretrained_model)
在200B token规模下，专用嵌入选出的数据使GLUE平均得分提升19.2%，远超通用嵌入的6.8%。
2.3 数据源区分能力
构建混淆矩阵评估嵌入模型区分不同来源数据的能力：
| 数据源 | arXiv | GitHub | StackExchange |
|---|---|---|---|
| 通用嵌入准确率 | 62% | 58% | 55% |
| 专用嵌入准确率 | 89% | 83% | 91% |
3. 关键技术实现细节
3.1 嵌入模型架构优化
采用双塔结构，在预训练过程中动态更新：
- 共享编码器：12层Transformer，768隐藏维度
- 对比损失：InfoNCE温度参数τ=0.05
- 训练策略：每50B tokens更新一次嵌入模型
3.2 多样性采样算法改进
传统k-center贪婪算法的时间复杂度为O(n²)，我们提出：
- 基于局部敏感哈希(LSH)的近似算法
- 分桶采样：将200B token分为400个桶并行处理
- 记忆库机制：保留最近10B tokens的嵌入避免重复计算
这使得处理速度提升17倍，内存占用减少83%。
4. 实战经验与避坑指南
4.1 数据选择中的典型误区
-
陷阱1：过度依赖余弦相似度
- 实际发现：预训练早期阶段欧式距离更有效
- 解决方案：动态调整距离度量（前50B tokens用L2，后续用余弦）
-
陷阱2：忽略数据量级效应
- 关键发现：<100B tokens时多样性优先，>100B时质量优先
- 调整策略：分阶段调整采样权重系数
4.2 工程实现优化技巧
-
嵌入缓存策略：
- 使用FAISS构建动态索引
- 采用LRU缓存最近1M样本嵌入
-
混合精度训练：
bash复制
torch.cuda.amp.autocast(enabled=True) optimizer.step(scaler.scale(loss).backward) -
内存优化：
- 梯度检查点：减少30%显存占用
- 分片优化器：将参数分组更新
5. 扩展应用与未来方向
当前框架已成功应用于：
- 多模态预训练数据选择（图像-文本对齐）
- 增量学习场景下的持续数据筛选
- 低资源语言的预训练语料构建
我们在内部实验中还发现：
- 相似度度量与模型架构存在耦合效应
- 动态嵌入更新频率影响最终性能
- 领域自适应采样可提升专业领域表现
重要提示：当处理超大规模数据时（>1T tokens），建议采用分层采样策略，先按粗粒度分类再执行细粒度选择，可降低90%计算开销。