【他山之石】CVPR-悉尼大学提出LaVin-DiT：扩散Transformer+时空VAE，20+ 视觉任务统一SOTA框架
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注！
来源：PaperEveryday
仅用于学术分享，若侵权请联系删除
引入时空变分自编码器：为处理视觉数据的高维性，LaVin-DiT纳入时空变分自编码器（ST-VAE）。它能将图像和视频数据编码到连续潜在空间，在保留关键时空特征的同时实现紧凑表示，降低计算需求，提升效率，且不影响模型捕捉复杂模式的能力。
提出联合扩散Transformer：在生成建模方面，LaVin-DiT对现有扩散Transformer改进，提出具有全序列联合注意力的联合扩散Transformer（J-DiT）。该模块通过并行去噪步骤合成视觉输出，减少序列依赖性，提高处理效率，还能维持视觉任务所需的空间连贯性。
采用上下文学习：为支持统一的多任务训练，LaVin-DiT引入上下文学习。输入-目标对作为任务上下文，引导扩散Transformer在潜在空间中使输出与特定任务对齐。
使用3D旋转位置编码：LaVin-DiT采用3D旋转位置编码（3DRoPE），将视觉数据视为连续序列，用3D坐标表示位置，为各种视觉任务提供统一且准确的时空位置编码，克服了1D位置嵌入在捕捉时空位置上的局限。
计算机视觉包括一系列任务，如目标检测和全景分割，这些任务通常由为特定输入-目标映射设计的专用模型处理。虽然这种专业化在单个任务中有效，但限制了模型在多个任务或不同视觉数据上的适应性和可扩展性。为了克服这一限制，作者旨在设计一个条件生成框架，将多个视觉任务统一在一个紧密结合的模型中。具体来说，给定一个查询（例如，一幅图像或一段视频），该框架在一组输入-目标对的条件下，生成相应的预测，以逼近目标。这些条件对提供了任务定义和指导，使模型能够根据提供的示例灵活地适应不同的任务。形式上，目标是对条件分布进行建模。
如图2（a）所示，所提出的大型视觉扩散Transformer（LaVin-DiT）框架将时空变分自编码器（ST-VAE）与联合扩散Transformer相结合，以统一多个视觉任务。给定一个视觉任务，例如全景分割，作者首先采样一组输入-目标对作为任务定义。然后，将这组对和其他视觉示例输入到ST-VAE中，编码为潜在表示。随后，将编码后的表示进行分块并展开为序列格式。这组对和输入视觉数据形成条件潜在表示，而目标则用随机高斯噪声进行扰动，得到噪声潜在表示。和都被输入到联合扩散Transformer（J-DiT）中，J-DiT对进行去噪，在共享潜在空间中恢复干净的潜在表示。最后，恢复的潜在表示通过ST-VAE解码器，在原始像素空间中重建目标。下面作者详细介绍ST-VAE和J-DiT的技术细节。
在原始像素空间中处理视觉数据计算成本很高。为了解决这个问题，作者提出使用时空变分自编码器（ST-VAE）。ST-VAE能够有效地压缩空间和时间信息，并将其从像素空间编码到紧凑的潜在空间中。如图2（b）所示，ST-VAE使用因果3D卷积和反卷积来压缩和重建视觉数据。它总体包括一个编码器、一个解码器和一个潜在正则化层。这些组件被结构化为四个对称阶段，交替进行2倍下采样和上采样。前两个阶段在空间和时间维度上都进行操作，而最后一个阶段仅在空间维度上起作用，实现了4×8×8的有效压缩，降低了计算负载。此外，作者应用Kullback-Leibler（KL）约束来正则化高斯潜在空间。
为了防止未来信息泄漏及其对时间预测的不利影响，作者在时间卷积空间的起始位置对所有位置进行填充。另外，为了支持图像和视频处理，作者将输入视频的第一帧单独处理，仅在空间上进行压缩以保持时间独立性。后续帧则在空间和时间维度上都进行压缩。ST-VAE的编码器将输入压缩到低维潜在空间，通过解码过程实现重建。ST-VAE的训练分两个阶段进行：作者首先仅在图像上进行训练，然后在图像和视频上联合训练。在每个阶段，作者使用均方误差、感知损失和对抗损失的组合来优化模型。
扩散Transformer（DiT）已成为一种强大的生成建模方法。作者的联合扩散Transformer（J-DiT）基于DiT构建，但进行了修改以支持任务条件生成。与原始DiT的一个关键区别是，作者考虑了两种概念上不同的潜在表示。条件潜在表示是干净的，而目标潜在表示受到高斯噪声的扰动，导致两者可能具有不同的值范围。为了处理这种差异并改善任务特定信息和视觉信息之间的对齐，作者为条件和目标潜在构建了单独的补丁嵌入。每个嵌入层使用2×2的补丁大小，以便为每种潜在类型定制表示。如图2所示，采样的时间步，以及条件和目标序列，被输入到一系列扩散Transformer层中。基于MM-DiT架构，作者引入了条件和目标特定的自适应RMS归一化（AdaRN），以独立调制每个表示空间。这通过在AdaRN层中为条件和目标设置不同的时间步嵌入来实现。
全序列联合注意力：全序列联合注意力是作者Transformer层的关键，它同时处理条件和噪声目标序列，以增强任务特定的对齐。如图2（c）所示，条件和目标序列被线性投影、连接，然后由双向注意力模块处理，使它们能够在各自的空间中操作，同时考虑对方。为了提高速度和内存效率，作者用分组查询注意力（grouped-queryattention）代替多头注意力，分组查询注意力将查询头分组，共享一组键值头。这种方法在减少参数的同时保留了表达能力，性能与标准多头注意力非常接近。此外，为了在更大的模型和更长的序列中稳定训练，作者在查询-键点积之前添加QK-Norm，以控制注意力熵的增长。遵循相关工作，作者还在每个注意力层和前馈神经网络（FFN）层之后应用三明治归一化，以在残差连接中保持激活幅度。
3D旋转位置编码：与一些方法不同，作者认为将视觉数据建模为一维序列不是最优选择，因为一维位置嵌入在捕捉精确的时空位置方面存在局限性。相反，通过将多个图像-注释对或视频片段视为单个连续序列，作者可以使用3D旋转位置编码（3DRoPE）简洁地表示时空关系。这样，视频中的每个位置都可以用一个3D坐标表示。通过引入3DRoPE，作者为各种视觉任务提供了统一且准确的位置编码时空表示。
J-DiT的训练过程：作者在潜在空间中使用流匹配（flowmatching）方法训练J-DiT。具体来说，给定一个表示和噪声，流匹配定义了一个基于线性插值的前向过程：，其中时间步。这个前向过程诱导了一个随时间变化的速度场，它沿着线性路径朝着的方向驱动流。速度场定义了一个常微分方程（ODE）：。作者使用由参数化的J-DiT来预测将噪声转换为干净潜在表示的速度场。流匹配的训练目标是直接回归目标速度场，得到条件流匹配（CFM）损失：
J-DiT的生成过程：在J-DiT训练完成后，作者使用它从噪声分布向表示分布进行积分，以生成新的表示。具体来说，从时的噪声开始，作者将学习到的J-DiT向后积分到，以获得表示。例如，使用欧拉方法，作者将时间间隔离散化为步，步长为，表示时间上的向后积分。在每一步，作者更新时间和生成的表示如下：
其中，，，。通过迭代应用这些更新，作者获得一个新的表示，用于后续ST-VAE的解码过程。
在完成LaVin-DiT的训练后，该模型变得通用，可应用于一系列下游任务。具体来说，当为任何选定的任务给定一个查询（例如，一幅图像或一段视频）时，作者随机采样一组定义该任务的输入-目标对。这些对与视觉输入和高斯噪声分量一起被输入到联合扩散Transformer（J-DiT）中。在J-DiT中，这些元素被处理以生成潜在表示。最后，这个潜在表示通过ST-VAE解码器，转换到原始像素空间，产生所需的预测。为了更好地理解这个推理过程，请参考图2（a）。
本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。
收藏，分享、在看，给个三连击呗！
