研究问题
论文链接：
代码链接：
模型链接：
GPT 和 LLaMA 这样的超大语言模型（LLMs）迅速获得了广泛关注并改变了该领域，展示了在统一框架内处理多种语言任务的强大能力。
这种将多种语言任务集成到单一大模型中的突破，激发了研发 Large Vision Models（LVMs）的动机。研发能够在多种视觉任务间泛化的 LVMs，代表了迈向更通用、可扩展且高效的视觉 AI 方法的希望之路。
然而，与 LLMs 相比，构建 LVMs 更为复杂，因为视觉数据本身具有多样性和高维特性，同时还需要处理任务中的尺度、视角和光照变化。
为了解决这些问题，近期的研究提出了一种基于序列建模的方法，将图像、视频和标注结果以统一的 “visual sentence” 形式表示，使模型能够从大规模数据集中预测连续的视觉标记，并完全脱离语言输入（如图 1（a）所示）。
尽管这一方法在多种视觉任务中取得了令人鼓舞的成果，但它面临两大主要挑战：一是自回归序列建模的效率限制，需要逐个标记预测，这对于高维视觉数据来说计算开销很大；二是将视觉数据转化为序列格式会破坏空间一致性，从而影响关键的空间依赖性保留。
本文中，我们提出了 Large Vision Diffusion Transformer（LaVin-DiT），旨在推动下一代 LVMs 的发展。LaVin-DiT 在计算效率上更优，同时有效地保留了视觉数据的空间关系，从而在多种视觉任务中取得了卓越性能（如图1（b）所示）。
技术上，为了解决视觉数据的高维特性，我们提出了 Spatial-temporal Variational Autoencoder（ST-VAE），能够将数据（例如图像和视频）编码到连续潜空间中，实现紧凑表示的同时保留关键的时空特征，从而减少计算需求并提高效率。
此外，在生成建模方面，我们改进了现有的扩散变换器，提出了一种 Joint Diffusion Transformer（J-DiT），通过并行去噪步骤生成视觉输出，有效减少了序列依赖性，同时保留了对视觉任务至关重要的空间一致性。
此外，为支持统一的多任务训练，我们加入了 In-context Learning 机制，利用 input-target pair 来引导 Diffusion Transformer 对齐特定任务的输出。在推理阶段，LaVin-DiT 利用 Task-specific context set 和 test sample 作为 query，无需微调即可适应各种任务。这种能力使得 LaVin-DiT 能够在多个复杂视觉应用中实现强大的泛化能力。
我们通过全面的实验验证了 LaVin-DiT 的优越性。结果表明，LaVin-DiT 在各种视觉基准测试中显著优于最强的 LVM 基线。例如，在 NYU-v2 深度估计中，AbsRel 降低了 24。此外，从 256×256 到 512×512 分辨率范围内，LaVin-DiT 推理速度比 LVM 快 1.7 至 2.3 倍。
不同模型规模的评估显示了 LaVin-DiT 在多个复杂视觉任务中的可扩展性和快速收敛性。最后，我们观察到，Task-specific context 长度的增加在各种任务中始终提升了性能。这些有希望的结果确立了 LaVin-DiT 作为一种高度可扩展、高效且多功能模型的地位，为大规模视觉基础模型指明了新路径。
方法
3.1 问题设定
如图 2（a）所示，所提出的 LaVin-DiT 框架结合了 ST-VAE 与 J-DiT，以统一多个视觉任务。对于一个视觉任务，例如全景分割，我们首先采样一组 input-target pair 作为任务定义。随后，该任务定义和其他视觉示例被输入到 ST-VAE 中并被编码为 latent representation。
接下来，这些 latent representation 被分块并展开为 sequence。任务定义与输入视觉数据形成 conditional latent presentation ，目标则被添加随机高斯噪声后生成一个 noisy latent representation 。随后， 和 被送入 J-DiT 通过去噪来恢复 clean latent representation。最后，这些 latent 通过 ST-VAE 解码器恢复到原始像素空间。
3.3 ST-VAE
直接处理像素空间中的视觉数据计算开销极大。为了解决这一问题，我们提出了 Spatial-temporal Variational Autoencoder（ST-VAE）。ST-VAE 能够高效地压缩时空信息，将其从像素空间编码为紧凑的潜在空间。
如图 2（b）所示，ST-VAE 使用因果 3D 卷积和反卷积来压缩和重建视觉数据。整体上，它包括一个编码器、一个解码器和一个潜在正则化层。这些组件被分为四个对称阶段，交替进行 2× 的降采样和上采样。前两个阶段作用于时空维度，而最后两个阶段仅作用于空间维度，实现了 4×8×8 的有效压缩，大幅降低了计算负担。
此外，我们应用了 Kullback-Leibler（KL）约束对高斯潜在空间进行正则化。
为防止未来信息泄露及其对时间序列预测的负面影响，我们在时间卷积空间的开始位置对所有位置进行填充。此外，为支持图像和视频处理，我们对输入视频的第一帧仅进行空间压缩以保持时间独立性，随后帧则同时压缩空间和时间维度。
ST-VAE 的编码器将输入压缩至低维潜在空间，解码器通过解码过程完成重建。ST-VAE 的训练分为两个阶段：首先单独在图像上训练，然后联合在图像和视频上训练。在每个阶段，我们通过 MSE Loss、Perceptual Loss 和 Adversarial Loss 的组合来优化 ST-VAE。
3.4 J-DiT
Diffusion Transformer（DiT）已经成为生成建模的有力方法。我们的 Joint Diffusion Transformer（J-DiT）基于 DiT，并引入了任务条件生成的改进。与原始 DiT 相比，一个关键区别是我们考虑了两种不同概念的潜在表示。Condition latent representation 是未加噪的，而 target latent representation 则与高斯噪声叠加，导致两者可能具有不同的模式。
为处理这种差异并提升任务特定与视觉信息之间的对齐，我们为条件和目标潜在表示分别构建了独立的分块嵌入层。每个嵌入层使用 2×2 的块大小，从而针对每种潜在类型调整表示。
如图 2 所示，采样的 timestep 连同条件序列和目标序列一起输入到一系列扩散变换器层中。基于 MM-DiT 架构，我们引入了条件和目标特定的自适应 RMSNorm（AdaRN），用于独立调节每种 latent representation space。这是通过 AdaRN 层中条件和目标的 timestep embedding 实现的。
3.5 Full-sequence joint attention
Full-sequence joint attention 是 J-DiT 的关键组件之一，能够同时处理条件序列和噪声目标序列以增强任务特定对齐。如图 2（c）所示，条件和目标序列被 linear project 和 concate，并通过 bi-direction attention 模块处理。这允许每个序列在保留自身特性时考虑另一个序列的信息。
为提高速度和内存效率，我们用 Group-query attention 代替了 Multi-head attention，将 query 分组以共享一组 key 和 value。这种方法减少了参数，同时保持了与标准的 Multi-head attention 相当的性能。
此外，为了在更大的模型和更长的序列中稳定训练，我们在 query-key dot-product 之前添加了 QK-Norm 以控制注意力熵的增长。根据以往研究，我们还在每个注意力层和前馈层后应用了 Sandwich Norm，以在残差连接中控制激活幅度。
与 LVM 不同，我们认为将视觉数据建模为一维序列并不理想，因为一维位置嵌入难以捕捉精确的时空位置。
相反，通过将多对 input-target pair 或 video clip 视为单个连续序列，我们可以使用三维旋转位置编码（3D RoPE）来简洁地表示时空关系。然后，视频中的每个位置可以用一个 3D 坐标表达。引入 3D RoPE 后，我们为各种视觉任务提供了统一且准确的时空位置表示。
LaVin-DiT 的算法流程，包括训练和推理的具体步骤，分别如下图所示：
实验
4.1 训练数据
为了统一多个计算机视觉任务，我们构建了一个大规模多任务数据集，涵盖室内和室外环境，并跨越真实和合成领域。该数据集包括约 320 万张独特图像和 60 万段独特视频，涉及 20 多种任务：
-
基于图像的任务：目标检测、实例分割、全景分割、姿态估计、边缘提取、深度估计、表面法线估计、图像修复（如去雨、去玻璃模糊和去运动模糊）、深度到图像生成以及法线到图像生成。
-
基于视频的任务：帧预测、视频深度估计、视频表面法线估计、视频光流估计、视频实例分割、深度到视频生成以及法线到视频生成。
为克服深度估计和表面法线估计中大规模标注的限制，我们利用 Depth-anything V2 和 Stable-Normal（turbo）分别在 ImageNet-1K 数据集上生成伪深度图和法线图。
4.2 实现细节
训练分两个阶段进行，逐步提高图像分辨率。在第一阶段，我们以 256×256 分辨率训练 100,000 步，使用 DeepSpeed ZeRO-2 优化器和梯度检查点技术以提升内存和计算效率。全局批量大小设置为 640，优化器采用 AdamW，学习率为 0.0001，动量参数 和 分别为 0.9 和 0.95，权重衰减为 0.01。该配置在无须额外正则化技术的情况下保证了稳定的训练。
在第二阶段，我们将分辨率提升至 512×512，继续训练 20,000 步，学习率调整为0.00005，其余超参数保持不变。这种两阶段策略有效提升了模型的分辨率适应能力，确保在多分辨率下的最佳性能。默认情况下，我们在推理时使用 20 timesteps（N=20）。所有实验均在 64 张 NVIDIA A100-80G 显卡上完成。
我们在图像和视频领域中覆盖的广泛计算机视觉任务上评估模型性能。按照既定协议，我们为每项任务报告标准指标。
定量分析：为验证所提方法的有效性，我们在广泛的计算机视觉任务上进行了深入实验，并默认报告 3.4B 参数模型的结果，详见表 1 和表 2。我们的模型在多项任务中持续优于现有基线，包括前景分割和单目标检测等具有挑战性的任务，展现出卓越的泛化能力和适应性。
在表 1 中，我们报告了前景分割和单目标检测在不同数据划分上的性能。LaVin-DiT 在所有划分上均显著优于基线方法。具体而言，在前景分割任务中，我们分别在四个划分上取得了 67.87%、75.80%、66.98% 和 66.90% 的 mIoU，比 LVM 和 MAE-VQGAN 等方法有大幅提升。
此外，在单目标检测中，我们在所有划分上均表现优异，特别是在 split-4 中，我们的 mIoU 达到 68.88%，较最佳基线 LVM 高出 19.96%。这些显著提升表明我们的模型在面对训练中未见过的任务时，能够高效分割和检测目标。
我们进一步评估了模型在图像着色任务中的表现，较低的 LPIPS 和 MSE 值表明更优的性能。根据表 1，我们的方法在着色任务中取得了 0.26 的 LPIPS 和 0.24 的 MSE，显著优于所有基线。这些结果反映了模型在从灰度图生成逼真自然颜色方面的能力，这在修复和艺术领域尤为重要。
为验证模型对 3D 场景几何结构的理解能力，我们在 NYU-v2 数据集上评估了深度估计和表面法线估计任务的性能（见表 2）。
在深度估计中，我们取得了 6.2 的 AbsRel 和 96.1% 的 精度，与 Marigold 和 DPT 等专家模型表现相当。在表面法线估计中，我们的模型实现了 15.901 的 MAE 和 58.382% 的 ＜11.25°C 精度，超过了强大的专家模型 StableNormal。
这些结果展示了模型在复杂环境中准确估计表面方向的能力，特别是在增强现实和 3D 重建等需要几何理解的任务中。此外，在 ImageNet-1K 验证集的 2500 张随机图像上进行的图像修复任务中，我们的模型取得了 1.65 的 FID，远优于 LVM 的 4.05。
定性分析：如图 3 所示，我们展示了模型在多种图像和视频任务中的定性结果。我们的模型能够持续跟随任务上下文，精准生成对应的预测结果。此外，在给定任务上下文的序列帧的基础上，我们的模型生成了后续 12 帧，展现了其在时间一致性和场景动态处理方面的出色能力。
4.5 扩展性
为研究 LaVin-DiT 的扩展性，我们针对三个模型规模（0.1B、1.0B 和 3.4B 参数）进行了实验，所有模型均训练 100,000 步。如图 4 所示，较大的模型持续实现更低的损失值。
此外，3.4B 模型收敛速度更快，在更少的训练步数内达到更低的损失值。这种快速收敛表明较大规模的模型能够更好地捕捉复杂数据模式，从而提升学习效率。
在下游任务中，模型规模对性能也有显著影响，特别是在着色和深度估计任务中。如图 5 所示，随着模型规模的增加，性能持续提升。在着色任务中，3.4B 模型的 MSE 为 0.273，显著优于 1.0B 和 0.1B 模型的 0.311 和 0.609。
同样，在深度估计任务中，3.4B 模型的 AbsRel 为 6.2，而 1.0B 和 0.1B 模型分别为 6.5 和 7.6。这些结果证明，较大规模的模型在多个任务上确实表现更优，表现了 LaVin-DiT 作为高性能可扩展框架的潜力。
4.6 推理延迟分析
如图 6 所示，我们对比了 LaVin-DiT 和 LVM（均为 7B 模型）在不同分辨率下的推理延迟，结果表明我们的模型始终更高效。在 256 分辨率下，LaVin-DiT 每个样本仅需 4.67 秒，而 LVM 需 8.1 秒；在更高分辨率（如 512）下，这一差距进一步扩大（20.1 秒对 47.2 秒）。
这一优势体现了扩散模型在视觉任务中的关键优势：不同于随输入规模增大而变得越来越耗时的自回归模型，扩散模型能够并行处理 sequence，从而更高效地扩展。这种并行特性使得 LaVin-DiT 成为大规模视觉应用的理想选择。
4.7 任务上下文长度的影响
上下文学习使得模型能够通过少量示例适应新任务，并且随着提供示例的增多，性能通常会有所提升。我们通过对十个下游任务评估任务上下文长度的影响，验证了这一点。
如图 7 所示，模型随着任务上下文的增加持续受益，性能显著提升。例如，在深度到图像生成任务中，较长的上下文带来了更低的 FID；在去运动模糊任务中，较长的上下文提升了 PSNR。这些结果表明 LaVin-DiT 能够有效利用扩展的任务上下文，进一步增强任务适应性和准确性。
结论
我们提出了 LaVin-DiT，一个可扩展的统一计算机视觉基础模型，它结合了 Spatial-temporal Variation Autoencoder（ST-VAE）和 Joint Diffusion Transformer（J-DiT），能够高效处理高维视觉数据，同时保留空间和视觉一致性。
通过上下文学习，LaVin-DiT 无需微调即可适应广泛的任务，展现出卓越的通用性和适应性。大量实验验证了 LaVin-DiT 在可扩展性和性能方面的优势，确立了其作为通用视觉模型框架的潜力。
局限性
尽管 LaVin-DiT 具有诸多优点，但仍受限于当前大规模训练数据、多样化任务标注以及计算资源的限制，特别是在与大规模语言模型的对比中。
尽管我们的模型在已知任务和相关未知任务上表现出色，但当任务定义显著偏离训练分布时，模型的泛化能力仍然较弱。这一局限性凸显了开发能够仅通过任务上下文有效泛化到全新任务的视觉模型所面临的核心挑战。
未来工作
未来研究应进一步探索在模型容量、数据集多样性和任务复杂性方面扩展 LaVin-DiT，以突破视觉泛化的界限。我们预计，随着这些因素的扩展，LaVin-DiT 及类似模型将有能力处理仅依赖少量输入-目标对定义的任意（超出训练范围的）视觉任务。
此外，研究自动选择最优任务上下文的方法可能为提升模型性能提供一种快速且高效的途径，确保模型能利用每项任务中最相关的示例。这些研究方向将推动开发更加稳健、适应性更强且高度通用的计算机视觉基础模型。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
（文：PaperWeekly）