训练文本到图像模型速度提高 3.6 倍

TL;DR Linum v2 的瓶颈在于其注意力上下文窗口的巨大尺寸。一个 720p、5 秒的剪辑花费了高达 11 万代币。从长远来看，法学硕士在 97% 的预训练中看到的样本少于 8K。注意力的成本是二次方的，因此我们加速模型训练的最大杠杆是修剪上下文窗口。大多数生成图像和视频系统都是潜在扩散模型（LDM）。他们将压缩和生成分成独立训练的模块：变分自动编码器（VAE）和 DiT（扩散变压器）。最近，像 JiT 这样的像素空间模型已被证明是一个有前途的替代方案。它将两个模型简化为一个模型，并允许扩散模型构建一个专门用于生成的潜在空间，而不是依赖于为重建而构建的潜在空间。当在我们的（图像，标题）数据集上进行训练时，JiT 似乎很难产生细粒度的细节。我们提出了一种新颖的编码器-解码器架构 (JiT-DDT)，它可以恢复此细节，并且比 LDM 对应物更有效地进行训练。与我们的 Linum v2 基线相比，JiT-DDT 训练文本到图像模型的 GPU 小时数减少了 3.6 倍，尽管它生成的图像具有 4 倍的像素。 Linum v2（我们的，以前的）* · 256×256 2.0B 潜在空间 DiT + VAE 256 潜在令牌 * 仅图像检查点 JiT-DDT（我们的，新） · 512×512 2.5B 活动像素空间 DiT 320 像素令牌 = 64 个编码器 + 256 个解码器 研究版本 JiT-DDT 代码和模型权重可在 Apache 下使用2.0 许可证。我们希望通过与更广泛的社区分享我们的发现，我们可以鼓励其他人也探索更有效的培训方法。这应该被视为研究工件，而不是完整的模型发布。请继续关注更多这样的研究检查点，在通往 Linum v3 的路上。撞上 VAE 压缩墙 几乎所有生成图像和视频模型都是潜在扩散模型 (LDM)。它们有两个关键组件：用于压缩的变分自动编码器 (VAE) 和用于生成的扩散变压器 (DiT)。以原始像素进行操作的成本太高（尤其是对于视频），因此我们首先需要找到一种方法将 RGB 像素减少为更少量的 DiT 令牌。这就是 VAE 的用武之地。它经过压缩和重建训练。具体来说，它通过概率编码器推送我们的像素空间样本，吐出维度标记，然后通过概率解码器推送这些潜在标记以返回像素空间。构建 LDM 时，您单独训练 VAE，然后冻结它（即没有梯度从 DiT 流到 VAE）。这样，潜在空间在整个 DiT 训练过程中保持静态。您运行 VAE 的编码器来嵌入数据，训练 DiT 来遍历 VAE 的潜在空间，然后使用 VAE 的解码器将 DiT 生成的潜在标记转换为像素空间。我们希望从 VAE 中获得尽可能多的令牌压缩，以便我们可以抑制 DiT 中的注意力成本。但如果您对流行的开源文本到图像模型（如 FLUX、Ideogram 和 Z-Image）进行调查，您会发现它们的上限都是 16×16 令牌减少。这与我们几年前对图像视频 VAE 的实验一致。不幸的是，在不降低重建效果的情况下，我们可以从标准 CNN VAE 中获得的压缩量似乎存在经验上限。使用统一模型解锁激进压缩去年秋天，Tianhong Li 和 Kaiming He 发表了一篇论文（JiT），通过完全抛弃 VAE 并将压缩任务推入 DiT 本身，实现了 32×32 令牌减少。 16×16 像素，3 个通道 (RGB)，每个切割成 4×4 像素补丁（16 个补丁），每个补丁独立标记：16 像素 × 3 通道成为一个 48 维标记，线性层 W ε ℝ 12×48 投影 48 暗度降至 12 ← 就绪 › 说明性。在 512px 的 JiT 中，我们使用 32×32 块，因此 512×512 图像变成 256 个标记，每个标记从 32·32·3 = 3,072 个暗点开始；瓶颈将其映射到 256。这种减少令牌数量的方法并不是特别新。它是五年前为视觉变换器 (ViT) 发明的，并且通常与 VAE 配合使用，以在令牌序列进入 DiT 之前进一步压缩它们。在 Linum v2 中，我们的 VAE 为我们提供了 8×8 (h×w) 压缩和 16 维潜在变量。在 DiT 的基础上，我们应用 2×2 补丁来获得 16×16 令牌压缩和 64 维潜在变量。我们在 Linum v2 中使用了它，FLUX 等模型也是如此。那么，为什么之前没有人尝试过呢？这感觉就像免费的午餐。你得到了一种（可能）无损的方式来降低注意力成本，而且非常简单。 2025 年初，VA-VAE 等论文证明 DiT 很难从高维输入中学习。有一些小技巧，例如在 VAE 训练期间使用外部模型作为正则化器（例如 DINOv3），这（可能）使 FLUX-2 等模型能够从 64 个潜在值实现飞跃