Krea 2 技术报告 (59分钟阅读)

# Krea 2 技术报告 来源：https://www.krea.ai/blog/krea-2-technical-report ## 引言 在过去几年中，图像生成领域取得了显著的进展。扩散模型和流匹配模型能够生成高分辨率图像，产生锐利的照片级真实感和稳定的结构，渲染密集的文字，编码广泛的世界知识，并精确遵循用户提示的细节。这些改进由多个相互作用的因素驱动，包括可扩展的Transformer架构、改进的描述生成和文本编码器、更好的潜在表示以及流水线式后训练技术。然而，随着该领域在这些能力上优化可靠性，许多系统已收敛到一套狭窄的默认美学风格。虽然这些模型作为有效的生产工具，但它们作为创意探索引擎的效果却有所减弱，因为用户常常需要在不同风格、情绪、构图和视觉方向之间进行搜索，而不是仅仅接收一个单一的精致默认输出。 为了解决这些限制，我们推出了**Krea 2**，一个专注于创意探索的基础模型系列。Krea 2 的模型建立在这样一个信念之上：图像生成应该是一种探索性媒介——表达力足够强以覆盖多种美学风格，且可控性足够高以便创作者能够驾驭它们。 我们从零开始构建了大规模数据基础设施和分布式训练框架，以策划一个具有广泛世界知识和风格覆盖的综合预训练数据集。 利用这一基础设施，我们通过多阶段流水线训练表达性模型，该流水线涵盖预训练、中期训练、监督微调（SFT）、偏好优化和强化学习（RL），每个阶段都旨在逐步优化模型的输出分布。通过彻底的消融实验，我们开发了一种简单而高效的扩散Transformer（DiT）（https://arxiv.org/abs/2212.09748）架构。我们的模型整合了多个加速收敛的组件，包括iREPA（https://arxiv.org/abs/2512.10794）、改进的VAE（https://bfl.ai/research/representation-comparison）和Qwen3-VL（https://arxiv.org/abs/2511.21631）。我们还集成了多项架构改进，包括分组查询注意力（GQA）（https://arxiv.org/abs/2305.13245）、Sigmoid门控注意力（https://arxiv.org/abs/2505.06708）、轻量级时间步调制（https://arxiv.org/abs/2503.10618）以及文本编码器特征的多层特征聚合（https://arxiv.org/abs/2510.12789），这些改进共同提升了训练的稳定性和效率。 一个强大的基础模型只有在用户能够可靠地到达其分布中关心的部分时才有用。在训练中，模型从丰富且精心构建的描述中学习，这些描述以密集的视觉细节描述图像。而在实际应用中，用户输入通常更短、更模糊，并受到多种不同表达习惯的影响。有些用户用自然语言描述场景；另一些用户则倾向于表达一种情绪、一种风格或一张参考图像。这就造成了模型学到的条件空间与创意意图在推理时的表达方式之间的差距。 为了缩小这一差距，我们构建了两个系统，使Krea 2 更具探索性，并且能够通过文本和图像输入进行引导：一个提示扩展器和一个风格参考系统。提示扩展器将简单或不明确的用户提示映射为更丰富的视觉方向，而不会覆盖用户的意图。它通过基于开源LLM的两阶段SFT和RL流水线进行训练，其目标不仅是提高图像质量，还要鼓励创意变化和可控探索。作为这一文本界面的补充，风格参考系统允许用户在语言不足以表达时通过图像来表达视觉意图。它允许用户注入一张或多张参考图像的风格或情绪，同时最大限度地减少内容泄露，并提供对风格强度和加权风格混合的精细控制。 这些组件共同将Krea 2 定义为一个用于探索性生成的基础模型。Krea 2 不仅优化单一的精致默认输出，还设计用于展现广阔的视觉空间，并给用户提供实用的方式——通过文本和基于图像的控制——在其中移动。Krea 2 在Artificial Analysis文本到图像排行榜上位列前十，并在独立实验室的模型中排名第二。Krea 2 提供了一个全面的基线，并在保持竞争性能的同时实现了创意生成体验。 ## 数据 ### 数据策展原则 在详细描述我们的数据流水线之前，有必要明确对我们目标而言什么是“好的数据混合”。一个好的混合不仅仅包含“高质量”图像。考虑到我们构建表达性强、风格多样的模型的目标，多样性和广泛的领域覆盖至关重要。我们认为，传统基于模型的过滤方法，即使用美学评分和图像质量评估（IQA）模型，会引入隐含偏见。例如，这类方法可能会将模糊图像归类为低质量，尽管运动模糊或柔和可以是刻意的艺术选择。 此外，我们认为，只要描述准确描述了其图像，即使是不理想的图像在后续用例中也可能有所帮助：因为模型精确理解了不想要的行为，这样的样本可以在后续用于引导生成远离该分布。 基于这些原因，我们通过仅过滤掉以下内容来构建预训练数据集： 1. 重复样本和过度代表的概念。 2. VLM始终未能捕捉图像重要方面的样本。 3. 导致不想要的偏差和伪影的样本。 4. 视觉复杂度太高，无法在低分辨率下可靠建模的样本。 5. AI生成的样本。 这些条件塑造了一个具有广泛覆盖面的预训练数据集，同时避免了较差的文本到图像对齐和伪影。 重要的是，我们在预训练混合中完全不使用AI生成的图像。合成数据和蒸馏可以是获取模型能力的有效捷径。然而，我们发现即使是小比例的AI生成图像也会向模型的输出分布引入偏差，因为合成图像通常更容易学习，这实际上对模型质量施加了一个上限。因此，我们设计了内部分类器来过滤掉此类图像。 ### 描述生成 我们采用多阶段方法生成描述。首先，我们在每个目标图像上运行OCR模型以提取任何可见文本。在第二阶段，我们将OCR结果和任何可用的元数据（相机设置、已知实体等）提供给描述生成模型，该模型生成一个包含世界知识和提取文本的丰富描述。 通用描述生成流水线 *通用描述生成流水线* 一旦获得上下文丰富、长格式的自然语言描述，我们使用一个更便宜的LLM将其重新格式化为各种长度和格式，使模型接触一系列提示风格。根据经验，我们发现使用长提示进行训练提供了密集的监督，导致更快的收敛和更低的训练损失。然而，对于许多下游和应用用例，短中长度提示的性能仍然很重要。因此，我们主要训练长描述，同时确保模型在整个训练过程中接触到短和中长度提示。 我们的整体训练流水线和数据阶段 *我们的整体训练流水线和数据阶段* ### 预训练数据 预训练数据涵盖256px、512px和1024px分辨率阶段。逐步扩大分辨率形成了一种课程学习策略：我们将大部分FLOPs分配给低分辨率阶段，以高效地构建核心模型能力，然后随着训练分辨率的提高，赋予模型高保真生成能力。 低分辨率预训练是学习基本文本-图像对齐和结构的阶段。在此阶段，数据集规模达到数十亿张图像，因此我们严重依赖廉价的基于CPU的过滤器来移除低质量图像。这些过滤器包括从简单的损坏文件、分辨率和宽高比过滤器（移除不合格图像），到拉普拉斯过滤器（移除极端纹理和噪声模式的图像）。 例如，我们在预训练K2时遇到的一个问题是模型倾向于生成纯色背景和边框伪影。为了缓解这一问题，我们使用了RGB熵、白/黑像素比率、自定义启发式规则和内部分类器来过滤掉引发此行为的样本。 在构建内部分类器时，一个有效的策略是使用大型VLM为过滤任务设计任务特定的系统提示（例如，检测特定模式或伪影），生成伪标签数据集，然后训练一个基于DINOv3（https://arxiv.org/abs/2508.10104）或SigLIP-2（https://arxiv.org/abs/2502.14786）的小型分类器以大规模运行过滤器。任何在低分辨率阶段需要GPU计算资源的过滤模型，出于效率考虑，其参数都控制在10亿以下。 对于低分辨率阶段的去重，我们主要使用廉价的基于哈希的方法，结合md5、phash和colorhash，以最小的计算量移除重复图像。我们发现默认的8x8 phash不考虑颜色因素，且误报率较高；因此我们将12x12 phash与colorhash结合，以实现更稳健的去重。 随着训练分辨率的提高，我们引入了图像质量和美学过滤器。重要的是，这些质量分数仅用于移除质量极差的图像，而不是根据分数对图像进行过采样。我们还使用了图像复杂度分数和文本密度（来自OCR结果）来排除那些文本和内容无法在低分辨率下有意义表示的图像。随着训练的进行，我们会调整质量、复杂度和文本密度阈值。 除了传统的质量过滤器，我们还使用在预训练语料库样本上计算的SigLIP-2嵌入训练了一个稀疏自编码器（SAE）（https://arxiv.org/abs/2309.08600）。在训练SAE后，我们使用VLM根据其top-k激活样本对每个SAE特征进行注释。这些注释特征形成了一个无监督标记系统，我们从每张图像中提取主要的SAE特征。这个标记系统对于过滤明显的视觉伪影非常有用，而无需训练显式分类器。 ### 中期训练数据 与预训练阶段不同，中期训练显式地选择已知提供良好风格覆盖和特定视觉领域高质量图像的特定图像源。预训练是一个从通用池开始的、自下而上的过程，而中期训练数据则是自上而下策展的：首先选择领域和来源。中期训练是一个关键阶段，它平滑地连接了通用预训练分布和高质量SFT分布。为了提高分布的质量，我们引入了语义聚类，并使用基于检索的策略来确保世界知识覆盖。 在《自动数据策展用于自监督学习》（https://arxiv.org/abs/2405.15613）方法的基础上，我们使用FAISS执行层次化k-means聚类，然后进行采样，以便保留长尾视觉概念，同时避免浪费计算资源对头部概念进行过采样。在计算层次聚类后，我们让VLM检查每个聚类质心最近的图像，以命名并在适当时标记该聚类。在对标记的聚类进行人工审查后，我们删除了几个低质量或有问题的聚类。我们还通过语义去重进一步移除冗余数据，计算每个剩余叶子聚类内图像之间的SigLIP相似度。 图像生成模型的一个重要能力是忠实地表示用户可能仅通过名称引用的已知实体。某些实体，如运动员或演员，可能落入包含许多其他实体的语义聚类中，这可能导致在简单的层次采样下被丢弃。为了解决这个问题，我们使用Danker（https://github.com/athalhammer/danker）在英语维基百科上运行了PageRank，并按排名保留了前90%的文章。然后，我们根据其维基数据元数据过滤掉所有描述不可表示主题的文章，对于剩余的约500万个概念，我们在数据集中所有描述中执行全文搜索以评估覆盖率。在采样时，我们优先处理描述引用稀有概念的图像。最后，我们对得到的样本重复了覆盖率分析，以确认初始数据集中存在的概念没有被完全丢弃。 ### 监督微调数据 对于监督微调（SFT）阶段，我们使用一个小的、手工策展的数据集，专注于各个视觉领域。我们发现，一旦达到足够的数量，数据集的质量远比其规模重要。 ## 架构 在我们的架构消融实验中，我们发现将每个消融的目标分类到以下类别之一是有用的： - **稳定性：** 它是否使训练更稳定？是否减少了损失和梯度尖峰？ - **性能：** 它是否使模型收敛更快？如果是，这种趋势是否在更长时间范围内和更高分辨率下持续？ - **效率：** 它是否在不影响模型质量的情况下减少了参数量、FLOPs、内存或通信需求？ - **简洁性：** 我们能否在不影响其他类别的情况下使模型更简单？ 值得注意的是，我们的许多架构决策都受到它们在LLM领域采用的指导。选择在LLM生态系统中已经成熟的架构，使我们能够利用现有的内核和优化，即使对于扩散模型也是如此。 基于这些目标，我们从以下基线开始。 Krea 2 架构基线 **组件** | **基线** | **消融变量** | **最终组件** --- | --- | --- | --- 注意力 | 多头注意力 | GQA, MLA, 门控Sigmoid注意力 | GQA + 门控Sigmoid注意力 MLP | GeLU MLP | SwiGLU | SwiGLU 残差 | 标准残差 | 值残差, Laurel | 标准残差 **文本编码器** | T5-XXL编码器 | T5, Gemma, Qwen 2.5 VL, Qwen 3 VL, umT5 | Qwen 3 VL 调制 | 每块的MLP调制 | 带偏置的轻量调制 | 带偏置的轻量调制 自编码器 | FLUX AE | Qwen Image VAE, DC-AE, FLUX 2 AE, 内部VAE | Qwen Image VAE & FLUX 2 AE 块设计 | 单流Transformer块 | 混合流, 并行单流, | 单流Transformer块 归一化 | 层归一化, QK归一化 | RMSNorm, 零中心RMSNorm, Derf | 零中心RMSNorm, QK归一化 位置编码 | 3D轴向RoPE | 金门RoPE, MRoPE, 归一化RoPE, 部分RoPE | 3D轴向RoPE ### Transformer块 我们首先将GeLU MLP（https://arxiv.org/abs/1606.08415）替换为扩展因子为4x的SwiGLU（https://arxiv.org/abs/2002.05202）层，这已成为LLM架构中的事实标准模块。引入SwiGLU带来了持续的性能提升，因此我们在所有后续消融实验中采用了它。 在修改了MLP设计后，我们考虑了GQA（https://arxiv.org/abs/2305.13245）、MLA（https://arxiv.org/abs/2412.19437）和门控Sigmoid注意力（https://arxiv.org/abs/2505.06708）作为多头注意力基线的替代方案。我们发现GQA在提供改进的计算效率的同时，引入了最小的性能下降。我们还探索了MLA，并观察到相对于GQA的轻微增益，但未采用它，因为它引入了额外的计算开销。我们在使用MLA时采用了上下投影进行KV压缩，并且没有使用解耦的RoPE，因为di