249个样本 · 2B Pyramid-JiT

点击任意图片查看完整分辨率及其提示词

对 Pyramid-JiT(P-JiT)进行压力测试

对我们来说,这次 2.1 万美元的训练运行只是通往下一代视频模型途中的概念验证。考虑到模型规模,我们对结果之好感到惊讶,所以决定以实验性发布的形式与全世界分享。

如果你之前没有关注过我们的工作,我们很想退后一步做个自我介绍。我们是一个两人的研究实验室,完全从零开始训练自己的生成式视频模型。最终目标是打造让动画创作大众化的创意工具。你可以在这里了解更多这里。

如今,生成式视频模型有两个主要缺陷:它们太慢/太贵,而且缺乏恰当的控制手段。创作是一个迭代的过程。你制作、你调整、你变得更好。有了这双重打击,视频模型(以它们目前的状况而言)对叙事来说实际上毫无用处。

今年早些时候我们着手构建 Linum v3 时,首要目标是大幅降低推理的成本门槛。如果我们的模型又快又便宜,那就是通往实用创意工具的第一步(众多步骤中的第一步)。

我们已经就这个方向上所采取的步骤写了大量内容,但长话短说:

  • 视频训练和推理的瓶颈是注意力,O(n²)。如果我们能更激进地压缩,就能缩小注意力的上下文窗口,从而大幅削减训练和推理的成本。
  • 通常,压缩由一个单独的专业模型来处理,比如变分自编码器,但这未必是最优的。VAE 在所能达到的压缩量上似乎存在一个经验性的上限。还有其他方法可以获得更多压缩,但它们需要对传统 VAE 做大幅改造(例如 Minimax 的 H3 中的 2B 解码器)。更重要的是,改进 VAE 在其重建任务上的表现可能会在下游扩散模型中产生更差的生成结果。
  • 如果完全去掉 VAE,你应该能够更激进地压缩并更快地学习。扩散模型是比 VAE 更强大的一类生成模型。我们应该能够用它们同时进行压缩和生成。扩散模型不应使用另一个模型的潜空间,而应学习一个专为其用例量身定制的潜空间。反过来,我们应该能够从这种压缩中榨取更多价值。

在上一篇博客中,我们迭代出了一种新架构,我们称之为 Pyramid-JiT(P-JiT)。它是一个像素空间模型,收敛所需的样本数比 Linum v2 少 11 倍,同时 token 缩减更激进 4 倍(即,我们用过去表示 256² 图像所需的相同 token 数来表示 512² 图像)。Linum v2 是一个相当标准的 Wan VAE + DiT 扩散模型。Linum v2 像大多数模型(例如 FLUX-2)一样使用了 16×16 的 token 缩减。P-JiT 实现了 32×32 的 token 缩减。这意味着注意力上下文窗口缩减了 4 倍。

关于我们如何走到这一步的细节,我们建议你看看我们之前的博客,关于JiT-DDT和P-JiT。但概括来说,我们为像素空间模型找到了一种噪声调度,能帮助模型更好地学习细粒度细节。而且,我们在 DiT 的主干上添加了读出层。

相邻的工作如REPA、RAE-v2和Self-Flow都指向 DiT 早期模块中的梯度较弱。通过在像素空间中操作,我们 DiT 的每个模块实际上都在预测最终图像的一个版本,因此我们可以通过预测最终图像的低分辨率版本来在网络更早处增强梯度信号。这有助于加速收敛。

该架构在我们最初的实验中表现良好,但我们想看看在训练更长时间后(即约 6 亿样本,而不是 5000 万至 1 亿样本)它的表现如何。作为背景,生成式视频模型几乎总是从生成式图像模型起步。你必须先学会名词(人、地点、事物),然后才能学会动词(动作、运动)。这算是在通往更大模型(以及视频)道路上的一种压力测试。

训练方案

尺寸读作 512² 阶段 / 1024² 阶段 · 花 = 一个示意性目标采样时 21.8 亿参数 · 紫色头仅用于训练,增加约 2.4 亿参数,训练时约 24.2 亿

文本和图像各自经过两个单模态的 "refiner" DiT 块,然后合并成一条由 22 个相同 DiT 块组成的流。文本表示来自一个冻结的 Qwen3.5-4B 模型,我们在第 7、15 和 27 层(它的三个全注意力层)读取它。我们将这些表示拼接并投影到文本 refiner 中。图像被切分成 32×32 的 patch,并压缩到 1024 维空间,然后再投影到图像 refiner 块中。这是 3:1 的压缩率。一个带 RGB 值的 32×32 patch 是 32×32×3 = 3,072 的表示,它会被投影到 1024 维。

在第 5 个块处,我们有一个 2 层的PixelREPA读出层,它将模型的潜空间与 DINOv3-L 对齐。我们遮蔽隐藏状态中 20% 的 token,并将其通过一个 2 层 transformer,以预测同一图像的 DINOv3-L 特征表示。我们在主干上第 10 和第 16 个块之后有两个线性读出头。如前所述,它们分别预测四分之一和一半分辨率的干净图像。总损失由各读出处的 MSE、PixelREPA 余弦距离、最终输出头处的 MSE,以及全分辨率图像上的感知损失(LPIPS和P-DINO)给出。

blue = AdaLN · s 在注意力/FFN 之前缩放输入,g 门控其残差写入隐藏大小 2,944 · 23 个头 × 128 · FFN 内部宽度在主干中为 5,376,在 refiner 中为 2,816

每个 DiT 块都通过低秩 AdaLN 以时间步为条件,遵循Z-Image。时间步嵌入经过一次降维投影到 256 维。图像 refiner 和主干共享一个降维投影;文本 refiner 有自己的降维投影。然后每个块都有自己的升维投影,输出一个缩放和一个门控,零初始化,因此每个块一开始都是恒等映射。

正如我们上一篇博客所述,我们在噪声调度上采用课程学习。前 80M 个样本我们集中在高噪声上,以便模型学习基本结构。然后,在剩余训练时间里我们拓宽噪声分布,使其能学习细粒度细节。详见附录 A。

H100 小时

512² ·01024² ·0

0

已见样本数

512² ·0M1024² ·0M

0M

在每个阶段分辨率下,对一组固定的 20,480 张方形图像计算验证损失。80M 处实时飙升到 0.157,是时间步调度放宽后的第一次评估。x 轴在 550M 之后放宽,以便让冷却阶段和 1024² 阶段清晰可读;柱状图按比例绘制。

我们在 100M 样本时开始 EMA,因此模型已经学到了对结构和细节相当不错的表征。当我们切换到 1024² 时,我们使用全新的优化器,并从 512² 的 EMA 权重继续训练。在 1024² 下,我们将噪声尺度从 2 加倍到 4,以补偿信息的冗余(见JiT,附录A)。更高的噪声水平以及在高分辨率下精确复现精细细节的需要,导致整体损失更高,即使图像质量在快速提升。我们在分辨率过渡中使用 RoPE 插值。我们保持一个与 512² 对齐的固定网格,并将 1024² 的 patch 索引乘以 ½,使它们对齐到 512² 的坐标轴。此外,我们在 2M 样本时启动 1024² 的 EMA,因此它在我们开始平均过程之前就已经适应了新的分辨率。

对于后训练,我们使用了有限差分流优化(FDFO)的一个变体,配合 JEV-Image和人类在环评审员。

预训练

$20,348

8,139H100 小时 × $2.50

512² 下 610M 样本,然后 1024² 下 50M

后训练

$690

276H100 小时 × $2.50

使用 JEV-Image 和人类在环评审员的 FDFO

总计

$21,038

$20,348+ $690·8,415H100 小时

在 64 块 H100 上预训练,包含重启按每 H100 小时 $2.50 计价

扩散模型的后训练仍处于早期阶段

直白地说,后训练在扩散模型中是一门非常年轻的学科,过去几个月我们才看到它初见成效的星星之火。

当然,在模型蒸馏和调度蒸馏方面已有大量工作,但就提升基础模型能力的后训练而言(就像过去 18 个月我们在 LLM 的编程能力上看到的那样),它在扩散模型上尚未真正开花结果。

DPO 已经存在一段时间了,但往往会严重过拟合到特定的审美偏好。Flow-GRPO是将 RLVR 风格框架引入扩散模型的最早阐述之一,而那只在去年秋天出现。自那以后,出现了几篇论文,比如DiffusionNFT以及今年 6 月的FDFO(在我们看来,这是迄今为止对扩散后训练最好的阐述)。

我们认为,适配这些技术之所以困难,归根结底是因为在扩散模型中我们无法像在自回归 LLM 中那样获得强似然估计。未来几周我们将更深入地探讨这些问题以及我们完整的后训练技术栈。

模型局限性

该模型未使用渲染文本的图像(没有信息图、PDF、书页、PowerPoint 幻灯片)进行训练,因此它无法书写文字或生成海报。正如我们在数据过滤博客中详述的那样,我们认为这些数据对我们动画工具的最终目标没有用处,因此这类文本密集的数据已从我们的数据集中完全过滤掉。

这是一个小模型(2B 参数的 DiT,4B 文本编码器),因此模型在几个方面表现吃力:全身、多实体(例如 6 辆车、人群)、多人交互以及世界知识。该模型偶尔能够正确生成这些图像。训练更大的模型(例如 7B 参数)并进一步后训练可能会缓解这些问题。

附录 A – 训练设置

训练进度

第 2 阶段从80M/660M样本开始

训练时间步在0.08到0.13· 悬停图表以移动

第 1 阶段 ·LogitNormal(0.8, 0.8)

<0.05%

第 2 阶段 ·LogitNormal(−0.2, 1.0)

3.0%

第 2 阶段 / 第 1 阶段

107×

设置512² 阶段1024² 阶段
样本数610.27M(297,985 步)49.97M(24,400 步)
全局批次大小2,0482,048
数据原生 512px 级尺寸的图像:368×512、512×512、512×384、512×336、640×360原生 1024px 级尺寸的图像:864×1216、1024×1024、1152×864、1248×832、1280×720
宽高比混合各占 20%,不缩放各占 20%,不缩放
加噪:噪声尺度随边长翻倍
时间步前 80M 样本(第 39,063 步)使用 LogitNormal(0.8, 0.8),之后使用 LogitNormal(−0.2, 1.0)全程使用 LogitNormal(−0.2, 1.0)
目标函数x-prediction、v-loss、权重在 t = 0.1 处截断 · 权重为 1 时的读出 · 最终头上对 t ≤ 0.7 使用 LPIPS 0.1 和 P-DINO 0.01 · PixelREPA 0.1
优化器Muon(动量 0.95,Nesterov)用于 2D 块矩阵 · AdamW(β = 0.9、0.999,ε = 1e-8)用于其余部分 · Muon 更新与 AdamW 做 RMS 匹配
学习率峰值 2e-4 · 在 1M 样本上线性预热 · 保持平坦至 550M · 在最后 60M 上按 1 − √ 衰减至 2.5e-5在 0.5M 样本上线性重新预热至 2.5e-5,然后保持平坦
权重衰减0.1(偏置上为 0)
梯度裁剪1.0(全局范数)
EMA半衰期 1,954 步(4M 样本),从 100M 样本(第 48,829 步)开始半衰期 977 步(2M 样本),从 2M 样本开始
尖峰感知 EMA当梯度范数超过其滚动中位数的 10 倍时,EMA 更新暂停 250 步
标题丢弃5%,用于无分类器引导
硬件64× H100(8 个节点)· bf16 · FSDP2 · 每 GPU 32 张图像相同,作为 4 个累积微批次,每批 8 个

有关架构和训练损失的更多细节,请参阅附录P-JiT。

本页所有文字均由我们撰写。我们使用 Claude Opus 5.5 帮助我们构建图片库和图表。

Huggingface 模型卡和 Github 仓库由 Claude Opus 5.5 自动撰写。我们将 Claude 指向我们内部的实验仓库,让它提取(并清理)必要的代码。

获取 Field Notes

关于从零构建生成式视频模型的技术深度解析,以及 Linum 新版本的更新。