文章链接：https://arxiv.org/pdf/2411.11505
亮点直击
- 高效处理高维视觉数据 LaVin-DiT 引入空间-时间变分自编码器，降低计算需求，同时保留关键空间和时间特征，实现高效视觉数据建模。
- 创新的联合扩散Transformer 通过并行去噪生成视觉输出，提升处理效率，减少顺序依赖，并保留视觉任务所需的空间一致性。
- 强大的任务泛化能力 借助上下文学习机制，无需微调即可适应多种任务，展现卓越的多任务泛化性能。
- 显著性能与效率提升 在多项基准测试中优于现有模型，推理速度提升 1.7∼2.3 倍，同时表现出良好的扩展性和快速收敛性。
总结速览
解决的问题
当前的大型视觉模型大多直接从自然语言处理架构中改造而来，使用低效的自回归技术，这种方法破坏了视觉数据中至关重要的空间关系，并且在高维度视觉数据的生成任务中表现有限。
提出的方案
- 空间-时间变分自编码器 (spatial-temporal variational autoencoder)
编码高维视觉数据到连续的潜在空间，降低数据维度。 - 联合扩散Transformer (joint diffusion transformer)
通过逐步扩散生成视觉输出，优化生成建模能力。 - 基于上下文学习的统一多任务训练 (in-context learning)
利用输入-目标对作为任务上下文，引导扩散Transformer在潜在空间中完成特定任务。推理时，通过任务上下文和测试数据实现跨任务的泛化，无需微调。
应用的技术
- 变分自编码器 (VAE) 技术用于高效编码视觉数据。
- 扩散模型与Transformer的结合，用于逐步生成高质量视觉结果。
- 上下文学习（in-context learning）机制，支持多任务统一训练和推理。
- 可扩展性：模型参数从 0.1B 扩展至 3.4B，充分验证了其可扩展性能。
达到的效果
- 统一框架：首次在生成框架中实现大规模、多任务视觉建模。
- 高效推理：通过上下文学习，在无需微调的情况下实现任务泛化。
- 开源贡献：计划开放代码和模型，为后续研究提供支持。
方法
问题设置
计算机视觉包含一系列任务, 例如目标检测和全景分割, 通常通过为特定输入-目标映射设计的专用模型来完成。虽然在单一任务上表现出色，但这种专用性限制了模型在多任务或多样化视觉数据中的适应性和扩展性。为了解决这一局限性, 我们旨在设计一个条件生成框架, 将多个视觉任务统一到一个模型中。具体而言，给定一个查询 （例如，一张图像或一段视频），框架生成对应的预测 来逼近目标 , 条件是有一组输入-目标对 。这些条件对提供了任务定义和指导, 使模型能够根据提供的示例灵活适应不同任务。形式化地, 目标是对条件分布 进行建模。
框架概述
如下图 2(a) 所示, 提出的大型视觉Diffusion Transformer（LaVin-DiT）框架将空间-时间变分自编码器
（ST-VAE）与联合扩散Transformer（J-DiT）结合起来，实现多个视觉任务的统一。给定一个视觉任务, 例如全景分割, 首先采样一组输入-目标对作为任务定义。随后, 这组样本和其他视觉示例被输入到 ST-VAE 中, 并编码为潜在表示。接着, 编码后的表示被分块并展开为序列格式。样本集和输入视觉数据形成条件潜在表示 , 而目标通过添加随机高斯噪声变为噪声潜在表示 。随后, 和 被输入到联合扩散Transformer（J-DiT）中, 通过去噪 恢复清晰的潜在表示, 并保持在共享的潜在空间内。最后, 恢复的潜在表示通过 ST-VAE 解码器还原为原始像素空间中的目标。
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(c) 所示，条件和目标序列被线性投影后连接在一起，并通过双向注意力模块处理，从而使两者在各自的空间中操作时也能互相考虑。
为提高速度和内存效率，用分组查询注意力（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的可扩展性和性能，确立了它作为开发通用视觉模型的有前景框架。
参考文献
[1] LaVin-DiT: Large Vision Diffusion Transformer
