PRX第四部分数据策略发布
速览
PRX发布第四部分博客,阐述其数据战略。该策略围绕数据治理、存储与共享展开,旨在支撑产品迭代。内容对技术团队有参考价值,但未涉及具体算法细节。
AI 深度解读
背景
本文是 Hugging Face 博客系列《PRX》的第四部分,聚焦于文本到图像生成模型 PRX 的数据策略。PRX 是一个 7B 参数规模的扩散模型,其训练数据的构建方式直接影响模型的视觉概念学习能力、生成质量以及后续微调空间。此前系列文章已经讨论了模型架构、训练配方等主题,本篇则详细拆解数据管道的设计原则、具体实现以及工程取舍。
核心内容
一句话概括:我们从公开和内部数据集的混合中组装训练数据,使用视觉语言模型(VLM)对图像重新生成描述,并将结果转化为可流式传输的语料库,用于训练 PRX。
高层级的数据管道流程如下(原文附有示意图,但文本无法展示,此处用文字描述其逻辑):
- 数据源聚合 → 2. 元数据与基线描述归一化 → 3. VLM 重新描述 → 4. 数据筛选与去重 → 5. 格式转换(Lance 构建 + MDS 流式传输) → 6. 训练时计算文本潜变量。
1. 指导原则
预训练需要多样化的数据集
目标是组装一个大规模、多样化的预训练数据集。在此阶段,模型学习的是世界的样子:视觉概念、物体和场景、事物如何构成和打光,以及图像可能包含的广泛内容。这是一个关于覆盖度和多样性的问题,而非单张图像的完美度。一个广泛、有代表性的语料库让模型学到的视觉世界结构远多于一个更小、更精致的语料库——即便许多单张图像只是普通的快照或略微压缩过。在此阶段过度过滤美学反而有害,会收窄分布,让模型失去后续无法恢复的概念和构图多样性。让生成结果看起来精致是另一个独立且后期的问题,我们通过在小而精心策划的数据集上进行微调和对齐偏好来解决。预训练追求广度,微调追求品味。
混合数据源
我们从公开和内部数据集的混合中组装预训练数据。此阶段的优先级是广度、多样性,以及利用已有的策展工作而非从头重做。如果某个来源已经经过质量过滤、去重、NSFW 和个人信息过滤,我们就直接基于这些工作,而不是大规模重复。不同来源以不同形态到达:有些自带图像数据,有些则只有元数据加基线描述,我们需要将其统一为共同格式。我们采取务实的方法:并非完全从头构建语料库,而是依赖现有数据集和自有工具快速组装。事后看,这不一定是最佳数据集,但对于预训练一个 7B 模型来说,它是一个坚实且轻量级的起点。
描述哲学
根据我们的经验,预训练中最重要的是使用长描述,准确描述图像中的所有内容。我们在第二部分直接看到了这一点:从短描述切换到长描述大幅提升了样本质量。如果描述是忠实的,我们就不用担心偶尔出现的截图、广告、标志或图像中的文字,因为这些内容也会在描述中被提及,模型将其视为有条件、可控的属性,而非无条件复制。准确的描述将“噪声”转化为你可以提示(prompt for)或提示避免(prompt away)的内容。这正是我们后续筛选故意保持轻量的原因:我们移除确实不可用的内容,而非一切不完美的东西。
数据格式
我们一直使用 Mosaic Streaming 和 Mosaic Data Shards(MDS)作为分布式训练的数据集格式。结合 Mosaic Composer,我们发现它是一个维护成本低、灵活且性能良好的分布式训练框架。此外,MDS 数据集可以轻松有效地混合和打乱,并支持直接从对象存储(如 S3 或 GCS)进行分布式训练。
然而,MDS 数据集非常僵化。添加一列或为某个筛选条件创建子集基本上意味着必须扫描并重写整个数据集。这就是我们在此类特征工程和数据集策展中使用 Lance 的原因。Lance 是一种列式数据格式,具有廉价的下推谓词、标量索引和向量搜索,是构建和探索数十亿行数据集的合适工具。
这两种格式在本帖以及 PRX 数据管道中相辅相成:Lance 用于构建,MDS 用于流式传输。
关于文本潜变量
在之前的训练中,我们使用 T5Gemma 作为文本编码器,并预计算文本潜变量,将其作为字节存储在 MDS 中。这次,在将文本编码器切换到 Qwen3-VL 后,我们决定在训练期间实时计算文本潜变量。在训练循环内运行文本编码器会降低吞吐量,但具体降低多少取决于模型:对于较小的去噪器来说影响可能显著,而在 PRX 7B 规模下,文本编码器的计算量相对于去噪器微不足道,我们测量到大约只有 3–4% 的吞吐量损失(在 30 天的训练中大约多出 1 天)。作为回报,我们得到两点好处:跳过预计算使 MDS 分片更小,小到足以将完整预训练数据集存储在 SLURM 集群的 SSD 支持的共享文件系统上,而无需通过网络从对象存储流式传输;同时,我们可以在以后自由更换文本编码器,而无需重写 TB 级的存储潜变量——这正是我们切换到 Qwen3-VL 时所做的更换。
关于图像编码
我们将所有图像编码为 JPEG(质量 92),而非 PNG 等无损格式。我们并未默认质量 92 是安全的,而是实际测量过。真实世界的图像通常已经过多次 JPEG 压缩,因此真正的问题是再一次重新编码是否有害。在高低分辨率(1–2 MP 和 0.25–0.5 MP)的真实图像上反复解码/编码循环中,第一次在质量 92 下重新编码几乎不可察觉。后续每次循环几乎不增加任何影响,因为 JPEG 会迅速收敛到稳定状态,即使经过 10 次循环,图像仍保持在不可察觉范围内,而 PNG 则会大 3–10 倍,却没有感官上的增益。以下是在质量 92 下对 100 张图像取平均,并以原始图像为基准测量的指标(PSNR 越高越好,LPIPS 越低越好):
(原文此处应有表格,但文本未给出具体数值,根据原文描述:第一次重新编码质量 92 与原始图像几乎无差异。)
由于大多数源图像已经是 JPEG 压缩的,将它们无损存储为 PNG 收益甚微,因此我们将所有图像转换为高质量 JPEG(质量 92)。
我们还检查了最关键的问题:在 JPEG 上训练是否会改变模型生成的结果。我们的语料库大部分本身就是 JPEG,因此 PNG 的唯一优势是避免引入新伪影。我们特意从高分辨率来源进行比较,预期即使是原始 JPEG 图像在高分辨率下也带有很少的伪影。我们训练了两个相同的 PRX 模型(1024px),使用相同的图像,一个以 PNG 存储,另一个以质量 92 的 JPEG 存储。两个模型在指标上训练效果相同,生成结果几乎无法区分,包括我们使用已知技术通过匹配量化表来估计输出看起来有多像 JPEG 压缩时:
(原文此处有检测率表格,同样未给出具体数值,描述为:检测率是指生成物中检测到任何量化结构的比例。平均和中位数估计 JPEG 质量仅针对那些被标记的图像。两个模型几乎无法区分:两者中大约每十个生成中只有一个显示出可检测的量化结构,且差异很小,因此我们有信心训练图像格式对输出质量影响极小。所以高质量 JPEG 存储不会在源数据已有的基础上增加任何可测量的影响。对于大规模文本到图像训练,这已经足够好了。对于用于训练 Photoroom 其他模型(如自定义 AI Shadows 模型)的伪影敏感数据,我们仍然依赖 PNG。)
2. 使用 Lance 构建数据集
你无法交互式地浏览包含数亿甚至数十亿行的 Parquet 表。因此我们将数据存储在 Lance 中,这样我们可以索引、查询和浏览。我们使用 Ray Data 进行数据摄取,并行读取源表并在集群上写入 Lance 表的多个片段。
一个值得分享的经验是关于碎片化。Lance 表被分割成多个片段,某些操作的规模与片段总数成正比,而非行数:一次扫描会打开每个片段的文件……(原文至此截断,未完整结束。后续内容可能涉及碎片化对性能的影响及优化策略,但根据已给文本,我们只能呈现到此。)
关键要点
- 预训练追求广度而非单张完美
