Krea 2 技术报告

在这份技术报告中，我们介绍 Krea 2：一系列旨在兼顾广泛审美多样性与用户创意控制的基础模型。我们将阐述：(1) 我们的数据整理原则；(2) 模型架构；(3) 多阶段训练流程；(4) 分布式基础设施；以及 (5) 未来工作。

引言

过去几年间，图像生成取得了显著进展。扩散模型与流匹配模型能够生成高分辨率图像、呈现锐利的照片级真实感与稳定结构、渲染密集文字、编码广泛的世界知识，并精确细致地遵循用户提示词。这些进步得益于多个相互作用的因素，包括可扩展的 Transformer 架构、改进的说明文字生成与文本编码器、更好的潜在表示，以及流水线化的后训练技术。然而，随着该领域在优化这些能力可靠性方面的推进，许多系统已收敛到一套狭隘的默认美学风格。尽管作为高效的生产工具，这使它们作为创意探索引擎的效果大打折扣——在创意探索中，用户常常需要跨风格、情绪、构图和视觉方向进行搜索，而非接收一个单一的、精雕细琢的默认输出。

为解决这些局限，我们推出 Krea 2，一系列专注于创意探索的基础模型。Krea 2 的模型建立在这样的信念之上：图像生成应成为一种探索性媒介——其表现力足以覆盖多种美学风格，其可控性足以让创作者自如驾驭。

我们从头构建了大规模数据基础设施和分布式训练框架，以整理一个涵盖广泛世界知识和风格覆盖的综合性预训练数据集。

利用这一基础设施，我们通过一个包含预训练、中期训练、监督微调（SFT）、偏好优化和强化学习（RL）的多阶段流水线训练高表现力模型，每个阶段都旨在逐步完善模型的输出分布。通过彻底的消融实验，我们开发了一个简单而高效的扩散Transformer（DiT）架构。我们的模型整合了多个加速收敛的组件，包括iREPA、改进的VAE和Qwen3-VL。我们还集成了多项架构改进，包括分组查询注意力（GQA）、sigmoid门控注意力、轻量级时间步调制以及针对文本编码器特征的多层特征聚合，这些共同提升了训练稳定性和效率。

一个强大的基座模型只有用户能够可靠地访问其分布中他们关心的部分时才有用。在训练中，模型从丰富且精心构建的标注中学习，这些标注用密集的视觉细节描述图像。在实践中，用户输入往往更简短、更模糊，并且受到多种不同表达习惯的影响。有些用户用自然语言描述场景；另一些用户则暗示一种情绪、一种风格或一张参考图像。这就在模型学习到的条件空间与推理时创意意图的表达方式之间造成了差距。

为了缩小这一差距，我们构建了两个系统，让 Krea 2 更具探索性，并且能同时通过文本和图像输入进行操控：一个提示词扩展器和一个风格参考系统。提示词扩展器将简单或未完全明确的用户提示映射为更丰富的视觉方向，同时不会覆盖用户的意图。它通过一个两阶段的 SFT 和 RL 流水线在开源大语言模型之上训练而成，其目标不仅是提升图像质量，还要鼓励创造性变化和可控的探索。作为这一文本界面的补充，风格参考系统让用户在文字不足以表达时，可以通过图像来表达视觉意图。它允许用户注入一张或多张参考图像的风格或氛围，同时尽可能减少内容泄漏，并提供对风格强度和加权风格混合的精细控制。

这些组件共同将 Krea 2 定义为一个面向探索式生成的基础模型。Krea 2 并非只针对某种单一的精美默认值进行优化，而是旨在呈现一个广阔的视觉空间，并为用户提供通过文本和图像控制来探索这一空间的实用方式。Krea 2 在 Artificial Analysis 文生图排行榜上位列前十，并在独立实验室的模型中排名第二。Krea 2 提供了一个全面的基线，在保持竞争性能的同时，实现了创意生成体验。

数据

数据整理原则

在详细介绍我们的数据流水线之前，有必要先明确：什么样的数据混合对我们的目标来说才是好的。好的混合不仅仅包含“高质量”图像。鉴于我们的目标是构建一个富有表现力、风格多样的模型，多样性和广泛的领域覆盖至关重要。我们认为，传统基于模型的过滤方法（使用美学评分和图像质量评估（IQA）模型）会引入隐含的偏见。例如，这类方法可能将模糊图像归类为低质量，即便运动模糊或柔和感可能是刻意的艺术选择。

此外，我们认为，只要一段描述文字能准确描述其对应的图像，即使是负面的图像也可能对下游应用场景有所帮助：因为模型精确地理解了这种不良行为，这样的样本随后可以用于引导生成结果远离该分布。

出于这些原因，我们通过仅过滤掉以下内容来构建预训练数据集：

1. 重复样本以及被过度呈现的概念。
2. 视觉语言模型始终未能捕捉到图像重要方面的样本。
3. 会引发不良偏差和伪影的样本。
4. 视觉复杂度过高、难以在低分辨率下可靠建模的样本。
5. AI 生成的样本

这些条件塑造了一个覆盖面广泛的预训练数据集，同时避免了较差的文本到图像对齐以及伪影问题。

重要的是，我们在预训练混合数据中完全没有使用 AI 生成的图像。合成数据和知识蒸馏可以是获取模型能力的有效捷径。然而我们发现，即使是小比例的 AI 生成图像也会给模型的输出分布引入偏差，因为合成图像往往更容易学习，这实际上对模型质量施加了一个上限。因此，我们设计了内部分类器来过滤掉这类图像。

描述文字生成

我们采用多阶段方法来生成描述文字。首先，我们在每张目标图像上运行一个 OCR 模型，以提取任何可见的文本。在第二阶段，我们将 OCR 结果以及任何可用的元数据（如相机设置、已知实体等）提供给描述生成模型，该模型会生成一条经过丰富处理的描述文字，其中融合了世界知识与提取出的文本。

通用描述生成流程

一旦获得了包含丰富上下文的长篇幅自然语言描述，我们会使用一个成本较低的大语言模型将其重新格式化为多种长度和格式，让模型接触到不同类型的提示词风格。根据经验，我们发现基于长提示词进行训练能够提供密集的监督信号，从而加快收敛速度并降低训练损失。然而，对于许多下游和实际应用场景而言，模型在短篇幅和中篇幅提示词上的表现仍然很重要。因此，我们主要基于长描述进行训练，同时确保在整个训练过程中模型也能接触到短篇幅和中篇幅的提示词。

我们的整体训练流程与数据阶段

预训练数据

预训练数据涵盖256像素、512像素和1024像素分辨率三个训练阶段。逐步提高分辨率形成了一种课程学习策略：我们将大部分FLOPs分配给低分辨率阶段，以高效地构建模型的核心能力，然后随着训练分辨率的提升，使模型具备高保真生成能力。

低分辨率预训练阶段是学习基础图文对齐与结构组成的阶段。在此阶段，数据集的规模达到数十亿张图像，因此我们大量依赖基于CPU的低成本过滤器来剔除低质量图像。这些过滤器包括简单的坏文件过滤、分辨率过滤和宽高比过滤（用于剔除不合格图像），以及拉普拉斯滤波器（用于剔除具有极端纹理和噪声模式的图像）。

举个例子，我们在预训练K2时遇到的一个问题是模型倾向于生成纯色背景和边框伪影。为了解决这一问题，我们使用了RGB熵、黑白像素比例、自定义启发式规则以及内部分类器，来过滤掉导致这种行为的样本。

在构建内部分类器时，一个有效的策略是使用大型视觉语言模型为过滤任务编写针对性的系统提示词（例如，检测某种特定模式或伪影），生成一个伪标签数据集，然后训练一个基于DINOv3或SigLIP-2的小型分类器，以便大规模运行该过滤器。在低分辨率阶段，任何需要GPU计算资源的过滤模型都被控制在10亿参数以下，以确保效率。

在低分辨率阶段的去重中，我们主要使用成本较低的哈希方法，结合 md5、phash 和 colorhash，以最小的计算量去除重复图像。我们发现默认的 8x8 phash 没有考虑颜色信息，且误报率较高；因此，我们将 12x12 phash 与 colorhash 结合，以实现更鲁棒的去重。

随着训练分辨率提升，我们引入了图像质量和审美过滤器。重要的是，这些质量分数仅用于剔除质量极差的图像，而非根据分数对图像进行过采样。此外，我们还使用图像复杂度分数和文本密度（来自 OCR 结果），以排除那些文本和内容在低分辨率下无法有效表示的图像。我们会在训练过程中动态调整质量、复杂度和文本密度的阈值。

除了传统的质量过滤器外，我们还基于预训练语料的一个样本，在 SigLIP-2 嵌入向量上训练了一个稀疏自编码器（SAE）。训练完成后，我们使用视觉语言模型（VLM）根据每个 SAE 特征的前 k 个激活样本来对其进行标注。这些标注后的特征形成了一套无监督标签系统，我们可以从中提取每张图像的主要 SAE 特征。这套标签系统在无需训练显式分类器的情况下，有效帮助过滤明显的视觉伪影。

中训练数据

与预训练阶段不同，中训练阶段会显式选择已知能提供良好风格覆盖范围和特定视觉领域高质量图像的来源。预训练是一个从通用数据池出发的自底向上过程，而中训练数据则是自顶向下策划的：先确定领域和来源。中训练是一个关键阶段，它平滑地衔接了通用预训练分布与高质量 SFT 分布。为了提升分布质量，我们引入了语义聚类，并使用基于检索的策略来确保世界知识覆盖。

在《自监督学习数据自动策展》方法的基础上，我们使用 FAISS 进行层次 K 均值聚类，然后对聚类结果进行采样，以保留长尾视觉概念，同时避免在头部概念上浪费算力进行过采样。计算完层次聚类后，我们让一个视觉语言模型（VLM）检视每个聚类质心附近的图像，以命名聚类并在适当时对聚类进行标记。经人工复核标记的聚类后，我们丢弃了若干质量较低或有问题的聚类。随后通过语义去重移除更多冗余数据——对每个剩余叶子聚类内的图像计算 SigLIP 相似度。

图像生成模型的一项重要能力是忠实还原用户仅凭名称即可引用的已知实体。部分实体（如运动员或演员）可能落入包含众多其他实体的语义聚类中，这使其在直接的层次采样下存在被丢弃的风险。为解决此问题，我们使用 Danker 对英文维基百科运行 PageRank，并按排名保留排名前 90% 的文章。随后，我们根据 Wikidata 元数据过滤掉所有描述不可表征主题的文章，并对剩余的约 500 万个概念，在我们的数据集中对所有字幕进行全文搜索以评估覆盖率。在采样时，我们优先选择字幕提及稀有概念的图像。最后，我们对所得样本重复此覆盖率分析，以确认初始数据集中存在的概念无一被完全丢弃。

监督微调数据

在监督微调（SFT）阶段，我们使用一个聚焦于单个视觉领域的小型人工策展数据集。我们发现，一旦达到足够的数据量，数据集的质量远比其规模重要。

架构

在进行架构消融实验时，我们发现有帮助的做法是将每次消融的目标分为以下类别之一：

• 稳定性：是否使训练更稳定？是否减少损失和梯度尖峰？
• 性能：是否使模型收敛更快？如果是，该趋势在更长的时间跨度和更高分辨率下是否依然保持？
• 效率：能否在不影响模型质量的前提下，减少参数量、FLOPs、内存或通信需求？
• 简洁性：能否在不影响其他类别的情况下，让模型变得更简单？

值得注意的是，我们的许多架构决策都受到大语言模型领域采用情况的影响。选择在大语言模型生态中已成熟的架构，使我们能够利用现有的内核和优化，即使对于扩散模型也是如此。

基于这些目标，我们从以下基线开始。

Transformer块

我们首先将GeLU MLP替换为扩展因子为4倍的SwiGLU层，这已成为大语言模型架构中的事实标准模块。引入SwiGLU带来了持续的性能提升，因此我们在后续所有消融实验中均采用了它。

修订完MLP设计后，我们考虑了GQA、MLA和门控Sigmoid注意力作为多头注意力基线的替代方案。我们发现GQA在带来极小性能损失的同时，提供了更高的计算效率。我们还探索了MLA，并观察到相比GQA有轻微提升，但由于它引入了额外的计算开销，我们并未采用。我们在MLA中使用了上下投影进行KV压缩，并且没有使用解耦RoPE，因为扩散模型是纯预填充过程，推理时不使用KV缓存。

在分组查询注意力（GQA）的基础上，我们增加了门控 Sigmoid 注意力机制，其方法借鉴自《面向大语言模型的门控注意力》。门控 Sigmoid 注意力机制仅增加极少的计算量和参数量。虽然它并未带来显著的性能提升，但使训练动态更加稳定，这一点在训练过程中的损失曲线和梯度范数曲线上均有所体现。

我们还对模态流设计进行了消融实验：

1. 单流设计：一种标准的 Transformer 模块，其中注意力层和 MLP 层的权重在文本 token 和图像 token 之间共享。
2. 双流设计：联合注意力机制，但文本 token 和图像 token 各自拥有独立的注意力层和 MLP 层权重。
3. 混合流设计：将前两者结合，网络前三分之一采用双流模块，剩余三分之二采用单流模块。

我们发现这三种设计在性能上没有显著差异，不过混合流设计略优于其他两种。然而，为了简洁起见，我们在最终架构中采用了单流模块。

时间步条件化

许多 MMDiT 架构使用逐模块的 MLP 来生成缩放、平移和门控因子。这些 MLP 模块可能占到总参数量的 20%–30%，我们认为对于注入一个标量条件而言，这样的参数量过度了。因此，我们将逐模块的 MLP 替换为逐模块的可调偏置项。这一改动使我们能够在不牺牲模型性能的前提下，将更多参数分配给注意力层和 MLP 层。

超越AdaLN调制，我们探索了两种替代方案：(1) 完全移除时间步条件化，以及(2) 通过时间步token进行上下文内时间步条件化。在低分辨率预训练实验中，完全移除时间步信息始终不如AdaLN基线。对于上下文内条件化，我们使用正弦嵌入创建时间嵌入，将它们拼接成统一的文本+图像+时间序列，并完全移除AdaLN层。在256像素预训练中，4-16个时间步token足以替代AdaLN。然而，在512像素和1024像素下，上下文内条件化相对于AdaLN基线表现较差。我们尝试通过增加时间步token数量来缓解这一问题，但观察到收益递减，且无法在更高分辨率下达到有竞争力的性能。

位置编码

在消融实验中，我们实现了多种RoPE方案。我们使用3D轴向RoPE，头部维度分别对应帧、高度和宽度。对于文本token，我们将RoPE索引设为零。在低分辨率下，我们没有观察到切换到Golden Gate RoPE、MRoPE、归一化RoPE或部分RoPE带来显著提升。对于部分RoPE，我们仅旋转头部维度的前半部分，剩余部分不旋转。正如预期，部分RoPE在将模型从256像素扩展到512像素时产生了更好的零样本推理结果，且未出现常见的重复伪影。尽管有这种初始分辨率泛化能力，但随着高分辨率训练的继续，部分RoPE最终表现不如基线RoPE设置。

自动编码器

近期研究表明，自编码器的潜空间设计能显著加速图像生成模型的训练。我们以FLUX.1-dev自编码器为基线，并将其与通义千问图像VAE、DC-AE、FLUX 2 VAE及我们内部自编码器进行对比。我们首先测试了DC-AE系列，因为它支持高达32倍的空间压缩，这对训练和推理效率均大有裨益。然而我们发现，由于其重构误差，DC-AE对扩散模型解析精细细节的能力施加了硬性上限。

相比之下，通义千问图像VAE和FLUX 2 VAE提供的潜空间在预训练消融实验中收敛速度显著更快，同时保持了优异的重构质量。因此，我们最初使用通义千问自编码器来扩展早期模型，后来在更大规模模型上采用了FLUX 2 VAE。我们还短暂探索了利用DINOv3进行语义对齐并结合轻量扩散损失训练内部自编码器的方法，该思路与REPA-E类似。我们验证了其性能与通义千问图像自编码器不相上下，但受时间限制，我们最终选择了已在规模化验证中表现良好的通义千问图像VAE和FLUX 2 VAE。

残差设计

我们默认使用标准残差连接。我们曾短暂尝试过Laurel，它通过添加低秩瓶颈分支来增强残差连接的表达能力，但未观察到明显改进。对于未来模型，我们计划探索NOBLE、delta注意力残差和mHC等替代方案，以改进扩散Transformer的残差设计。

归一化

RMSNorm 已成为大语言模型架构的标准组件，但尚未完全融入近期的主流扩散 Transformer 架构中。我们从 LayerNorm 基线出发，将所有归一化层替换为 RMSNorm，观察到质量下降非常微小。因此，我们将 RMSNorm 作为默认归一化模块（例如用于 prenorm 和 QKNorm）。我们使用零中心化的 RMSNorm，并对其可学习参数施加权重衰减。我们还尝试了更高效的变体（如 Derf），但发现存在不可忽视的质量下降。

文本编码器

我们以 T5-XXL 作为基线文本编码器。从一开始，我们就有意保持架构简洁，仅使用单个文本编码器。值得注意的是，我们发现 T5-XXL 相对于 T5Gemma、umT5、Qwen 2.5 VL 和 Qwen 3 VL 仍然极具竞争力。最终，我们选用 Qwen 3 VL 作为最终文本编码器，因为视觉语言模型（VLM）提供了更丰富的输入空间（文本和图像）以及更强的多语言泛化能力。

此外，受 Unifusion 启发，我们并未直接采用 VLM 特征的最后隐藏层，而是引入了一个浅层注意力层，用于跨层聚合隐藏特征。这种设计使模型能够动态选择从粗到细的文本表征。自回归大语言模型的最后层特征对我们的目标而言并非最优，因为它们是为下一 token 预测而非图像生成而优化的。在这层跨层特征聚合的基础上，我们还沿 token 轴添加了轻量级双向 Transformer 层，以降低表征空间中的自回归偏置。

优化

我们在整个流程中使用 AdamW 作为主要优化器。最初，将 Muon 应用于 MMDiT 架构时，我们观察到结果参差不齐。默认情况下，我们采用 Dion 中的 Muon 实现，并结合 Moonlight 中的 RMS-matched 设置来迁移 AdamW 超参数。

在我们的探索中，Muon 在初始步骤中比 AdamW 收敛得更快，但在更长的周期内表现不如 AdamW。我们还遇到了 Muon 的一些稳定性问题，包括训练过程中频繁的损失和梯度范数尖峰。我们发现必须将 MMDiT 的第一层和最后一层线性层从 Muon 参数中排除；这与大语言模型文献中的做法一致，即嵌入层和 LM 头参数也被排除在 Muon 之外。排除这些层并添加 Nesterov 动量后，Muon 在低分辨率和高分辨率下均持续优于 AdamW 基线。由于时间限制，我们在最近的预训练运行中并未采用 Muon，但鉴于这些强劲的结果，我们计划在下一轮预训练中采用它。

训练

我们的训练流程采用受现代大语言模型训练流程启发的多阶段结构。

预训练

预训练建立模型的基础能力，包括文本-图像对齐、文本渲染、风格覆盖和结构一致性。我们逐步将分辨率从 256px 提升到 512px，再到 1024px。对于最终模型，我们使用 v-参数化下的标准整流流损失进行训练。为了加速早期阶段，我们在 256px 阶段的第一个 epoch 中使用 iREPA，然后移除它，这促使 MMDiT 学习自身的表征，同时显著加快初始收敛。我们还探索了其他加速策略，如 TREAD，但收益甚微。

在 256px 和 512px 阶段，我们使用 8 位训练，相较于 bf16 基线观察到 15%-20% 的训练速度提升，训练损失和评估指标衰减极小。在 256px 阶段，我们使用基于张量级缩放的 8 位训练；在 512px 阶段，我们使用更细粒度的行级缩放。从 1024px 开始，直到最后的强化学习阶段，我们使用标准的 bf16 训练。

高分辨率预训练的另一个重要方面是调整与分辨率相关的时间偏移调度。我们在训练和推理中都使用移动的对数正态采样调度，并随着分辨率的提高逐步增加偏移量。参照 FLUX 2 VAE 博客，我们在每个分辨率下搜索最优的训练时间偏移。我们只搜索训练时的偏移量，同时保持推理时的偏移调度不变，因为某些自编码器对推理时间偏移不那么敏感。

在预训练期间，我们采用热身‑稳定‑衰减的学习率调度，并按照《大语言模型预训练中的模型合并》一文应用 PMA 方法。我们验证了 PMA 在达到与 EMA 相当的性能的同时，避免了 EMA 显著的内存开销。我们没有观察到合并方法之间的显著差异，尽管调整合并检查点的数量和合并间隔可以在下游指标上带来微小的提升。

中间训练

中间训练在大语言模型文献中已经变得常见，我们将类似阶段引入到我们的流程中。其重点是在监督微调阶段之前预热模型的分布。我们发现，中间训练通常是流程中最后一个可以为模型配备下游能力的阶段，这些能力包括高保真、高分辨率生成、强大的领域覆盖以及文本渲染。

监督微调（SFT）

在监督微调阶段，我们整理了一个小型、专门的高美感图像集。目标是进一步引导模型朝着美学上更理想的方向发展。我们发现这一阶段特别有助于提升整体检查点质量，并解决早期检查点中普遍存在的高饱和度和纹理问题。

在训练了特定领域的 SFT 检查点之后，我们使用模型合并来生成一个通用型 SFT 检查点。在流程的后期阶段，模型合并的收益逐渐递减，因为不同检查点的改进方向开始相互冲突。

偏好优化（PO）

偏好优化（PO）是我们在后训练堆栈中的第一阶段，由两阶段流水线构成。在第一阶段，我们运行一个大规模合成偏好对生成流水线，用于初步优化，采用类似于差分学习的策略；我们确保大多数偏好对至少包含一个来自当前策略（on-policy）的样本。第二阶段是一个校准阶段，仅使用人工标注。这些标注完全由熟悉该模型具体优势、劣势和特性的内部人员进行采集。

偏好优化中的一个常见现象是策略偏离。从高层次来看，诸如DPO等偏好优化方法鼓励模型增大其生成偏好样本的似然与生成非偏好样本的似然之间的差距（相对于参考模型）。在实践中，观察不同偏好数据集混合的情况，我们发现模型通过降低生成两种样本的似然（但下降速度不同）来实现这一目标。如果当前模型分布中胜出样本和失败样本的质量都较低，这种做法是理想的；但这一假设并不总是成立，具体取决于偏好集是如何整理的。此外，这种偏离会使模型远离通用的预训练分布，在训练后期表现为高频伪影。为了缓解这一问题，我们设计了一种DPO的变体，称之为STPO，它增加了一个辅助损失，并对原始DPO公式进行了修改，以减少这种偏离。

强化学习（RL）

强化学习（RL）是训练流水线的最后阶段。我们采用一种多奖励的GRPO风格方法，使用多个奖励模型：（1）通用美学模型，（2）指令遵循奖励模型，（3）文本渲染奖励模型，（4）伪影与结构奖励模型。通用美学模型是通过在偏好优化阶段收集的偏好数据上微调一个开源视觉语言模型（VLM）得到的。我们精心设计了奖励结构并调整了数据混合比例，以防止奖励破解引入的伪影。

与本质主观的一般美学奖励不同，提示词遵循和文本渲染能提供更具体的信号，因为它们可以对照用户明确的意图进行核查。挑战在于，这种意图在不同提示词之间差异很大。为应对这一点，我们采用了一种受大语言模型训练中基于评分标准评估启发的、针对特定提示词的评分标准奖励。我们不要求评判模型给出单一的整体评分，而是将每个提示词拆解为可验证的要求，并据此评估生成的图像。这为RL阶段提供了更结构化的信号，用于与用户意图对齐，使模型能更好地满足细粒度提示词约束，同时不会将提示词遵循降级为通用图像质量。

我们还发现，仅优化美学和提示词遵循可能导致奖励劫持。模型可能学会生成乍看合理、但包含结构性伪影的图像，例如多出的手指、畸形的肢体或扭曲的文本。这些缺陷对人类来说显而易见，但往往会被通用VLM评判模型遗漏。为解决这一问题，我们训练了一个专门的伪影奖励模型，用于检测这些结构性错误，并阻止RL阶段为提升基准测试指标而牺牲视觉正确性。

在强化学习阶段，我们发现成功不仅取决于奖励模型的质量，还取决于训练计算量在提示词上的分配效率。奖励模型定义了改进方向，而提示词池则决定了模型在何处获得有用的学习信号。因此，我们精挑细选了一个涵盖多样风格、概念、场景和主题的广泛提示词池，然后持续分析生成组的奖励统计信息，以识别哪些提示词信息量最大。那些已经过于简单、始终过难、或在样本间方差很小的提示词，贡献的信号有限，会被降低优先级或移除。在实践中，有效的强化学习需要将提示词选择视为一个资源分配问题，即训练过程应将更多计算量投入到模型尚能学习的示例上，而对提供饱和或噪声反馈的示例投入更少。

扩散模型强化学习中的另一个实际考虑是如何处理无分类器引导（CFG）。生成的推演和训练都可以在有或无 CFG 的情况下进行，不同的选择会在对齐、稳定性和效率之间做出不同的权衡。经过消融实验后，我们发现保持推演分布和训练分布对齐、同时避免不必要的计算开销非常重要。因此，我们在整个强化学习阶段都在无 CFG 的情况下进行训练。这种设置能快速改善条件模型分布，使得无 CFG 的样本在训练早期就接近有引导的样本。在推理时，CFG 仍然可以作为额外的控制旋钮启用，在需要时进一步提高质量。

时间步蒸馏

在强化学习阶段之后，我们增加了一个可选的 timestep 蒸馏阶段，在该阶段中同时应用引导蒸馏和 timestep 蒸馏。我们考虑了多种蒸馏技术，包括 DMD、DMD2、解耦 DMD、piFlow 和 APT，但最终采用了轨迹分布匹配（TDM），原因如下。我们寻找一种易于调节、超参数最少的技术，这排除了基于 GAN 的方法和 piFlow（后者需要将模型改造为多 timestep 预测模型）。我们选择 TDM，因为它提供了一种快速、无数据且支持灵活多步蒸馏的方法。

DMD 通过匹配真实样本与生成样本在干净图像分布上的分布来蒸馏教师模型。因此，标准 DMD 使用少步学生模型预测干净图像，然后对预测结果重新加噪以训练学生模型（见上图）。与仅匹配干净图像分布的 DMD 不同，TDM 在多个 timestep 上应用 DMD，实际上是在轨迹层面而非样本层面执行分布匹配。由于我们的目标是获得灵活的多步学生模型，我们发现 TDM 是最适合我们场景的方法。

提示词扩展

密集提示词能够可靠地产生更好的图像生成结果，但用户很少会写出类似训练时所使用的丰富描述。我们将此问题视为一个分布映射问题：图像模型最适合以接近其训练分布的详细描述为条件，而用户的实际提示词往往简短、口语化且信息不足。因此，我们开发了一个提示词扩展器，用于解读用户意图，并将输入提示词映射为更丰富、更符合模型需求的描述。

我们首先对现有开源大语言模型进行监督微调。为了整理训练数据，我们使用另一个语言模型从长描述中生成合成的“用户描述”：更简短、更具对话性、半指令性的提示词，有意省略目标描述中许多视觉细节。由此产生了一组成对数据，形式为“未充分说明的用户提示词 → 扩展后的、模型友好的描述”。我们还合成思考轨迹以保持模型的推理能力，因为我们发现中间意图重建步骤能改善下游行为。除此之外，我们还进行了少量有针对性的分布塑造。从宏观角度看，我们对视觉丰富和艺术性的图像进行过采样，使扩展器覆盖更多创意和审美提示风格，并为应当扩展为逼真描写的提示词添加了轻微的摄影介质偏好。其目的不是强加某种固定风格，而是确保扩展后的提示词分布既涵盖表现性的、艺术导向的图像，也涵盖直接的逼真请求。

监督微调让扩展器接近所需的描述分布，但匹配合成目标不等于改善最终图像。因此，我们使用强化学习直接通过扩展器生成的图像来优化它。在此阶段，目标从模仿目标描述转变为生成能够在保留用户意图的同时提升图像质量的扩展。我们在多奖励目标下使用GDPO进行训练：图像级奖励衡量生成结果的质量和偏好，而提示词级可验证奖励则检查扩展是否忠实于原始请求。我们还加入了安全性和约束检查，以防止扩展器引入明显未提示或不可接受的内容。由于这些检查本质上是稀疏的，我们将其用作整体奖励的闸门，而非密集优化信号。

我们力求在匹配用户输入的真实分布与确保每组包含高对比度的好坏内容之间取得平衡。为此，RL提示词混合将逼真的用户类流量与挖掘出的困难案例相结合。逼真部分包括实际提示词和之前观察到的失败案例，并辅以人工筛选的失败案例以及已知失败模式的合成增强。人工策划的示例来自内部踩、错误报告和手动重写的提示词，并按桶分组，以便每个已知失败类别都能得到体现。合成示例从手工编写的奖励触发探针开始，扩展为多种变体，并使用RL期间使用的相同奖励模式进行离线评分。我们选择那些“困难但并非无望”的提示词，生成具有中间综合分数和实际奖励差异的组，这样GDPO就能接收到有意义的偏好信号，而不仅仅是显而易见的成功或彻底失败。

我们明确优化的一种失败模式是多样性崩溃。提示词扩展器可能学会一种单一安全、高奖励的固定风格，尤其是在图像奖励占主导时。针对这一点，我们为提示词组添加了一个简单的DINOv3嵌入多样性分数，在奖励质量和对齐的同时，也奖励组内的视觉多样性。我们曾短暂尝试对多样性奖励进行退火处理，但发现一旦其权重变得过小，模型就会迅速朝多样化程度降低的方向崩溃。实践中，在整个训练过程中保持多样性奖励处于活跃状态对于维持多样性是必要的。

风格参考系统

我们的风格参考系统基于基础模型构建。它允许用户根据文本生成图像，同时使用一张或多张参考图像来引导输出风格。我们设计该系统以支持（1）多种风格的流畅语义混合，（2）对每种风格参考强度的连续控制，以及（3）对复杂风格的最先进遵循。

风格迁移之所以困难，是因为图像中“风格”与“内容”的界定存在歧义。最常见的失败模式之一是风格图像中的内容和主体物渗入最终生成图像。此外，与可以从视频等来源挖掘数据的传统编辑任务不同，风格迁移数据很难在我们所针对的保真度下大规模获取。

为了应对这些挑战，我们设计了一种新颖的自监督技术来训练风格参考模块，随后通过偏好优化步骤进一步对齐输出结果。

训练基础设施

我们的分布式训练框架完全基于 PyTorch 从头构建。我们主要依赖 DTensor 抽象以及 torchtitan 项目支持的 torch 原生特性。在大多数预训练和后训练运行中，我们同时使用 FSDP2 和 Megatron-LM 风格张量并行。对于 TP 大小大于 2 的设置，我们通过 torch.compile 标志启用异步 TP，相比朴素 TP 能带来适度的加速。由于自编码器参数增加的内存开销很小，我们将其在所有设备上复制，仅对文本编码器和主 MMDiT 骨干进行分片。节点内连接使用 NVLinkSharp，节点间连接使用 InfiniBand。

为了训练效率，我们采用了更宽、隐藏维度更大的模型，原因有二。其一，更大的隐藏尺寸增加了每层的计算强度，使得 FSDP2 预取更容易隐藏延迟；减少层数也减少了 all-gather 和 reduce-scatter 操作的数量。这一改变显著减少了整个预训练过程中与 NCCL 相关的错误。其二，更大的矩阵乘法规模有助于分摊 8 比特训练中量化和反量化的开销。

我们主要依赖 torch.compile 作为核心优化策略。对于注意力机制，我们默认使用最新的 cuDNN 内核，并按需采用 FlexAttention 或 FlashAttention 3。在低分辨率下，我们使用选择性激活检查点；而在更高分辨率下，当激活值开始占据内存主导时，我们使用全激活检查点。

在数据加载方面，我们使用 Parquet 作为主要格式。对于每一行，我们存储图片引用（例如本地路径或 S3 位置）、裁剪和缩放尺寸、文字描述以及任何其他相关元数据。在大型运行中，我们预先对行进行打乱和打包，以便每个数据加载器工作进程加载一批具有相同宽高比的图片。这种打包方式使我们能够在单次自编码器前向传播中对潜在表示进行编码。

这种实现方式有多个好处。通过预先打乱数据，我们可以对磁盘进行顺序扫描以实现高性能的数据加载，同时确保全局打乱得当。预先打乱对于可复现性和调试也至关重要，因为数据可以按精确顺序重放，以识别任何一个可能导致损失尖峰的样本。

在我们最大的预训练运行中，遇到了各种基础设施和容错挑战。传统的大规模分布式训练在集群中引入了许多全局同步点（例如 DP 副本之间的梯度全局归约），并且本质上是不稳定的：单卡 GPU 故障或掉队者就可能导致整个运行中断。虽然存在如 torch-ft 和解耦的 DiLoCo 等容错解决方案，但在我们的规模下，我们通过快速、频繁的检查点和改进的启动时间来优化平均故障间隔时间（MTBF）和平均恢复时间（MTTR），发现这是一种令人满意的解决方案。

影响可靠性的另一个关键因素，是在所有训练设备上维持 I/O、CPU 和 GPU 的负载均衡。我们有意识地对数据加载器进行设计，使每个 CPU 和 GPU 承受大致相等的负载。在早期的大规模运行中，低分辨率阶段使用的高分辨率图像会在运行时被即时裁剪并调整至低分辨率；然而，在最大规模的运行中，这种做法会因原始图像是否为高分辨率而给不同设备带来不均衡的 CPU 和 I/O 负载。为解决此问题，我们提前将所有图像裁剪并调整至目标训练分辨率。我们还确保每个 GPU 都接收到填充至完全相同形状的张量输入，从而均匀负载。

对于我们的强化学习基础设施，我们将奖励模型推理与主训练过程分离。由于 Krea 2 是我们的首次大规模 RL 迭代，我们采用了一种简单设计，训练 GPU 和 rollout GPU 共享。在不久的将来，我们计划实现一种分离式训练和推理架构，以支持类似 PipelineRL 的异步 RL 训练技术。

系统基础设施

我们的研究运行在一个单一的 Kubernetes 集群内，该集群中的 GPU 与生产环境推理共享。系统设计使得研究任务在需要时能够占用整个 GPU 池：如果集群中的每一个 GPU 都被分配给某个训练运行，那么 Krea 的推理工作负载会自动迁移到别处。这样一来，我们在发起训练运行时可以不必考虑生产容量，因为系统会自动处理流量故障切换，即使本地没有可用 GPU，也能保持生产环境的响应能力。

这一能力并非一开始就具备。调度和管理系统随着研究团队在研究周期内的需求演变而逐步发展。在接下来的几个小节中，我们将描述主要组件：使用 Kueue 进行工作负载调度、在集群外部扩展推理、将两者结合在一起的调度策略、我们的训练启动流程，以及对大规模预训练至关重要的可观测性栈。

使用 Kueue 进行调度

Kueue 在我们的架构中一直处于核心地位，尽管我们的使用方式已发生显著变化。它提供了一套双层优先级系统，结合了 Kueue 的工作负载优先级与 Kubernetes 的 Pod 优先级。如果配置得当，这种双层设计能够产生有用的调度语义。

我们曾考虑过 Volcano 等自定义调度器，也考虑过直接修改 kube-scheduler，这样本可带来一些额外期望的特性。但最终，默认的 Kubernetes 调度器结合 Kueue 已经足够满足我们的需求。我们预计，随着规模持续增长，我们最终将不得不放弃 Kueue 和默认调度器，转向自定义调度器——考虑到我们工作负载的特殊性，这很可能在不久的将来发生——但眼下尚未到达那个阶段。

Kueue 支持组调度（gang-scheduling），这对于多节点训练是必要的。此外，“借入”、“借出”和“回收”等排队原语有助于最大化利用率。不过，我们希望 Kueue 能够动态推断每个队列的 GPU 数量，而不是手动指定。当节点数量发生变化时，这种手动操作一直是个令人困扰的问题。

在集群外部扩展推理

第二个组件解决的是：当所有 GPU 都分配用于研究时，推理如何在其他地方进行扩展。我们最初的设计引入了相当高的复杂性，最终我们采用了一种更简单的方法，基于一个虚拟 Kubelet（Virtual Kubelet，VK）来模拟 Kubernetes 节点。由于最终系统完全以 Kubernetes 原语表达，所需的大部分行为直接从 Kubernetes 继承而来。

我们首先评估了同样基于虚拟 Kubelet 构建的 InterLink，但发现它不适合我们的用例。因此，我们在 VK 之上构建了自己的层，以提供我们系统所需的语义，并为集成新的 GPU 提供商提供清晰的接口——前提是每个提供商都提供以编程方式伸缩容量的能力。

系统运行方式如下：在集群中注册一个虚拟 Kubernetes 节点。当一个 Pod 被调度到该节点上时，我们的代码会将 Pod 规格转换为与目标提供商兼容的形式，部署到目标侧，并在提供商侧发生故障时对两边进行协调。

此设计的一个关键优势是：扩缩容行为继承自 Kubernetes。假设配置了一个 HPA，目标副本数为十个，且全部在集群外运行，其中某个副本发生故障。协调循环不会尝试修复该副本，而是检测到故障，将 Pod 标记为失败（随后将其垃圾回收，以避免积累死 Pod），并允许 HPA 调度一个替换副本。因此，系统并不直接尝试修复故障；它会检测故障、向 Kubernetes 传播故障，然后委托 Kubernetes 进行恢复。

该集成需要实现一小套原语：Pod 创建、Pod 删除、exec（必要时），以及节点的启动行为——例如，当我们的代码崩溃而 Pod 仍在提供商侧运行时如何恢复。一个专用层吸收了针对我们系统的特定需求，而新提供商则通过一个清晰的接口接入这些原语，从而使业务逻辑不会泄露到提供商实现中。

启动训练任务

这个组件比其他的更简单，但仍然让我们在人体工学和哪些值得自动化方面学到了很多。随着 GPU 数量增长，更多运维问题出现，我们的启动流程也随之演变。

在大规模训练期间，我们经常遇到少量故障节点（下文将进一步讨论），它们会中断训练，直到这些节点被轮换出去。最初我们在启动前将已知故障节点封锁，让 Pod 调度到健康节点池上。然而，随着时间的推移，故障节点集不断扩大，每次崩溃都需要大量人工干预。这个故障节点列表最初只是以纯文本文件形式保存节点名称。

我们还观察到，在节点上运行额外进程——甚至是那些不访问 GPU 的进程——随着训练扩展到更多节点，会引入不稳定性。因此，手动流程逐渐演变为：先识别出一组干净的节点，为它们打上标签，以便训练 Pod 能够通过亲和性调度，再对这些节点进行污点标记，以防止其他工作负载（训练除外）调度到其上，最后将它们排空。

污点标记仅在需要最大稳定性的超大规模训练中才必要。因此，该流程被部分自动化：启动 CLI 会检索故障节点列表，排除这些节点以及已经在运行训练或开发机的节点，从剩余节点中选出所需数量，并自动为其打标签和污点标记。在拆除时，它会移除这些标签和污点。

随后，这个列表从文本文件迁移到了节点标签上，从而使我们能够直接对故障节点表达软反亲和性或硬反亲和性。这还催生了一个 Kubernetes 算子“Packerman”，它可以将开发机等 Pod 打包调度到故障节点上，从而让健康节点维持空闲用于训练。在大多数情况下，运行在故障节点上的开发机是可以接受的：在开发过程中，节点中某块 GPU 比其他 GPU 温度高 5–10 摄氏度是可以容忍的，但在训练中则不行。

可观测性

指标

可观测性是我们在大规模预训练中学到最多的领域。每次崩溃都有一组反复出现的原因，但整个过程中不断出现新的、意料之外的故障模式，随着时间的推移，我们变得善于诊断给定崩溃的原因。如果没有针对 GPU、PCIe、NVLink、InfiniBand 及相关子系统的指标，如此规模的训练是不可能实现的。我们通过 DCGM 和自定义 DaemonSet 相结合的方式收集这些指标，后者还会输出其他感兴趣的指标。

下面我们描述几个最有用的指标，以及当训练偏离预期值时它们的行为表现。在原始报告中，每个指标都附有一张来自实际生产运行的图表。

DCGM_FI_DEV_GPU_TEMP GPU 对温度极为敏感。我们发现，任何 GPU 温度超过 75–78 摄氏度都会增加训练不稳定性，因为此时 GPU 开始降频，整体吞吐量下降。

DCGM_FI_DEV_GPU_UTIL 我们发现这个指标常常具有误导性，因为它往往呈现的是症状而非原因。许多分布式训练错误始于单个进程崩溃，我们发现这一指标常报告某个 GPU 利用率骤降为 0%，随后很快崩溃，这样的报告几乎没什么帮助。此外，我们认为这个“利用率”指标总体上具有误导性，因为它只报告“任何 CUDA 内核在 GPU 上运行的时间占比”，而非 GPU 能力的“真正”使用情况。

DCGM_FI_PROF_PIPE_TENSOR_ACTIVE Tensor core 利用率是我们判断训练是否达到预期的主要指标。我们在不同训练阶段观察到明显的相关性：分辨率越高的图像，tensor core 利用率越高。它还可以作为其他状况的代理指标：出现过热 GPU 的节点会表现出异常高的 tensor core 利用率（与直觉相反），而一次运行过程中利用率持续下降通常表明存在故障节点或其他不稳定性。

DCGM_FI_DEV_FB_USED 内存使用量很有用，不过其原因与通常预期的不同。OOM 事件很少发生，一旦配置稳定后，实际上从未发生过。这个指标的作用在于：少数 GPU 在内存分配时失败或停滞——大部分 GPU 在训练开始时分配了内存，而少数 GPU 的显存卡在约 5 GiB。纠正措施是重启受影响的节点。

DCGM_FI_DEV_XID_ERRORS 与内存使用量类似，XID 错误在少数情况下有用。大多数崩溃不会产生 XID，但一旦出现 XID，就能直接确定原因。

DCGM_FI_DEV_CORRECTABLE_REMAPPED_ROWS、DCGM_FI_DEV_UNCORRECTABLE_REMAPPED_ROWS 和 DCGM_FI_DEV_ROW_REMAP_FAILURE 这些指标对识别有故障内存模块的 GPU 很有用；当根因是内存问题时，信号通常非常明确。

DCGM_FI_DEV_PCIE_REPLAY_COUNTER — 该指标报告 GPU 与主板之间连接的质量。在特定 GPU 上出现的一连串重放一致地出现在崩溃之前。零散分布于随机 GPU 的偶发性重放是无害的；故障模式表现为单个 GPU 上出现一个巨大的尖峰。

NVLink 错误（自定义）— DCGM 默认不导出这些指标，因此我们实现了一个自定义 DaemonSet 来收集 CRC、重放和恢复错误。在极少数情况下，与 PCIe 重放类似，这些指标使我们能够识别出有故障的节点。

InfiniBand 指标（自定义）— IB 指标可以说是我们收集的最重要的指标，尤其是因为集群网络结构并非始终稳定。我们收集了以下指标：

VL15_dropped
excessive_buffer_overrun_errors
link_downed
link_error_recovery
local_link_integrity_errors
multicast_rcv_packets
multicast_xmit_packets
port_rcv_constraint_errors
port_rcv_data
port_rcv_errors
port_rcv_packets
port_rcv_remote_physical_errors
port_rcv_switch_relay_errors
port_xmit_constraint_errors
port_xmit_data
port_xmit_discards
port_xmit_packets
port_xmit_wait
symbol_error
unicast_rcv_packets
unicast_xmit_packets

这些指标对于诊断各类 InfiniBand 相关问题至关重要。网络结构不稳定是导致运行崩溃的最大单一因素，包括链路振荡、数据包错误、拥塞、符号错误以及设备间吞吐量差异。这些指标使我们能够定位网络结构中行为异常的部分——例如，局限于 mlx5_10 和 mlx5_11 上的数据包错误，或 mlx5_0 上的吞吐量下降。我们发现 ibtop 对于调试这些问题非常有用。

经验教训

大规模训练让我们对集群的行为以及更广义的大规模训练有了深刻的了解。早期，我们对崩溃的反应很激进：更换节点、重新启动、更改配置。随着时间的推移，我们学会了区分值得处理的问题和暂时的异常。在某些情况下，一次训练运行会在两个小时内反复崩溃，然后除了重新启动之外没有任何变化，却能连续运行十六个小时而未出现任何问题。

我们还遇到了扩展 GPU 数量时的困难，我们的经验可能与其他预训练工作不同。通常作为规模函数报告的故障率与我们的观察并不相符：将 GPU 数量翻倍所产生的不稳定性远超预期。举例来说，在相同的代码和数据条件下，少于 128 个 GPU 的运行非常稳定，常常能连续运行数天不出事故，崩溃通常可归因于真正的硬件故障。然而，随着我们扩展 GPU 数量，运行开始更加频繁地崩溃，并且在超大规模下，我们没有一次运行能超过 24 小时而不崩溃。其中许多崩溃没有明显原因；它们悄无声息地出现，例如 NCCL 超时，而此时所有指标看起来都正常。

文件系统和检查点

我们早期最严重的错误之一就是采用了 Ceph，它在我们的规模或使用场景下表现不佳。因此，我们切换到了 Weka，并且从未重新考虑过这个决定。与文件系统相关的问题和宕机时间急剧下降，而性能则得到了相应的提升。这次迁移并非完全没有摩擦。在安装期间以及文件系统的整个生命周期中，我们遇到了许多怪癖和边界情况，但 Weka 始终积极响应，帮助我们解决了这些问题。鉴于我们采用时 Kubernetes 产品尚未完全成熟，这一点尤其宝贵——尽管自那以后它已经有了很大改进，增加了许多新功能和错误修复。

正如在可靠性部分所讨论的，Weka 是训练 Krea 2 的关键推动因素，我们在训练中非常激进地进行检查点保存。该文件系统跟上了这一需求：一个检查点大约在 30 秒内完成，因此几乎没有时间浪费在检查点上。这个单一文件系统存储了所有内容——研发数据、图像、检查点、数据集以及所有其他工件——并且在每种情况下都表现出色。

一个值得注意的小问题（并非出在 Weka 身上）与我们使用 IPoIB 模式有关：我们怀疑它会导致集群中 InfiniBand 网络结构不稳定。不过，该网络结构在引入 Weka 之前就已经不稳定了，所以最多只是加剧了原有的问题。

Weka 集群的最大读取吞吐量

Weka 集群的最大写入吞吐量

数据基础设施

为了摄取和整理 K2 的数据，我们围绕一组 PostgreSQL 服务器构建了自定义的数据仓储和队列系统。我们把每台 Krea 平板服务器称为一个 krablet。每个 krablet 由以下部分组成：

• 一个持有数据分片的 Postgres 实例，以及
• 该分片上的一个“漏斗”服务器部署，用于异步批量处理和排队变更操作，以尽量减少锁竞争。

所有读取操作都通过大规模部署的“RPC”服务器进行代理，取代了像 PgBouncer 这样的传统连接池。每台 RPC 服务器都维护到数据库每个分片的连接池。

krablet 服务器保存我们训练数据集的所有元数据，并存储引用对象存储中图像的键。

虽然 krablet 系统并非传统架构，但它已扩展到仅元数据就达到 208 TB，并且每秒可以处理数万笔存在竞争的 UPSERT 事务。更重要的是，它为所有研究数据提供了单一事实来源，并让我们的流处理层与数据层保持一致。

Krea 的一个典型作业处理流程如下。我们首先创建一个表，例如：

CREATE TABLE images (
	image_name TEXT PRIMARY KEY,

	contains_text BOOLEAN DEFAULT NULL,
	ocr_last_tried_at TIMESTAMP DEFAULT NULL,

	embedding_path TEXT DEFAULT NULL,
	embed_last_tried_at TIMESTAMP DEFAULT NULL
)

该表定义了每张图像可以执行的两个作业：OCR 和嵌入。OCR 工作节点运行以下查询来查找待处理的行：

-- OCR worker
WITH picked AS (
	SELECT image_Name FROM images
	WHERE contains_text IS NULL
	ORDER BY ocr_last_tried_at NULLS FIRST
	LIMIT 8
	FOR UPDATE SKIP LOCKED
)
UPDATE images i
SET ocr_last_tried_at = NOW()
FROM picked p
WHERE i.id = p.id
RETURNING i.*

嵌入工作节点则运行：

-- Embed worker
WITH picked AS (
	SELECT id FROM images
	WHERE embedding_path IS NULL
	AND contains_text = FALSE -- Only process rows which have completed OCR and passed a filter
	ORDER BY embed_last_tried_at NULLS FIRST
	LIMIT 16
	FOR UPDATE SKIP LOCKED
)
UPDATE images i
SET embed_last_tried_at = NOW()
FROM picked p
WHERE i.id = p.id
RETURNING i.*

隐式地，这个处理流程是一个 DAG（有向无环图）：

许多大规模数据处理系统（包括 Ray、Spark 和 Daft）都可以处理 DAG 工作流。然而，在此模型中将这些工作流视为队列，可以给我们带来几个实际好处：

• 重试机制：与 Kafka 或 Ray 不同，该系统在失败时不会丢弃数据行，也不会将其发送到死信队列（DLQ）。任何因任何原因处理失败的数据行，都会在队列末尾重新尝试（依托于 last_tried_at 的原子更新），这同时还避免了队头阻塞问题。
• 容错性：任何工作节点可随时崩溃，而不会导致整个任务失败或进度丢失。许多类型的故障——例如坏的 GPU 或临时不可用的 API——实际上都能通过系统的最终重试机制自动得到处理。
• 动态工作节点数：我们使用 Kubernetes 部署处理任务，任何部署都可以任意扩缩容，而无需重新分片或移动数据；一个任务可以只有一个工作节点，也可以有一千个。对于许多工作负载，我们会基于 claim 查询的 COUNT(*)（例如，`COUNT(*) WHERE embedding_path IS NULL AND contains_text = FALSE`）暴露一个 Prometheus 扩缩容指标，从而能根据可用工作量自动扩缩管道的各个部分。
• 部分处理：由于每一行结果都会立即持久化到 Postgres，用户可以按需处理表中任意多或少的数据；没有任何一个任务需要运行到完成。这使我们能够对许多处理任务利用可中断和竞价优先级类别，仅在有空闲容量时运行它们。
• 即时可见性：研究人员可以立即在可视化仪表盘中查看任何任务的结果，并进行持续迭代。他们还可以检查队列任意部分的吞吐量，例如使用 `SELECT COUNT(*) FROM images WHERE embed_last_tried_at > NOW() - INTERVAL '5 minutes'`。
• 异构批次大小：每个工作节点可以选择自己合适的批次大小。
• 持续增量处理：我们可以持续向系统添加新数据，使其自动流经各个阶段，而无需手动重新运行任务、回填旧数据行或跟踪不同阶段。

为方便研究人员，我们在此基础上暴露了一个名为“pluck”的系统，它提供了适合在笔记本中使用的全局映射 API。例如：

def embed_func(batch):  # users may also provide stateful Actors rather than one-off functions
    ...

t = pluck.Table(
    "images",                                            # table name
    "image_name",                                        # primary / shard key
    "embedding_path IS NULL AND contains_text = FALSE",  # condition
)
t.map(embed_func)  # returns a handle the user can attach to in order to see live progress

在底层，pluck 使用 TABLESAMPLE 和统计估计来对表的主键空间进行分区，并创建一个批量队列，然后通过前面展示的 FOR UPDATE SKIP LOCKED 语义来消费该队列。用户自定义函数（UDF）会使用 cloudpickle 库进行序列化，并在远程工作节点上执行。

在我们下一代研究中，团队正在构建一个继承系统，该方案保留了 krablet 和 FOR UPDATE SKIP LOCKED 队列语义，但将数据存储在基于对象存储的 LSM 树中。如果你对此工作感兴趣，我们的超级计算/分布式系统团队正在招聘！

讨论与未来工作

尽管我们已经在整个技术栈的多个方向进行了探索，但仍有大量具有前景的研究问题有待解决。

规模扩展

对于 Krea 2，我们在架构和优化器方面做出了相对保守的选择，以优先保证稳定性和迭代速度。在下一个预训练周期中，我们希望将现代大语言模型的 Transformer 设计适配到扩散 Transformer 中，包括采用 MoE；利用稀疏注意力扩展到原生 2K—4K 分辨率；在 NVFP4 精度下进行预训练；以及使用 Muon 优化器进行扩展等方向。除了架构和算法层面的改动，我们还发现当前模型存在训练不足的问题，延长训练时间将使其受益。

多教师在线策略蒸馏（MOPD）

我们当前的训练流程以一个多奖励强化学习阶段收尾。在近期的大语言模型文献中，训练领域专家模型并通过密集的专家监督来蒸馏其能力已成为常见做法。与模型融合和离线蒸馏等方法相比，MOPD 已被证明能够将多个专家的能力蒸馏到单个学生模型中，且不同能力之间不会相互冲突。从组织角度来看，MOPD 也具有很高的可扩展性，因为它允许不同的强化学习团队专注于特定领域的能力，而无需担心在其他方面造成回退。我们已通过内部专家验证了 OPD 和 MOPD 都是针对扩散模型的有效蒸馏方法，并希望很快分享更多成果。

架构简化

目前，大多数生产级扩散模型都需要一套相互依赖的复杂模型组。运行一个潜在扩散模型通常需要自编码器、扩散Transformer、文本编码器和提示词扩展模型，根据技术栈的不同，还可能添加风格参考模型或超分辨率模块等额外组件。维护多个需要独立训练且相互依赖的组件，使得研究团队协调工作变得非常困难。这与大语言模型训练流程形成鲜明对比——后者围绕单一统一模型展开，研究人员可以通过数据团队并行工作，或在MOPD阶段进行独立的专家训练；更重要的是，这允许研究团队将扩展工作集中在一个模型上。基于这些原因，我们计划在下一轮预训练周期中简化架构，将各个组件统一到单一模型之下。

新能力

在Krea 2中，我们主要聚焦于面向创意探索的图像生成。我们计划扩展Krea模型的能力，使其包含强大的编辑、图像参考以及原生2K/4K生成功能。我们也越来越发现，传统的自然语言提示词已不足以支持用户的全方位请求。在我们用户的提示词中，观察到了多种多样的提示风格，包括自然语言、标签、详细JSON、边界框、指令、视觉指南和Markdown等。虽然提示词扩展可以解决部分问题，但我们认为原生理解这类提示词也应该是模型的核心能力。

结论

我们介绍了Krea 2，这是我们为创意探索打造的首个基础模型系列。在这份技术报告中，我们分享了该模型系列背后的基础设施、数据系统、训练流程和研究选择。我们认为基础图像模型仍处于早期阶段，期待后续分享更多研究成果。

1. 去噪扩散概率模型。Ho, J., Jain, A. 和 Abbeel, P., 2020. 神经信息处理系统进展, 第 33 卷, 第 6840—6851 页. 链接

2. 基于潜在扩散模型的高分辨率图像合成。Rombach, R., Blattmann, A., Lorenz, D., Esser, P. 和 Ommer, B., 2022. IEEE/CVF 计算机视觉与模式识别会议论文集, 第 10684—10695 页. 链接

3. 流直且快：学习使用整流流生成和迁移数据。Liu, X., Gong, C. 和 Liu, Q., 2022. arXiv 预印本 arXiv:2209.03003. 链接

4. 用于生成建模的流匹配。Lipman, Y., Chen, R.T.Q., Ben-Hamu, H., Nickel, M. 和 Le, M., 2022. arXiv 预印本 arXiv:2210.02747. 链接

5. 视觉自回归建模：通过下一尺度预测实现可扩展图像生成。Tian, K., Jiang, Y., Yuan, Z., Peng, B. 和 Wang, L., 2024. 神经信息处理系统进展, 第 37 卷, 第 84839—84865 页. 链接

6. Infinity：扩展按位自回归建模以实现高分辨率图像合成。Han, J., Liu, J., Jiang, Y., Yan, B., Zhang, Y., Yuan, Z., Peng, B. 和 Liu, X., 2025. 计算机视觉与模式识别会议论文集, 第 15733—15744 页. 链接

7. Emu：多模态下的生成式预训练。Sun, Q., Yu, Q., Cui, Y., Zhang, F., Zhang, X., Wang, Y., Gao, H., Liu, J., Huang, T. 和 Wang, X., 2024. 国际学习表征会议, 第 2024 卷, 第 12352—12380 页. 链接

8. 从像素开始的生成式预训练。Chen, M., Radford, A., Child, R., Wu, J., Jun, H., Luan, D. 和 Sutskever, I., 2020. 国际机器学习大会, 第 1691—1703 页. PMLR. 链接

9. 扩展自回归模型以实现内容丰富的高质量文本到图像生成。Yu, J., Xu, Y., Koh, J.Y., Luong, T., Baid, G., Wang, Z., Vasudevan, V., Ku, A., Yang, Y., Ayan, B.K. 等, 2022. arXiv 预印本 arXiv:2206.10789. 链接

10. Maskgit：掩码生成式图像 Transformer。Chang, H., Zhang, H., Jiang, L., Liu, C. 和 Freeman, W.T., 2022. IEEE/CVF 计算机视觉与模式识别会议论文集, 第 11315—11325 页. 链接

11. 基于Transformer的可扩展扩散模型。Peebles, W. 和 Xie, S., 2023. IEEE/CVF国际计算机视觉大会论文集，第4195—4205页。链接

12. Pixart-alpha: 用于逼真文本到图像合成的扩散Transformer快速训练。Chen, J., Yu, J., Ge, C., Yao, L., Xie, E., Wang, Z., Kwok, J., Luo, P., Lu, H. 和 Li, Z., 2024. 国际学习表征会议，2024卷，第57611—57640页。链接

13. 用于高分辨率图像合成的缩放整流流Transformer。Esser, P., Kulal, S., Blattmann, A., Entezari, R., Muller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F. 等，2024. 第四十一届国际机器学习大会。链接

14. Sit: 基于可缩放插值Transformer的流与扩散生成模型探索。Ma, N., Goldstein, M., Albergo, M.S., Boffi, N.M., Vanden-Eijnden, E. 和 Xie, S., 2024. 欧洲计算机视觉大会，第23—40页。Springer。链接

15. Sdxl: 改进用于高分辨率图像合成的潜在扩散模型。Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Muller, J., Penna, J. 和 Rombach, R., 2024. 国际学习表征会议，2024卷，第1862—1874页。链接

16. FLUX。BlackForest，2024。链接

17. FLUX.2: 前沿视觉智能。Black Forest Labs，2025。链接

18. FLUX.1 Kontext: 潜在空间中上下文图像生成与编辑的流匹配。Labs, B.F., Batifol, S., Blattmann, A., Boesel, F., Consul, S., Diagne, C., Dockhorn, T., English, J., English, Z., Esser, P. 等，2025. arXiv预印本 arXiv:2506.15742。链接

19. FLUX.1 Krea [dev]。Black Forest Labs 和 Krea，2025。链接

20. Lumina-Image 2.0: 统一高效的图像生成框架。Qin, Q., Zhuo, L., Xin, Y., Du, R., Li, Z., Fu, B., Lu, Y., Yuan, J., Li, X., Liu, D. 等，2025. arXiv预印本 arXiv:2503.21758。链接

21. Hunyuan-dit: 具有细粒度中文理解能力的高性能多分辨率扩散Transformer。Li, Z., Zhang, J., Lin, Q., Xiong, J., Long, Y., Deng, X., Zhang, Y., Liu, X., Huang, M., Xiao, Z. 等，2024. arXiv预印本 arXiv:2405.08748。链接

22. Qwen-image 技术报告。吴辰、李嘉骏、周杰、林佳、高凯、闫凯、尹思敏、白硕、徐鑫、陈杨等人，2025年。arXiv 预印本 arXiv:2508.02324。链接

23. Longcat-image 技术报告。ML团队、马浩、谭浩、黄洁、吴杰、何建宇、高磊、肖帅、魏翔、马晓等人，2025年。arXiv 预印本 arXiv:2512.07584。链接

24. JoyAI-Image：在多模态统一理解与生成中唤醒空间智能。京东探索研究院，2026年。链接

25. HunyuanImage-2.1。腾讯混元，2025年。链接

26. Hunyuanimage 3.0 技术报告。曹松、陈浩、陈鹏、程远、崔宇、邓旭、董宇、龚凯、顾涛、徐翔等人，2025年。arXiv 预印本 arXiv:2509.23951。链接

27. Z-image：采用单流扩散Transformer的高效图像生成基础模型。蔡昊、曹松、杜瑞、高鹏、许少辉、侯振、黄松、江东、金鑫、李磊等人，2025年。arXiv 预印本 arXiv:2511.22699。链接

28. Seedream 3.0 技术报告。高宇、龚磊、郭强、侯鑫、赖哲、李飞、李琳、连翔、廖晨、刘磊等人，2025年。arXiv 预印本 arXiv:2504.11346。链接

29. Seedream 2.0：一个原生中英双语图像生成基础模型。龚磊、侯鑫、李飞、李琳、连翔、刘飞、刘磊、刘伟、陆炜、施展等人，2025年。arXiv 预印本 arXiv:2503.07703。链接

30. Seedream 4.0：迈向下一代多模态图像生成。Seedream团队、陈宇、高宇、龚磊、郭敏、郭强、郭志、侯鑫、黄伟、黄宇等人，2025年。arXiv 预印本 arXiv:2509.20427。链接

31. Seedream 5.0 Lite。字节跳动豆包，2025年。链接

32. GPT Image 1.5。OpenAI，2025年。链接

33. Nano Banana Pro。Google，2025年。链接

34. Qwen3-vl 技术报告。白硕、蔡宇、陈睿、陈凯、陈晓、程志、邓磊、丁伟、高畅、葛晨等人，2025年。arXiv 预印本 arXiv:2511.21631。链接

35. 自编码变分贝叶斯。Kingma, D.P. 和 Welling, M.，2013年。arXiv 预印本 arXiv:1312.6114。链接

36. Flow-GRPO：通过在线强化学习训练流匹配模型。刘杰、刘刚、梁健、李宇、刘佳、王翔、万鹏、张东和欧阳文，2026年。神经信息处理系统进展，第38卷，第40783—40818页。链接

37. Diffusionnft：基于前向过程的在线扩散强化学习。作者：Zheng, K., Chen, H., Ye, H., Wang, H., Zhang, Q., Jiang, K., Su, H., Ermon, S., Zhu, J. and Liu, M.Y.，2025。arXiv预印本 arXiv:2509.16117。链接

38. 改进的美学预测器（LAION-Aesthetics Predictor V2）。作者：Schuhmann, C.，2022。GitHub仓库，christophschuhmann/improved-aesthetic-predictor。链接

39. ArtiMuse：基于联合评分与专家级理解的细粒度图像美学评估。作者：Cao, S., Ma, N., Li, J., Li, X., Shao, L., Zhu, K., Zhou, Y., Pu, Y., Wu, J., Wang, J. et al.，2025。arXiv预印本 arXiv:2507.14533。链接

40. UniPercept：面向美学、质量、结构与纹理的统一感知级图像理解。作者：Cao, S., Li, J., Li, X., Pu, Y., Zhu, K., Gao, Y., Luo, S., Xin, Y., Qin, Q., Zhou, Y. et al.，2025。arXiv预印本 arXiv:2512.21675。链接

41. DINOv3。作者：Simeoni, O., Vo, H.V., Seitzer, M., Baldassarre, F., Oquab, M., Jose, C., Khalidov, V., Szafraniec, M., Yi, S., Ramamonjisoa, M. et al.，2025。arXiv预印本 arXiv:2508.10104。链接

42. SigLIP 2：改进语义理解、定位与密集特征的多语言视觉-语言编码器。作者：Tschannen, M., Gritsenko, A., Wang, X., Naeem, M.F., Alabdulmohsin, I., Parthasarathy, N., Evans, T., Beyer, L., Xia, Y., Mustafa, B. et al.，2025。arXiv预印本 arXiv:2502.14786。链接

43. 自监督学习的自动数据整理：一种基于聚类的方法。作者：Vo, H.V., Khalidov, V., Darcet, T., Moutakanni, T., Smetanin, N., Szafraniec, M., Touvron, H., Couprie, C., Oquab, M., Joulin, A. et al.，2024。arXiv预印本 arXiv:2405.15613。链接

44. GLU变体改进Transformer架构。作者：Shazeer, N.，2020。arXiv预印本 arXiv:2002.05202。链接

45. GQA：从多头检查点训练通用化多查询Transformer模型。作者：Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F. and Sanghai, S.，2023。arXiv预印本 arXiv:2305.13245。链接

46. DeepSeek-V2：一个强大、经济、高效的混合专家语言模型。作者：DeepSeek-AI et al.，2024。arXiv预印本 arXiv:2405.04434。链接

47. 大语言模型的门控注意力：非线性、稀疏性与无注意力沉溺。Qiu, Z., Wang, Z., Zheng, B., Huang, Z., Wen, K., Yang, S., Men, R., Yu, L., Huang, F., Huang, S. 等人，2025 年。arXiv 预印本 arXiv:2505.06708。链接

48. SANA：使用线性扩散 Transformer 进行高效高分辨率图像合成。Xie, E., Chen, J., Chen, J., Cai, H., Tang, H., Lin, Y., Zhang, Z., Li, M., Zhu, L., Lu, Y. 等人，2024 年。arXiv 预印本 arXiv:2410.10629。链接

49. 关于 N 维旋转位置嵌入。Xiong, J.，2025 年。博客文章，jerryxio.ng。链接

50. 用于高效高分辨率扩散模型的深度压缩自编码器。Chen, J., Cai, H., Chen, J., Xie, E., Yang, S., Tang, H., Li, M., Lu, Y. 和 Han, S.，2025 年。国际学习表征会议，第 2025 卷，第 96539—96560 页。链接

51. 均方根层归一化。Zhang, B. 和 Sennrich, R.，2019 年。神经信息处理系统进展，第 32 卷。链接

52. UniFusion：视觉语言模型作为图像生成中的统一编码器。Li, K., Brack, M., Katakol, S., Ravi, H. 和 Kale, A.，2025 年。arXiv 预印本 arXiv:2510.12789。链接

53. 大语言模型预训练中的模型合并。Li, Y., Ma, Y., Yan, S., Zhang, C., Liu, J., Lu, J., Xu, Z., Chen, M., Wang, M., Zhan, S. 等人，2025 年。arXiv 预印本 arXiv:2505.12082。链接

54. 使用直接偏好优化的扩散模型对齐。Wallace, B., Dang, M., Rafailov, R., Zhou, L., Lou, A., Purushwalkam, S., Ermon, S., Xiong, C., Joty, S. 和 Naik, N.，2023 年。arXiv 预印本 arXiv:2311.12908。链接

55. FireRed-Image-Edit-1.0 技术报告。超级智能团队，Qiao, C., Hui, C., Li, C., Wang, C., Song, D., Zhang, J., Li, J., Xiang, Q., Wang, R. 等人，2026 年。arXiv 预印本 arXiv:2602.13344。链接

56. 在可微分奖励上直接微调扩散模型。Clark, K., Vicol, P., Swersky, K. 和 Fleet, D.J.，2023 年。arXiv 预印本 arXiv:2309.17400。链接

57. ImageReward：学习和评估文生图的人类偏好。Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang, J. 和 Dong, Y.，2023 年。arXiv 预印本 arXiv:2304.05977。链接

58. 直接对齐完整扩散轨迹与细粒度人类偏好。Shen, X., Li, Z., Yang, Z., Zhang, S., Zhang, Y., Li, D., Wang, C., Lu, Q. and Tang, Y., 2025. arXiv 预印本 arXiv:2509.06942. 链接

59. 无分类器扩散引导。Ho, J. and Salimans, T., 2022. arXiv 预印本 arXiv:2207.12598. 链接

60. 基于分布匹配蒸馏的单步扩散。Yin, T., Gharbi, M., Zhang, R., Shechtman, E., Durand, F., Freeman, W.T. and Park, T., 2024. IEEE/CVF 计算机视觉与模式识别会议，第 6613—6623 页。链接

61. 用于快速图像合成的改进分布匹配蒸馏。Yin, T., Gharbi, M., Park, T., Zhang, R., Shechtman, E., Durand, F. and Freeman, W.T., 2024. 神经信息处理系统进展，第 47455—47487 页。链接

62. 解耦 DMD：以 CFG 增强为矛，以分布匹配为盾。Liu, D., Gao, P., Liu, D., Du, R., Li, Z., Wu, Q., Jin, X., Cao, S., Zhang, S., Li, H. et al., 2025. arXiv 预印本 arXiv:2511.22677. 链接

63. pi-Flow：基于模仿蒸馏的策略驱动少步生成。Chen, H., Zhang, K., Tan, H., Guibas, L., Wetzstein, G. and Bi, S., 2025. arXiv 预印本 arXiv:2510.14974. 链接

64. 用于单步视频生成的扩散对抗后训练。Lin, S., Xia, X., Ren, Y., Yang, C., Xiao, X. and Jiang, L., 2025. arXiv 预印本 arXiv:2501.08316. 链接

65. 通过轨迹分布匹配学习少步扩散模型。Luo, Y., Hu, T., Sun, J., Cai, Y. and Tang, J., 2025. arXiv 预印本 arXiv:2503.06674. 链接

66. HPSv3：迈向广谱人类偏好评分。Ma, Y., Shui, Y., Wu, X., Sun, K. and Li, H., 2025. arXiv 预印本 arXiv:2508.03789. 链接

67. EvalMuse-40K：用于文生图模型评估的可靠且细粒度的基准，包含全面人工标注。Han, S., Fan, H., Fu, J., Li, L., Li, T., Cui, J., Wang, Y., Tai, Y., Sun, J., Guo, C. et al., 2024. arXiv 预印本 arXiv:2412.18150. 链接

68. Adam：一种随机优化方法。Kingma, D.P. and Ba, J., 2014. arXiv 预印本 arXiv:1412.6980. 链接

69. Muon 可扩展用于大语言模型训练。Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., Qin, Y., Xu, W., Lu, E., Yan, J. 等，2025。arXiv 预印本 arXiv:2502.16982。链接

70. 表示对齐的关键因素：全局信息还是空间结构？Singh, J., Leng, X., Wu, Z., Zheng, L., Zhang, R., Shechtman, E. 和 Xie, S.，2025。arXiv 预印本 arXiv:2512.10794。链接

71. Kimi K2.6。月之暗面，2026。链接

72. Delta 注意力残差。Luo, C., Cai, Z. 和 Hu, J.，2026。arXiv 预印本 arXiv:2605.18855。

73. 降噪生成模型是否必然需要噪声条件？Sun, Q., Jiang, Z., Zhao, H. 和 He, K.，2025。arXiv 预印本 arXiv:2502.13129。

74. LAuReL：学习型增强残差层。Menghani, G., Kumar, R. 和 Kumar, S.，2024。arXiv 预印本 arXiv:2411.07501。

75. 改进均值流：关于快进生成模型的挑战。Geng, Z., Lu, Y., Wu, Z., Shechtman, E., Kolter, J.Z. 和 He, K.，2025。arXiv 预印本 arXiv:2512.02012。

76. 注意力残差。Kimi Team，Chen, G., Zhang, Y., Su, J., Xu, W., Pan, S., Wang, Y., Wang, Y. 等，2026。arXiv 预印本 arXiv:2603.15031。

77. NOBLE：通过非线性低秩分支加速 Transformer。Smith, E.，2026。arXiv 预印本 arXiv:2603.06492。

78. mHC：流形约束超连接。Xie, Z., Wei, Y., Cao, H., Zhao, C., Deng, C., Li, J., Dai, D., Gao, H. 等（DeepSeek-AI），2025。arXiv 预印本 arXiv:2512.24880。

79. Gemma：基于 Gemini 研究与技术的开放模型。Gemma Team，2024。arXiv 预印本 arXiv:2403.08295。

80. 层归一化。Ba, J.L., Kiros, J.R. 和 Hinton, G.E.，2016。arXiv 预印本 arXiv:1607.06450。

81. 高斯误差线性单元 (GELUs)。Hendrycks, D. 和 Gimpel, K.，2016。arXiv 预印本 arXiv:1606.08415。

82. TREAD：面向高效架构无关扩散训练的 Token 路由。Krause, F., Phan, T., Gui, M., Baumann, S.A., Hu, V.T. 和 Ommer, B.，2025。arXiv 预印本 arXiv:2501.04765。

83. AReaL：用于语言推理的大规模异步强化学习系统。Fu, W., Gao, J., Shen, X., Zhu, C., Mei, Z., He, C., Xu, S., Wei, G. 等，2025。arXiv 预印本 arXiv:2505.24298。

84. PipelineRL：面向长序列生成的更快在线策略强化学习。作者：Piche, A., Kamalloo, E., Pardinas, R., Chen, X. 和 Bahdanau, D.，2025年。arXiv预印本编号：arXiv:2509.19128。

引用

@misc{krea-2-2026,
    author={Sangwu Lee, Erwann Millon, Le Zhuo, Matthew Newton, Andrei Filatov, Abhinay Devarinti, Andrei Filatov, Elea Zhong, Avram Djordjevic, Gabriel Menezes, Will Beddow, Titus Ebbecke, Mihai Petrescu, Owen Fahey, Gian Saß, Felix Gil, Victor Perez},
    title={{Krea 2}},
    year={2026},
    howpublished={\url{https://www.krea.ai/blog/krea-2-technical-report}},
}