摘要：当前的Large Vision Models多为直接从自然语言处理架构改编，但这些模型依赖效率较低的自回归技术，同时容易破坏视觉数据中至关重要的空间关系，从而限制了在高维视觉任务中的性能和效率。此外，多任务统一建模也是一个重要挑战，因为视觉任务的多样性和复杂性

研究方向｜视觉表征学习，多模态表征学习

编辑丨极市平台

论文链接：https://arxiv.org/abs/2411.11505

代码链接：https://github.com/DerrickWang005/LaVin-DiT

模型链接：https://huggingface.co/DerrickWang005/LaVin-DiT

当前的Large Vision Models多为直接从自然语言处理架构改编，但这些模型依赖效率较低的自回归技术，同时容易破坏视觉数据中至关重要的空间关系，从而限制了在高维视觉任务中的性能和效率。此外，多任务统一建模也是一个重要挑战，因为视觉任务的多样性和复杂性需要一个能够高效处理不同任务的通用生成框架。因此，本研究旨在解决如何高效统一建模多种视觉任务。

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作为一种高度可扩展、高效且多功能模型的地位，为大规模视觉基础模型指明了新路径。

直接处理像素空间中的视觉数据计算开销极大。为了解决这一问题，我们提出了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。

Diffusion Transformer（DiT）已经成为生成建模的有力方法。我们的Joint Diffusion Transformer（J-DiT）基于DiT，并引入了任务条件生成的改进。与原始DiT相比，一个关键区别是我们考虑了两种不同概念的潜在表示。Condition latent representation是未加噪的，而target latent representation则与高斯噪声叠加，导致两者可能具有不同的模式。为处理这种差异并提升任务特定与视觉信息之间的对齐，我们为条件和目标潜在表示分别构建了独立的分块嵌入层。每个嵌入层使用2×2的块大小，从而针对每种潜在类型调整表示。

如图2所示，采样的timestep t连同条件序列和目标序列一起输入到一系列扩散变换器层中。基于MM-DiT架构，我们引入了条件和目标特定的自适应RMSNorm（AdaRN），用于独立调节每种latent representation space。这是通过AdaRN层中条件和目标的timestep embedding实现的。

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 的算法流程，包括训练和推理的具体步骤，分别如下图所示：

为了统一多个计算机视觉任务，我们构建了一个大规模多任务数据集，涵盖室内和室外环境，并跨越真实和合成领域。该数据集包括约320万张独特图像和60万段独特视频，涉及20多种任务：

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

我们在图像和视频领域中覆盖的广泛计算机视觉任务上评估模型性能。按照既定协议，我们为每项任务报告标准指标。

定量分析： 为验证所提方法的有效性，我们在广泛的计算机视觉任务上进行了深入实验，并默认报告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%的

定性分析： 如图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模型的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）下，这一差距进一步扩大（20.1秒对47.2秒）。这一优势体现了扩散模型在视觉任务中的关键优势：不同于随输入规模增大而变得越来越耗时的自回归模型，扩散模型能够并行处理sequence，从而更高效地扩展。这种并行特性使得LaVin-DiT成为大规模视觉应用的理想选择。

上下文学习使得模型能够通过少量示例适应新任务，并且随着提供示例的增多，性能通常会有所提升。我们通过对十个下游任务评估任务上下文长度的影响，验证了这一点。如图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及类似模型将有能力处理仅依赖少量输入-目标对定义的任意（超出训练范围的）视觉任务。此外，研究自动选择最优任务上下文的方法可能为提升模型性能提供一种快速且高效的途径，确保模型能利用每项任务中最相关的示例。这些研究方向将推动开发更加稳健、适应性更强且高度通用的计算机视觉基础模型。

来源：极市平台