JointDiT:使用扩散变换器增强RGB-深度联合建模

JointDiT: Enhancing RGB-Depth Joint Modeling

with Diffusion Transformers

一、论文基本信息

****论文标题：****JointDiT: Enhancing RGB-Depth Joint Modeling with Diffusion Transformers

****作者：****Kwon Byung-Ki, Qi Dai, Lee Hyoseok, Chong Luo, Tae-Hyun Oh

****机构：****POSTECH, Microsoft Research Asia, KAIST

****发表会议：****ICCV 2025

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

****项目主页：****https://byungki-k.github.io/JointDiT/

****代码仓库：****https://github.com/kaist-ami/JointDiT

二、研究背景

近年来，扩散模型在图像生成领域取得了举世瞩目的成就。以Stable Diffusion、DALL-E、FLUX等为代表的文生图模型，不仅能够根据文本描述生成高质量、高分辨率的图像，更在视觉理解与生成领域树立了新的标杆。这些模型通过学习海量图像数据的分布，掌握了丰富的视觉先验知识，包括物体形状、纹理细节、空间布局以及光照阴影等方面的深层理解。这些预训练的视觉先验为下游视觉任务提供了强大的基础能力。

文生图扩散模型的视觉先验已被广泛应用于多种下游任务中，包括逆问题求解（如超分辨率、去模糊、图像修复）、图像编辑与操控、深度估计以及深度条件图像生成等。这些应用的核心思想是利用预训练扩散模型中蕴含的丰富视觉知识，通过微调或条件引导的方式将其迁移到目标任务上，从而在无需从头训练的情况下获得优异的性能表现。

然而，现有的RGB-Depth联合建模方法（如LDM3D、JointNet等）在生成质量上仍存在明显不足。这些方法大多基于较早期的潜在扩散模型架构，无法充分利用最先进的扩散Transformer（Diffusion Transformer, DiT）架构所带来的性能提升。此外，传统的条件生成方法需要针对每个特定任务单独训练模型，例如深度估计需要一个模型，深度条件图像生成又需要另一个模型，这种范式不仅训练成本高昂，而且无法充分利用RGB与深度信息之间的内在关联。

相比之下，联合分布建模具有显著的优势：通过学习RGB图像与深度图的联合概率分布，一个统一的模型即可处理多种组合任务，包括联合生成、深度估计和深度条件图像生成。这种范式的核心挑战在于：如何在保持高保真图像生成的同时，生成几何上合理且准确的深度图。JointDiT正是针对这一挑战提出的创新解决方案。

三、核心方法详解

JointDiT的核心思想是在最先进的扩散Transformer架构基础上，通过引入并行的深度分支和精心设计的训练策略，实现RGB图像与深度图的高质量联合建模。以下从五个方面详细阐述其方法设计。

3.1 整体架构

JointDiT基于FLUX模型构建，FLUX是目前最先进的扩散Transformer模型之一。FLUX的架构由两种核心模块组成：MM-DiT（多模态扩散Transformer）块和P-DiT（并行扩散Transformer）块。MM-DiT块负责处理文本条件与图像特征之间的交互，而P-DiT块则通过并行处理进一步提升模型的表达能力和计算效率。

在架构设计上，JointDiT在预训练的FLUX模型的RGB分支旁引入了一个并行的深度分支。为了使预训练模型能够处理额外的深度域信息，研究者在MM-DiT和P-DiT块中添加了LoRA（低秩自适应）模块。LoRA是一种参数高效的微调方法，通过在预训练权重矩阵旁添加低秩分解矩阵，可以在不显著增加参数量的情况下适配新的任务域。在训练过程中，预训练骨干网络的权重被完全冻结，仅训练LoRA参数和联合连接模块，这既保护了预训练模型中宝贵的视觉先验知识，又使得训练过程更加高效和稳定。

3.2 联合条件流匹配（JCFM）

JointDiT的另一个核心创新是联合条件流匹配（Joint Conditional Flow Matching, JCFM）框架。该方法将传统的流匹配（Flow Matching）框架扩展到多模态联合分布学习的场景。具体而言，JCFM学习一个联合向量场 v(t_x, t_y)(x, y | x_1, y_1)，其中x和y分别表示RGB图像和深度图，x_1和y_1表示目标数据分布中的样本，t_x和t_y是两个独立的时间步变量。

使用两个独立时间步的设计是JCFM的关键创新之一。通过灵活地调整t_x和t_y的值，可以在推理时实现不同的任务模式：当t_x=0且t_y=0时，模型执行联合生成任务，同时生成RGB图像和深度图；当t_x=1且t_y=0时，模型接收噪声化的RGB输入和纯噪声的深度输入，执行深度估计任务；当t_x=0且t_y=1时，模型接收纯噪声的RGB输入和噪声化的深度输入，执行深度条件图像生成任务。这种通过简单的时间步控制实现多种任务切换的设计，极大地提升了模型的灵活性和实用性。

JCFM的损失函数基于条件流匹配的目标函数进行设计，通过最小化模型预测的向量场与目标向量场之间的差异来优化模型参数。由于两个模态使用独立的时间步，损失函数需要对两个模态的重建误差进行联合优化，确保模型能够同时学习RGB和深度两个域的分布特征。

3.3 自适应调度权重

在RGB-Depth联合建模中，一个关键问题是如何在两个分支之间有效地传递信息。JointDiT提出了自适应调度权重（Adaptive Scheduling Weight）机制来解决这个问题。其核心思想是：根据RGB和深度分支的相对"干净度"来自适应地调度信息传递的强度。在扩散过程的去噪过程中，更干净的数据（即更接近目标分布的数据）通常包含更多有用的结构信息，能够为另一个模态的生成提供更准确的联合生成引导。

具体而言，自适应权重被应用于联合交叉注意力的输出上。数学上，RGB分支的特征更新可以表示为：G_x = Attn(S_x) + w_x(t_x, t_y) · JointAttn(S_x, S_y)，其中S_x和S_y分别是RGB和深度分支的特征表示，Attn表示自注意力操作，JointAttn表示联合交叉注意力操作，w_x(t_x, t_y)是根据两个时间步计算得到的自适应权重。当某个分支的数据更干净时，其对应的自适应权重会增大，使得该分支能够向另一个分支传递更多的有用信息。这种设计确保了信息传递的方向和强度始终与去噪进程的实际需求相匹配。

3.4 非平衡时间步采样策略

由于JCFM框架使用了两个独立的时间步变量，在训练过程中需要从二维时间步空间中采样。如果采用均匀采样策略，可能会导致不同模态之间的噪声级别出现严重不平衡的情况。例如，当t_x接近1（高度噪声化）而t_y接近0（几乎无噪声）时，两个模态处于截然不同的去噪阶段，这种不平衡可能导致训练不稳定或模型性能下降。

为解决这一问题，JointDiT提出了非平衡时间步采样策略。该策略通过精心设计的时间步采样分布，确保在训练过程中两个模态的噪声级别保持相对平衡，从而提高训练效率和模型性能。这种采样策略使得模型能够在各种噪声级别组合下都能获得充分且均衡的训练信号，避免了某些极端情况下的训练不足问题。

3.5 联合连接模块

联合连接模块是JointDiT实现RGB与深度分支信息交互的核心组件。该模块在每个DiT块的注意力机制中引入了RGB和深度分支之间的特征交换机制。具体而言，联合连接模块采用联合交叉注意力（Joint Cross-Attention）的设计，使得每个分支不仅能够通过自注意力处理自身模态内的特征依赖关系，还能够通过交叉注意力从另一个模态中获取互补信息。

研究指出，自注意力机制在图像生成中扮演着关键角色，它能够捕获图像中的长程依赖关系和全局结构信息。在联合建模的场景下，RGB分支的自注意力主要关注图像的外观和纹理特征，而深度分支的自注意力则更关注空间结构和几何关系。通过联合交叉注意力，两个分支能够互相补充彼此缺失的信息，从而实现更高质量的联合生成。这种设计使得RGB分支能够利用深度信息来改善空间一致性，同时深度分支也能够利用RGB信息来生成更精细的几何细节。

四、实验结果分析

4.1 联合生成

在联合生成任务上，JointDiT取得了显著优于先前方法的性能表现。与LDM3D、JointNet等现有方法相比，JointDiT生成的RGB图像具有更高的保真度和更丰富的细节，同时生成的深度图在几何准确性和结构合理性方面也有明显提升。实验结果表明，基于最先进的DiT架构进行联合建模，能够有效克服传统基于潜在扩散模型的方法在生成质量上的瓶颈。

特别值得注意的是，JointDiT在卡通图像和像素艺术插画等挑战性领域表现出色。这些领域的图像具有独特的风格化特征，与自然图像存在显著差异，传统方法往往难以生成高质量的结果。JointDiT通过充分利用预训练FLUX模型中丰富的视觉先验知识，结合联合分布建模的优势，在这些困难场景下依然能够生成高质量的RGB-Depth配对结果。此外，通过3D提升（Lifting）技术将生成的深度图转换为3D点云的结果显示，JointDiT生成的深度图能够产生更加合理和一致的3D结构。

4.2 深度估计

在零样本深度估计任务上，JointDiT取得了与现有专门方法可比的结果。这一发现具有重要的理论意义：它证明了联合分布建模不仅可以用于生成任务，还可以作为条件生成（即给定RGB图像估计深度图）的可行替代方案。通过将t_x设为1（保持RGB输入不变）并将t_y设为0（从噪声生成深度图），联合训练的模型能够在不需要额外训练的情况下完成深度估计任务，展现了联合分布建模范式的强大泛化能力。

4.3 深度条件图像生成

在深度条件图像生成任务上，JointDiT同样取得了与现有方法可比的结果。通过将t_x设为0（从噪声生成RGB图像）并将t_y设为1（保持深度输入不变），模型能够根据给定的深度图生成与之对应的RGB图像。实验结果表明，生成的图像在空间结构上与输入深度图保持良好的一致性，同时在视觉质量上达到了较高的水平。这一结果进一步验证了联合分布建模在多任务统一处理方面的优势。

4.4 消融实验

消融实验充分验证了自适应调度权重和非平衡时间步采样策略的有效性。实验结果表明，移除自适应调度权重会导致联合生成质量明显下降，证明了该机制在跨模态信息传递中的关键作用。同样，移除非平衡时间步采样策略也会导致性能退化，特别是在某些极端噪声级别组合下的表现更为明显。

此外，消融实验还揭示了RGB和深度分支之间有趣的互补行为：深度分支倾向于捕获场景的全局结构信息，而RGB分支则更多地关注纹理和外观细节。这种互补性是联合建模能够成功的关键因素之一，也暗示了扩散模型内部可能存在某种形式的模块化特征表示，不同分支自然地分化为处理不同类型的视觉信息。

五、与相关工作的对比

JointDiT与现有的RGB-Depth联合建模方法存在本质性的差异。LDM3D是基于潜在扩散模型（Latent Diffusion Model）的方法，虽然开创了RGB-Depth联合生成的先河，但其基于U-Net的架构在生成质量和细节保真度方面存在固有限制。JointNet虽然也采用了联合建模的思路，但需要并行运行两个独立的扩散模型，计算开销较大，且难以充分利用最先进扩散模型的架构优势。

在深度估计方面，Marigold和GeoWizard等方法通过微调预训练扩散模型来实现单目深度估计，虽然取得了优异的性能，但它们本质上是条件生成方法，需要针对深度估计任务进行专门的训练和优化。相比之下，JointDiT通过联合分布建模，在训练一个模型的同时即可获得深度估计的能力，无需额外的任务特定训练。

与OneDiffusion等需要大规模重训练的方法相比，JointDiT采用LoRA微调策略，仅需要训练少量参数即可实现RGB-Depth联合建模，训练成本显著降低。同时，由于预训练骨干网络被冻结，JointDiT能够完整保留FLUX模型的强大生成能力。这种基于最先进DiT架构的联合分布建模方法，在生成质量、任务灵活性和训练效率之间实现了良好的平衡。

六、优缺点分析

优点

首次利用最先进的扩散Transformer进行RGB-Depth联合建模，成功将FLUX模型的强大生成能力迁移到联合建模任务中，显著提升了生成质量。
通过简单的时间步控制（调整t_x和t_y的值）即可实现多种组合任务的切换，包括联合生成、深度估计和深度条件图像生成，极大地提升了模型的实用性和灵活性。
自适应调度权重和非平衡时间步采样策略的设计巧妙，有效解决了多模态联合建模中的信息传递和训练平衡问题。
在卡通图像和像素艺术插画等挑战性领域表现出色，展现了模型强大的泛化能力。
联合分布建模被证明可以作为条件生成的可行替代方案，为多模态视觉任务提供了新的解决思路。

缺点

依赖预训练的FLUX模型，如果需要迁移到其他扩散模型架构（如Stable Diffusion 3或DALL-E 3），可能需要进行额外的架构调整和超参数搜索。
训练过程需要RGB-Depth配对数据，这类数据的获取成本较高，可能限制了模型在更多场景下的应用。
推理时需要同时运行RGB和深度两个分支，计算开销相对较高，可能影响实时应用场景中的部署。
论文未充分讨论模型在真实场景深度数据（如LiDAR采集的深度数据）上的泛化能力，在真实世界应用中的表现有待进一步验证。

七、个人见解和未来展望

JointDiT的工作为多模态联合分布建模提供了一个极具参考价值的范例。从更宏观的视角来看，联合分布建模与条件生成代表了两种不同的多模态学习范式。条件生成方法针对每个任务单独建模，虽然可以在特定任务上取得优异性能，但缺乏对模态间深层关联的建模。而联合分布建模通过学习模态间的联合概率分布，能够捕获更深层次的跨模态依赖关系。JointDiT的成功证明了联合分布建模的巨大潜力，这种范式很可能成为未来多模态统一模型的主流方向。

消融实验中揭示的RGB和深度分支互补行为是一个值得关注的重要发现。深度分支自然地倾向于捕获结构信息，而RGB分支则关注纹理和外观，这种分工暗示扩散模型内部可能存在某种模块化的特征表示机制。这一发现对于理解扩散模型的内部工作机制具有重要意义，也为未来的模型设计和优化提供了新的思路。

对3D理解领域而言，JointDiT的高质量RGB-Depth联合生成能力可以为3D重建、场景理解等下游任务提供丰富的合成训练数据，特别是在真实深度数据稀缺的场景下具有重要的应用价值。展望未来，该方向有以下几个值得探索的研究方向：第一，扩展到更多模态的联合建模，如表面法线、语义分割图等，构建更加全面的多模态联合表示；第二，探索视频级别的RGB-Depth联合建模，利用时序信息提升深度估计的时序一致性；第三，结合3D高斯溺射（3D Gaussian Splatting）或神经辐射场（NeRF）等技术，从联合生成的RGB-Depth数据中实现更高质量的3D重建；第四，研究联合分布建模在其他视觉任务中的应用，如RGB-红外、RGB-事件相机等多模态场景。

对于从事扩散模型和多模态学习研究的学者而言，建议密切关注多模态联合分布建模这一新兴方向。随着扩散Transformer架构的不断演进和训练数据的持续增长，联合分布建模有望在更多视觉任务中展现其独特优势，推动计算机视觉领域向更加统一和通用的方向发展。