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

# 文本到图像模型训练速度提升 3.6 倍 来源：https://www.linum.ai/field-notes/jit-ddt 摘要 Linum v2 的瓶颈在于其巨大的注意力上下文窗口大小。一段 720p、5 秒的片段就需要多达 11 万 tokens。作为对比，大语言模型在预训练阶段，97% 的样本 token 数少于 8K (https://arxiv.org/pdf/2512.13961)。注意力机制的计算成本是二次方的，因此我们加速模型训练的最大杠杆就是裁剪上下文窗口。 大多数图像和视频生成系统都是潜在扩散模型。它们将压缩和生成拆分为独立训练的模块：变分自编码器和扩散变换器。最近，像 JiT 这样的像素空间模型被证明是一个有前景的替代方案。它将两个模型合二为一，并允许扩散模型专门为生成构建一个潜在空间，而不是依赖于为重建构建的空间。 当在我们（图像，描述）数据集上训练时，JiT 似乎难以生成精细细节。我们提出了一种新颖的编码器-解码器架构 (JiT-DDT) 来恢复这些细节，其训练效率比 LDM 对应版本高得多。与我们的 Linum v2 基准相比，JiT-DDT 训练一个文本到图像模型所需的 GPU 小时数减少了 3.6 倍，即使它生成的图像像素数是前者的 4 倍。 Linum v2（我们，之前）* 示例，256×256 **Linum v2（我们，之前）**·256×256 2.0B 潜在空间 DiT + VAE 256 个潜在 tokens * 仅图像检查点 JiT-DDT（我们，新）示例，512×512 **JiT-DDT（我们，新）**·512×512 2.5B 活跃像素空间 DiT 320 个像素 tokens = 64 编码器 + 256 解码器 GPU 小时数 见样本数 更多比较，请参阅附录 (https://www.linum.ai/field-notes/jit-ddt#appendix) 研究发布 JiT-DDT 代码和模型权重在 Apache 2.0 许可证下提供。我们希望通过与更广泛的社区分享我们的发现，鼓励其他人也探索更高效的训练方法。这应被视为研究工件，而非完整模型发布。敬请期待更多类似的研究检查点，我们正朝着 Linum v3 迈进。 ## 碰上 VAE 压缩瓶颈 几乎所有生成式图像和视频模型都是潜在扩散模型。它们有两个关键组成部分：用于压缩的变分自编码器和用于生成的扩散变换器。 直接处理原始像素成本过高（尤其是视频），因此我们首先需要找到一种方法将 RGB 像素减少为更少的 token 数量供 DiT 使用。这就是 VAE 的作用。它通过压缩和重建来训练。具体来说，它将我们的像素空间样本输入概率编码器，生成维度较小的 token，然后将这些潜在 token 输入概率解码器，最终还原回像素空间。 准备就绪 梯度（紫色）到达每个权重* 简化：实际上 KL 项约等于 0，之前我们训练 σ-VAE 时使用了 L1 重建损失加 LPIPS 和 GAN 损失；参见我们的 VAE 文章 (https://www.linum.ai/field-notes/vae-reconstruction-vs-generation)。 构建 LDM 时，你会单独训练 VAE，然后将其冻结（即 DiT 的梯度不会流向 VAE）。这样，在整个 DiT 训练过程中潜在空间保持静态。你运行 VAE 的编码器来嵌入数据，训练 DiT 在 VAE 的潜在空间中遍历，然后使用 VAE 的解码器将 DiT 生成的潜在 token 转换为像素空间。 准备就绪 VAE 冻结（虚线）· DiT 可训练（紫色）· 梯度在 DiT 处停止 · t=0 干净图像，t=1 纯高斯噪声 我们希望尽可能地从 VAE 中榨取 token 压缩，以抑制 DiT 中注意力机制的成本。但如果你调查一下流行的开源文本到图像模型，如 FLUX、Ideogram 和 Z-Image，你会发现它们都限制在 16×16 的 token 压缩。这与我们几年前关于图像-视频 VAE 的实验 (https://www.linum.ai/field-notes/vae-reconstruction-vs-generation) 一致。不幸的是，似乎存在一个经验上限，即在不降低重建质量的前提下，标准 CNN VAE 能实现的最大压缩量。 ## 通过统一模型解锁激进压缩 去年秋天，Tianhong Li 和 Kaiming He 发表了一篇论文 (JiT (https://arxiv.org/pdf/2511.13720))，通过完全抛弃 VAE 并将压缩任务整合到 DiT 本身，实现了 32×32 的 token 压缩。 16×16 像素，每个 3 通道（RGB） 准备就绪 示意图。在 JiT 的 512px 版本中，我们使用 32×32 的 patch，因此一张 512×512 的图像变为 256 个 token，每个 token 起始于 32·32·3 = 3,072 维；瓶颈层将其映射到 256 维。这种减少 token 数量的方法并非特别新颖。它早在半世纪前就为视觉变换器发明了，并且通常与 VAE 结合使用，在 token 序列进入 DiT 之前进一步压缩。在 Linum v2 中，我们的 VAE 提供了 8×8（高×宽）压缩和 16 维潜在向量。在 DiT 底层，我们应用 2×2 的分块，实现了 16×16 的 token 压缩和 64 维潜在向量。我们在 Linum v2 中使用了它，FLUX 等模型也是如此。 *那么，为什么之前没人尝试过？* 这感觉像是一顿免费的午餐。你获得了一种（可能）无损的方式来降低注意力成本，而且它极其简单。 在 2025 年初，像 VA-VAE (https://arxiv.org/pdf/2501.01423) 这样的论文表明 DiT 很难从高维输入中学习。有一些小技巧，例如在 VAE 训练期间使用外部模型作为正则化器（例如 DINOv3），这（可能）使像 FLUX-2 (https://bfl.ai/research/representation-comparison) 这样的模型能够从 64 个潜在维度跨越到 128 个潜在维度。但这些策略只是回避了 DiT 内部一个明确的学习性问题。激进的分块明确地将信息推入通道维度，因此会引发这种不稳定性。但事实证明，这并非架构固有的问题。相反，它是下游的 v-prediction，即过去几年大家用来训练扩散模型的 v-loss 流匹配目标导致的。 ### 流匹配快速回顾 在传统的 2022 年代的去噪扩散模型中，我们迭代地向样本添加噪声 (https://lilianweng.github.io/posts/2021-07-11-diffusion-models/#what-are-diffusion-models)，并训练神经网络来去除噪声。这样在推理时，我们可以使用神经网络在一系列步骤中将高斯噪声转换为来自我们数据分布的样本。这种公式有一系列问题（例如生成中的过度饱和、不稳定的学习、蒸馏崩溃），因此在过去几年中，该领域已经从它转向了流匹配。 在流匹配中，我们在数据分布中的每个样本和一个高斯噪声样本之间构建一条直线路径：噪声和样本之间的路径不一定是直线。但在实践中，我们都是这样做的。 在 t=0 时，我们恢复 x₀。在 t=1 时，我们得到 x₁ = ε，其中 ε ~ N(0, I)。我们在整篇文章中遵循 DDPM 的约定：x₀ 是数据，x₁ 是噪声。一些流匹配论文运行方向相反，将 ε 视为噪声，x₀ 视为数据。两种表述是等价的。然后我们训练一个网络来近似该路径上的速度： 我们称之为 v-prediction，v-loss，因为神经网络显式地预测速度，并且它在速度预测和真实条件速度场之间的均方误差上训练。 ### V-prediction 与维度灾难 如果你正在训练一个流匹配模型，你不一定需要训练你的神经网络来预测和回归速度。这三个术语是线性可重排的；因此你可以混合搭配 x̂，ε̂ 和 v̂ 在预测和回归目标中： 三个目标，每个目标线性依赖于另外两个 重排一个恒等式以填充每个非对角单元格 选择网络预测的内容（列）和损失度量的内容（行）。每个非对角单元格是上述三个恒等式之一，重排后将预测转换为损失目标。x-prediction 与 v-loss 是我们使用的单元格。表格基于 Li & He (2025)。在 JiT 中，Li 和 He 重新审视了自流匹配诞生以来该领域一直在做的 v-prediction，v-loss 决策。他们采用了一个玩具分布（螺旋线上的点），然后将这些点从二维投影到不同大小的高维空间。对于每个空间，他们分别用 x-prediction，ε-prediction 和 v-prediction 训练了流匹配模型，并发现只有 x-prediction 模型能够在高维度下准确生成来自螺旋线分布的样本。DiT 一直难以从高维输入中学习，原因在于维度灾难。 D=2、8、16 和 512 时 x-、epsilon- 和 v-prediction 的螺旋样本网格；只有 x-prediction 在 D=512 时保持有效 一个二维螺旋通过随机投影嵌入到 D 维空间中。随着 D 的增长，ε- 和 v-prediction 崩溃，而 x-prediction 持续恢复螺旋线。图 2 来自 Li & He (2025)。速度是 v = dx/dt。当我们做 v-prediction 时，我们的神经网络必须隐式学习信号 x₀ 和噪声 ε。噪声是一个覆盖整个 D 维空间的随机高斯。因此，随着 D 增大，拟合噪声（在速度项内）的问题变得指数级困难。这就是为什么激进的过去失败了，以及为什么 LDM 一直难以从高维 VAE 潜在向量中学习。随着通道维度增大，我们最终陷入 DiT 在努力学习高维高斯噪声的退化情况。 通过切换到 x-prediction，我们可以尝试绕过维度灾难。如果我们相信图像和视频自然存在于一个低维流形上，我们应该能够即使在高 D 的情况下也能有效地预测 x₀。 ε ~ N(0, I_D) n = **0** · 内在维度 = **512** x₀ ~ 图像流形 n = **0** · 内在维度 = **3** 1 个点 = 1 个样本 · D = 512 · 每个点绘制在其 512 维中的坐标 1、2 和 3 上 · 相同的球体，两侧 n 相同 · 流形是示意图，并非真实图像集 在高维空间 (D=512) 中，噪声是真正随机的。它分布在整个空间，不可压缩，无法用任何更小的维度数来描述（左）。图像本质上是低维的，因此即使在高维空间中，它们也聚集在一个小子空间中（右）。 通过预测 v，模型必须同时学习噪声 ε 和结构 x₀。噪声是两者中更难的部分，D 越大，模型用于拟合它的容量就越多。通过切换到 x-prediction，模型可以将其全部容量用于低维信号，即使 D 变得很大。 经验上这是有效的，如果你做 x-prediction，v-loss。网络预测 x̂₀。我们将该预测转换为速度，并与真实速度计算均方误差。由于 v = x₁ - x₀，这实际上是按 (1-t) 缩放的 x-loss：相同的目标，权重偏向小 t（低噪声，接近干净图像）。这反过来解锁了我们应用激进分块、扩大通道维度并将压缩问题推入 DiT 的能力。 ## 将 JiT 扩展用于文本到图像模型 当我们读到 JiT 时，我们非常兴奋地尝试它，因为它能够实现 32×32 的 token 压缩。但是，我们会说这是实现这种压缩程度的唯一方法，或者每个人都认为这是实现这种压缩程度的最佳方法，这是我们的疏忽。 LTX 已经在他们的视频模型中通过将 VAE 的解码器更改为显式去噪器（即他们使用流匹配目标微调 VAE 解码器）实现了这一点。最近，Minimax H3 通过将标准的约 8000 万-1.5 亿参数 CNN 解码器替换为一个 20 亿参数的变换器（大约相当于我们整个 Linum v2 模型的大小），在他们的 VAE 中实现了 32×32 的压缩。而且在像 ImageNet 这样的玩具基准测试上，LDM 仍然略优于像素空间模型。尽管如此，我们认为我们可以克服原始 JiT 论文中存在的一些局限性，并在生成式图像和视频中匹配 LDM 的性能。 在理念上，我们相信在大规模训练神经网络时，简单往往胜过复杂。像去年秋天的 E2E-VAE (https://arxiv.org/pdf/2504.10483) 这样的论文已经表明，允许你的 DiT（智能地）反向传播到 VAE 可以提高生成结果并显著加速收敛。 对我们来说，逻辑上讲得通的是，如果我们能专门为生成量身定制“潜在空间”，而不是依赖于为压缩构建的空间，我们就能得到更好的结果。而且我们还得到了一个额外好处，即拥有一个统一的模型，而不是两个分离的模型。 我们也认为 JiT 在 ImageNet 基准测试上的表现低估了它的潜力。JiT 可能能够实现比同等 VAE 更好的压缩，因为它可以利用文本提示提供的脚手架。 ### 文本到图像基准 当我们预训练 Linum v2 时，我们依赖于一个提供 16×16 token 压缩的 VAE + 分块堆栈。因此，我们在引入 180p 视频并扩展到 512px 分辨率之前，在 256px 分辨率下训练了约 6 亿样本。 对于我们的 JiT 基准，我们希望了解在相同图像-潜在 token 预算下的输出质量。这意味着我们在 512px 图像上训练，使用 32×32 的 token 压缩。 Linum v2（我们，之前）* 示例，256×256 **Linum v2（我们，之前）**·256×256 2.0B 潜在空间 DiT + VAE 256 个潜在 tokens * 仅图像检查点 JiT（宽）示例，512×512 **JiT（宽）**·512×512 2.0B 活跃像素空间 DiT 256 个像素 tokens GPU 小时数 见样本数 JiT 能够以相同的 token 数量（256 tokens）学习 512×512 的图像。收敛似乎快得多，但人脸看起来有磨皮感和过度饱和。 ### 我们的 JiT 设置 通过从 LDM 转向像素空间，我们从 v-prediction, v-loss 转向了 x-prediction, v-loss。但是，我们也对网络进行了一系列其他调整： 单流 DiT · ~2B 参数 x̂₀ 512×512×3 反分块输出头 线性层 2,944 → 3,072 最终 LayerNorm 变换器块 23×LayerNorm 1 Mod (scale, shift) 自注意力 23 头 × 128 · QK-Norm ×门控 + LayerNorm 2 Mod (scale, shift) SwiGLU FFN 2,944 → 7,936 → 2,944 ×门控 + sigmoid 门控（每头） 块 8 之后 像素重建损失 遮罩 20% 图像 tokens 变换器 2 层 余弦退火 vs DINOv3-L (x₀) v-加权 MSE ‖(x̂₀ - x₀)/max(t, 0.1)‖₂ LPIPS ×0.1 P-DINO ×0.01 23 层 · 维度 2,944 · 23 头 × 128 优化器 Muon 用于块矩阵 · AdamW 用于其余 上下文内文本条件 无交叉注意力 ‖[图像；文本] = 512 tokens AdaLN-Zero 正弦位置编码 (256) 共享头 SiLU · 线性 → 6·维度 ×2 位移 · ×2 缩放 ×2 门控 描述 Qwen3.5-4B 隐藏状态 7·15·27 → 7,680 MLP 投影 → 2,944 256 个文本 tokens x_t 512×512×3 32×32 分块 → 256 个 tokens 线性层 3,072 → 256 瓶颈 线性层 256 → 2,944 3D RoPE 256 个图像 tokens 灰色 = 冻结 · 虚线 = 损失连接 · 加法 · 乘法 1. **单流骨干网络** 与其交替使用自注意力（图像/视频）和交叉注意力（文本到图像/视频）的块，我们将视觉 tokens 和文本 tokens 连接成一个流，通过 DiT。这增加了每个块中的注意力序列长度，并增加了每个 token 的 FLOPs，但应该允许文本和图像 tokens 之间有显著更强的表达关系。之前（v2）：自注意力 → 交叉注意力 → FFN；Q：图像；K/V：文本。之后（v3）：自注意力