可扩展多任务视觉基础模型 LaVin-DiT：融合时空 VAE 与 DiT

LaVin-DiT 模块

ST-VAE

直接在原始像素空间处理视觉数据计算量巨大。为了解决这个问题，本文提出使用空间 - 时间变分自编码器（ST-VAE）。ST-VAE 能有效压缩空间和时间信息，将其从像素空间编码到紧凑的潜在空间。如上图 2(b) 所示，ST-VAE 使用因果 3D 卷积和反卷积来压缩和重建视觉数据。它总体上包含一个编码器、一个解码器和一个潜在正则化层。这些组件被结构化为四个对称阶段，交替进行下采样和上采样。前两个阶段作用于空间和时间维度，而最后阶段仅作用于空间维度，实现了的压缩，有效减少了计算量。此外，应用了 Kullback-Leibler (KL) 约束来正则化高斯潜在空间。

为防止未来信息泄漏及其对时间预测的不利影响，在时间卷积空间开始时对所有位置进行填充。此外，为了支持图像和视频处理，将输入视频的第一帧独立对待，仅对其进行空间压缩以保持时间独立性。后续帧同时在空间和时间维度上进行压缩。

ST-VAE 的编码器将输入压缩到低维潜在空间，通过解码过程实现重建。ST-VAE 的训练分为两个阶段：首先单独对图像进行训练，然后联合图像和视频进行训练。在每个阶段，我们使用均方误差、感知损失和对抗性损失的组合来优化模型。

J-DiT

Diffusion Transformer（DiT）已成为生成建模的强大方法。联合 Diffusion Transformer（J-DiT）基于 DiT 构建，并引入修改以支持任务条件生成。与原始 DiT 的关键区别在于我们考虑了两种概念上不同的潜在表示。条件潜在表示是干净的，而目标潜在表示被高斯噪声扰动，这导致两者可能具有不同的值范围。为了处理这种差异并改善任务特定信息与视觉信息之间的对齐，为条件和目标潜在表示分别构建了独立的分块嵌入层。每个嵌入层使用 的块大小，能够针对每种潜在类型定制表示。如前面图 2 所示，采样的时间步以及条件和目标序列被输入到一系列扩散 Transformer 层中。基于 MM-DiT 的架构，在 AdaRN 层中引入了针对条件和目标的自适应 RMS 正则化（AdaRN），分别调节每种表示空间。这通过为条件和目标分别生成的时间步嵌入实现。

全序列联合注意力
全序列联合注意力是 Transformer 层的核心，能够将条件和噪声目标序列一起处理，以增强任务特定的对齐。如前面图 2© 所示，条件和目标序列被线性投影后连接在一起，并通过双向注意力模块处理，从而使两者在各自的空间中操作时也能互相考虑。

为提高速度和内存效率，用分组查询注意力（grouped-query attention）替代多头注意力。这种方法将查询头分组，以共享一组键值头，从而减少参数量，同时保持表达能力，与标准多头注意力性能接近。此外，为了在处理更大的模型和更长的序列时稳定训练，我们在查询 - 键点积之前添加了 QK-Norm 来控制注意力熵的增长。按照 [53] 的方法，还在每个注意力和前馈层后应用了夹心正则化（sandwich normalization），以在残差连接中保持激活幅度。

3D 旋转位置编码与 [4] 不同，将视觉数据建模为一维序列是不理想的，因为一维位置嵌入在捕捉精确的时空位置方面具有局限性。相反，通过将多组图像 - 标注对或视频片段视为一个连续序列，可以使用三维旋转位置编码（3D RoPE）简洁地表示时空关系。这样，视频中的每个位置都可以用三维坐标表示。引入 3D RoPE 后，为各种视觉任务提供了一种统一且精确的时空位置编码表示。

J-DiT 的训练过程在潜在空间中使用流匹配（flow matching）来训练 J-DiT。具体而言，给定表示和噪声，流匹配定义了一种基于线性插值的前向过程：

其中时间步。该前向过程在沿着线性路径的方向上诱导了一个时间相关的速度场。该速度场定义了一个常微分方程（ODE）：

使用由参数参数化的 J-DiT 来预测将噪声转换为干净潜在表示的速度场。流匹配的训练目标是直接回归目标速度场，从而得到条件流匹配（Conditional Flow Matching, CFM）损失：

J-DiT 的生成过程完成 J-DiT 训练后，通过从噪声分布向表示分布积分，生成新的表示。具体而言，从时的噪声开始，向反向积分学习到的 J-DiT，以获得表示。例如，使用欧拉方法，将时间区间离散为个步骤，时间步长为负值，以表示时间上的反向积分。在每一步，时间和生成的表示更新如下：

其中，，，。通过反复应用这些更新步骤，我们可以获得用于 ST-VAE 解码过程的新表示。

LaVin-DiT 推理

完成 LaVin-DiT 的训练后，该模型具有多功能性，可以应用于多种下游任务。具体而言，对于任何选定任务，在给定查询（例如，一张图像或一段视频）时，我们随机采样一组定义任务的输入 - 目标对。这些对与视觉输入以及高斯噪声组件一起被输入到联合 Diffusion Transformer（J-DiT）。在 J-DiT 中，这些元素被处理以生成一个潜在表示。最终，该潜在表示被传递到 ST-VAE 解码器中，将其转换为原始像素空间，从而生成所需的预测结果。

实验

设置

训练数据
为了统一多个计算机视觉任务，构建了一个大规模的多任务数据集，涵盖了室内和室外环境，跨越了现实世界和合成领域。该数据集包含约 320 万张独特图像和 60 万段独特视频，涵盖了超过 20 个任务：

基于图像的任务：物体检测、实例分割、全景分割、姿态估计、边缘提取、深度估计、表面法线估计、图像修复、图像着色、图像恢复任务（如去雨、去玻璃模糊、去运动模糊）、深度到图像生成、法线到图像生成。

基于视频的任务：帧预测、视频深度估计、视频表面法线估计、视频光流估计、视频实例分割、深度到视频生成、法线到视频生成。

为了克服深度和表面法线估计的大规模标注限制，利用 Depth-anything V2 和 Stable-Normal（turbo）在 ImageNet-1K 上生成了伪深度图和法线图。

实现细节
分两个阶段进行训练，逐步增加图像分辨率。在第一阶段，以 256×256 分辨率训练 100,000 步，利用 DeepSpeed ZeRO-2 优化和梯度检查点技术来管理内存和计算效率。使用 640 的全局批量大小，并采用 AdamW 优化器，学习率为 0.0001，betas 设置为 0.9 和 0.95，权重衰减为 0.01。该设置提供了稳定的训练，无需热身或额外的正则化技术。在第二阶段，将分辨率上调到 512×512，并继续训练 20,000 步，同时将学习率调整为 0.00005，其他超参数保持不变。该两阶段策略能够高效扩展，确保不同分辨率下的最佳性能。默认情况下，在推理过程中使用 20 个时间步（N=20）。所有实验均在 64 个 NVIDIA A100-80G GPU 上进行。

评估协议
在一系列涵盖图像和视频领域的计算机视觉任务上评估我们的模型。按照已建立的协议，报告每个任务的标准评估指标。

主要结果

定量分析
为了评估本文提出方法的有效性，在一系列计算机视觉任务上进行了广泛的实验，并报告了 3.4B 模型的结果，如下表 1 和表 2 所示。本文的方法在多个任务上始终优于现有基准方法，包括一些具有挑战性的案例，如未见过的前景分割和单物体检测，展示了模型在各种场景中的卓越泛化能力和适应性。除非另有说明，否则告的是 LaVin-DiT (3.4B) 的性能。

如上表 1 所示，报告了在不同分割上的前景分割和单物体检测任务的性能。LaVin-DiT 在所有分割上都取得了显著的提升。在前景分割任务中，在四个分割上分别达到了 67.87%、75.80%、66.98% 和 66.90% 的 mIoU，始终大幅超过了 LVM 和 MAE-VQGAN 等先前方法。

此外，对于单物体检测，本文的模型在所有分割上都表现出色，特别是在第 4 分割中，我们获得了 68.88% 的 mIoU，相较于表现最好的基准 LVM，提升了 19.96%。这些显著的提升凸显了我们模型在多种场景中有效地进行物体分割和检测的能力，尤其是在面对训练中未见过的任务时。根据先前的工作 [4, 7]，我们还在着色任务中评估了我们的模型，其中较低的 LPIPS 和 MSE 值表示更优的性能。

如前面表 1 所示，本文的方法取得了 0.26 的 LPIPS 和 0.24 的 MSE，显著优于所有基准方法。这些结果凸显了我们模型从灰度图像中生成逼真自然颜色的能力，这对于恢复和艺术领域至关重要。

为了验证本文模型理解 3D 场景几何结构的能力，在 NYU-v2 深度估计和表面法线估计任务上进行了评估，如表 2 所示。由于 Bai 等人未在他们的论文中报告相关结果，我们使用他们官方的 7B 模型进行了评估。对于深度估计，本文的模型达到了 6.2 的 AbsRel 和 96.1% 的δ1 阈值准确度，表现与 Marigold 和 DPT 等专家模型相当。在表面法线估计任务中，本文的方法取得了 15.901 的 MAE 和 58.382 的<11.25°阈值准确度，超越了强大的专家模型 StableNormal。这一表现突显了我们模型在准确估计表面方向方面的能力，增强了其在需要精确几何理解的任务中的应用，如增强现实和 3D 重建。这些结果反映了我们模型在复杂环境中精确理解 3D 场景几何结构的能力，这对现实世界的应用如 3D 场景重建和空间感知至关重要。

此外，还在图像修复任务上将 LaVin-DiT 与 LVM 进行了比较。通过使用从 ImageNet-1K 验证集中随机选择的 2,500 张图像，本文的模型在 FID 上取得了 1.65 的成绩，相较于 LVM 获得的 4.05，取得了显著的改进。

定性分析如下图 3 所示，展示了在多种基于图像和基于视频的任务中的定性结果。我们的模型始终遵循任务上下文，并准确地生成相应的预测。此外，给定带有任务上下文的连续帧，模型能够生成接下来的 12 帧预测，展示了其有效处理时间一致性和场景动态的能力。

可扩展性

为了研究所提出的 LaVin-DiT 的可扩展性，进行了三个模型大小（0.1B、1.0B 和 3.4B 参数）的实验，训练了这三个模型 100,000 步。下图 4 展示了训练损失曲线，结果显示，较大的模型在训练过程中始终能够取得较低的损失值。此外，3.4B 模型收敛速度更快，在更少的训练步骤内达到了较小的损失值。这一加速收敛表明，较大的模型更适合捕捉复杂的数据模式，从而提高了学习效率。观察到的训练动态强调了在复杂视觉任务中扩展模型容量的优势，较大的模型能够更有效地捕捉多样的数据特征。

除了训练动态外，模型大小对下游任务性能也有显著影响。这在着色和深度估计任务中得到了体现，这些任务因其在捕捉颜色保真度和空间结构上的不同要求而被选中。如下图 5 所示，随着模型规模的增加，性能持续提高。具体而言，在着色任务中，3.4B 模型的 MSE 为 0.273，显著优于 1.0B 和 0.1B 模型的 MSE，分别为 0.311 和 0.609。同样，在深度估计中，3.4B 模型的 AbsRel 为 6.2，而 1.0B 和 0.1B 模型分别为 6.5 和 7.6。这些结果表明，较大的模型在多个任务中确实能提供更强的性能，证明了 LaVin-DiT 是一个可扩展且适应性强的高性能视觉应用框架。

推理延迟分析

如下图 6 所示，比较了 LaVin-DiT 和 LVM（两者均为 7B 模型）在不同分辨率下的推理延迟，结果表明本文的方法始终更高效。在 256 分辨率下，LaVin-DiT 每个样本仅需 4.67 秒，而 LVM 需要 8.1 秒；在更高分辨率下（例如 512 分辨率时，LaVin-DiT 为 20.1 秒，LVM 为 47.2 秒），这一效率差距进一步扩大。这个差异突显了扩散模型在视觉任务中的一个关键优势：与逐步处理 tokens 的自回归模型不同，扩散模型可以并行处理 tokens，从而在面对较大输入时能够更有效地扩展。这种并行性使得我们的 LaVin-DiT 在大规模视觉应用中成为更合适的选择。

任务上下文长度的影响

上下文学习使得模型能够通过少量示例适应新任务，随着提供更多示例，性能通常会提高。我们通过评估任务上下文长度对十个下游任务的影响来进行研究。如下图 7 所示，模型在任务上下文变长的情况下，性能持续提升，取得了显著的性能增益。例如，随着更多输入 - 目标对的引入，LaVin-DiT 在深度到图像生成任务中取得了更低的 FID，并在去除运动模糊任务中实现了更高的 PSNR。这些结果表明，LaVin-DiT 能够有效利用扩展的任务上下文，突出其利用额外信息来增强任务适应性和准确性的能力。

结论

本文提出了 LaVin-DiT，一个可扩展且统一的计算机视觉基础模型，集成了空间 - 时间变分自编码器和扩散 Transformer，能够高效处理高维视觉数据，同时保持空间和视觉一致性。通过上下文学习，LaVin-DiT 能够在不进行微调的情况下有效适应广泛的任务，展现了其显著的多功能性和适应性。广泛的实验验证了 LaVin-DiT 的可扩展性和性能，确立了它作为开发通用视觉模型的有前景框架。