专栏名称: 极市平台

极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台，为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯，行业动态，在线分享信息，线下活动等。网站: http://cvmart.net/

无需微调即可适应20多种视觉任务！爱诗科技与悉尼大学联合推出LaVin-DiT：大规模视觉扩散Transformer

极市平台 · 公众号 · · 2024-12-04 22:00

正文

↑ 点击蓝字关注极市平台

作者丨王兆卿博士悉尼大学，爱诗科技

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

编辑丨极市平台

极市导读

本文提出一种新型的大规模视觉扩散变换器LaVin-DiT，它通过结合空间-时间变分自编码器和联合扩散变换器，有效地处理高维视觉数据并保留空间关系，同时支持多任务学习，展现出在多个视觉基准测试中的优越性能。该模型通过上下文学习机制，无需微调即可适应各种任务，具有强大的泛化能力和高效的推理速度。 >> 加入极市CV技术交流群，走在计算机视觉的最前沿

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

方法

1. 问题设定

计算机视觉包含各种各样的任务，如目标检测和全景分割，这些任务通常会针对特定的input-target映射关系来设计专家模型。尽管这种专家模型在单任务中表现出色，但限制了模型在多样视觉数据中的适应性和扩展性。为克服这一限制，我们旨在设计一个条件生成框架，将多种视觉任务统一到单一的模型中。具体来说，给定查询（例如图像或视频），该框架生成对应的预测，以近似目标，条件是input-target pair集合。这些条件提供了对具体视觉任务的定义和引导，使模型能够根据提供的samples灵活适应不同的下游任务。形式化地，目标是建模条件分布。

2. 架构总览

如图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. 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。

4. J-DiT

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实现的。

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，以在残差连接中控制激活幅度。

6. 3D rotary position encoding

与LVM不同，我们认为将视觉数据建模为一维序列并不理想，因为一维位置嵌入难以捕捉精确的时空位置。相反，通过将多对input-target pair或video clip视为单个连续序列，我们可以使用三维旋转位置编码（3D RoPE）来简洁地表示时空关系。然后，视频中的每个位置可以用一个3D坐标表达。引入3D RoPE后，我们为各种视觉任务提供了统一且准确的时空位置表示。

7. Training & Inference procedure

LaVin-DiT 的算法流程，包括训练和推理的具体步骤，分别如下图所示：

实验

1. 训练数据

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

基于图像的任务：目标检测、实例分割、全景分割、姿态估计、边缘提取、深度估计、表面法线估计、图像修复（如去雨、去玻璃模糊和去运动模糊）、深度到图像生成以及法线到图像生成。
基于视频的任务：帧预测、视频深度估计、视频表面法线估计、视频光流估计、视频实例分割、深度到视频生成以及法线到视频生成。

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

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. 评估协议

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

4. 主要结果

定量分析： 为验证所提方法的有效性，我们在广泛的计算机视觉任务上进行了深入实验，并默认报告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^{\circ}$ 精度，超过了强大的专家模型StableNormal。这些结果展示了模型在复杂环境中准确估计表面方向的能力，特别是在增强现实和3D重建等需要几何理解的任务中。此外，在ImageNet-1K验证集的2500张随机图像上进行的图像修复任务中，我们的模型取得了1.65的FID，远优于LVM的4.05。

定性分析： 如图3所示，我们展示了模型在多种图像和视频任务中的定性结果。我们的模型能够持续跟随任务上下文，精准生成对应的预测结果。此外，在给定任务上下文的序列帧的基础上，我们的模型生成了后续12帧，展现了其在时间一致性和场景动态处理方面的出色能力。