JointConn-v2：利用扩散变换器学习联合向量场，实现跨尺度连接与双时间步建模

《Neural Networks》：JointConn-v2: Learning a Joint Vector Field with Diffusion Transformers for Cross-Scale Connectivity and Dual-Timestep Modeling

　　摘要：本研究重新探讨了用于相对深度条件下的联合图像-深度合成的扩散变换器，重点解决了两个瓶颈问题：（1）跨模态注意力在边缘和结构区域表现不佳，从而导致几何失真；（2）深度分支容易受到语义分支的过度影响，进而引发过度耦合与不稳定现象。为解决这些问题，我们提出了带有GCM-WFM（门控跨模态加权流匹配）的JointConn-v2框架，该框架可用于联合图像合成及基于深度条件的图像合成。JointConn-v2通过Swap-Q交叉注意力强化双向引导，利用几何掩码偏差将二维相对位置和边缘能量直接注入注意力得分中，通过区域路由在标记级别决定何时以及何处融合跨模态特征，并借助带残差融合的内容门控在样本层面和结构层面稳定跨分支信息传递。在训练方面，我们引入了GCM-WFM机制，它在紧凑序列空间中回归联合向量场，同时将时间、几何、门控及路由项以对角权重形式纳入目标函数。该方法实现了边缘可控性与跨模态一致性之间的更好平衡。当前模型将深度视为标准化后的相对几何信号，其目的并非保持绝对度量尺度。相关代码可在https://github.com/haizhu12/JointConn-v2获取。

引言：文本到图像（T2I）模型如Shi等人（2024年）、Wang等人（2024年）、Yin等人（2024年）的研究成果正逐步从创意内容生成领域拓展到注重几何特性的现实应用场景，例如VR/AR系统（Savva等人，2019年；Yang等人，2024b年）、动画制作（Xing等人，2024年；Xu等人，2024年；Zhong等人，2025年）、图像编辑（Qi等人，2023年；Zhao等人，2023a年）以及机器人技术（Lu等人，2024年）。在这些应用场景中，仅靠RGB图像往往不够：物体放置、遮挡推理以及符合物理规律的光照等下游任务都需要深度图，而这类深度图不仅要在视觉上合理，还要与图像结构紧密契合，尤其是在物体边界附近。这促使我们研究如何以稳定、可控且高效的方式共同建模文本、RGB图像与深度信息。如图1所示，本文提出的联合建模方法既支持联合生成，也支持基于深度条件的生成。在本研究中，我们聚焦于相对深度而非绝对深度，这一选择与此处考虑的文本条件生成场景相契合，因为该场景的核心目标是保持物体布局、遮挡结构以及RGB图像与深度边界的一致性。由于MS COCO这类大规模图像-标题数据集并未提供经过校准的深度标注，因此我们在主要实验中使用的深度信号被视为尺度不确定的相对几何信息。相应地，我们的方法不应被视作绝对深度预测工具或精确的3D重建系统。尽管已有诸多进展，但要实现可靠的RGB图像与深度一致性仍然面临挑战。一个常见的问题是，跨模态交互在边缘和结构区域会下降，从而导致几何失真（Hsiao等人，2024年；Yang等人，2023年）；而过度强的耦合则可能使几何分支受语义信号的牵制，进而引发不稳定现象和伪影（Li等人，2025c年；Zhai等人，2024年）。当深度信号存在噪声、尺度不确定或属于异常分布时，这些问题会更为突出，而这类情况在实践中十分常见。现有方法大致可分为两类：基于控制器的条件注入方法（Brooks等人，2023年；Huang等人，2023年；Lin等人，2024年；Mou等人，2024年；Zhang等人，2023b年）通过额外的条件路径或适配器将深度信息引入几乎未被修改的T2I模型主干结构中。这类方法通常具有较高的参数效率且可控性较强，但信息注入往往仅在少数固定层中进行，可能无法适应不同时间步和空间位置的变化，因此可能无法充分利用边界附近的几何信息，从而导致纹理与几何结构不匹配（Cao等人，2023年；Ke等人，2024年；Zhong等人，2025年）。相比之下，端到端的联合建模方法（Li等人，2024b年；Ma等人，2023年；Xu等人，2023年）会在相关模态上学习单一模型以实现一致性，但这可能会加剧跨模态干扰：在平坦或不可靠的区域强制紧密耦合可能会传播错误的几何信息，使优化过程变得敏感，从而限制方法的稳健性和可控性。我们提出了一种统一框架，在共享的扩散变换器主干结构（Peebles和Xie，2023年）基础上，同时支持（i）基于文本条件的联合RGB图像-深度生成，以及（ii）基于深度条件的T2I生成。我们的核心思路是将跨分支交互视为一种有选择性的、注重几何特性的信息传递过程（Mou等人，2024年），而非固定的全局耦合机制。我们设计了一种轻量级的跨分支连接器（JointConn-v2），它能根据每个样本和每个空间位置自适应调节Go等人（2023年）、Tang等人（2023年）、Xu等人（2023年）所提出的耦合强度。这种设计能够在结构可靠的区域，如物体边界附近，增强信息交互，而在不可靠区域则减少过度耦合。为在几何可靠性不同的情况下稳定学习过程，我们设计了GCM-WFM（门控跨模态加权流匹配）目标函数（Lipman等人，2022年；Tong等人，2023年），该函数能够重点处理难度较大或信息丰富的区域，从而提升优化过程的稳定性。我们的方法在基于深度引导的生成任务中提升了语义质量与几何一致性，相比现有强基线模型，实现了更好的边界对齐效果和更低的深度误差。例如，与最新的联合DiT基线模型相比，我们的完整模型将深度均方根误差从0.150降低到0.138，同时GradCorr值从0.728上升至0.745，Edge-F1值从0.580提升到0.607，FID值也从16.80降至15.60，这些指标都表明我们的生成结果更具真实性且结构更加一致。我们的贡献可总结如下：1. JointConn-v2：一种注重几何特性、具备可控性的跨分支路由器。它通过结合样本级内容门控与标记级区域路由功能，实现针对每个样本和每个位置的自适应耦合，同时优先保证边界处的结构对齐，从而提升RGB图像与深度信息的一致性及可控性。2. GCM-WFM：一种稳定且统一的训练目标函数。它将流匹配监督任务转化为一系列自适应的正权重回归替代任务，通过将时间权重、门控机制、路由规则以及边缘能量整合为标准化的标记权重，从而稳定优化过程。3. 我们设计了一种参数效率较高的训练框架。该框架仅在FLUX风格的DiT模型主干结构上训练LoRA（分支内部优化）和JointConn-v2（分支间连接优化）模块；结合紧凑型潜变量表示与时间步控制器，该设计充分利用了扩散变换器的特性，可同时用于联合生成和基于深度条件的生成任务。在我们的实现中，冻结的基础模型采用的是Black Forest Labs开发的FLUX.1-dev DiT主干结构，该模型的DiT主干参数总量为11,901,408,320，即119亿个参数。

相关研究：联合生成。近期的扩散模型已从单模态合成发展为学习多个相关模态之间的联合分布，例如RGB图像与深度/几何信息。早期的多模态扩散模型旨在通过共享通用潜空间或学习跨模态的统一去噪模型，将不同的条件生成任务整合到同一框架中（Bao等人，2023年；Hu等人，2022b年；Li、Herrmann、Chan、Li、Sun、Ma和Yang，Li、Li、Shi、Farimani、Kluger）。

条件流匹配。条件流匹配（CFM）算法（Chen等人，2018年；Lipman等人，2022年；Liu等人，2022年；Tong等人，2023年）将生成过程建模为微分方程形式：dz/dt = vθ(zt, t, c)，其中t∈[0,1]，zt∈Rd为潜在状态，c为条件参数。给定条件c后，CFM算法会根据边际分布μ和πc采样端点对(z0, z1)，并采用线性路径zt=(1?t)z0+tz1，同时以oracle速度ut(z0,z1)=z1?z0作为路径演化规则。该算法的训练目标为简单的回归任务：LCFM(θ)=∫E{z0,z1～γc}E[t～ρ][∥vθ(zt,t,c)?ut(z0,z1)∥2]/2。在测试阶段，可随机采样z0～μ。

方法概述。我们的目标是在一个统一的网络中，建模不同噪声水平下RGB图像与深度图的联合分布。目前该网络支持两类任务：（i）联合图像-深度生成；（ii）通过调整各模态的噪声水平来实现基于深度条件的图像生成。我们采用了FLUX架构中的Esser等人（2024年）提出的两阶段DiT流程（MM-DiT → P-DiT），让RGB图像分支与深度分支并行发展。除JointConn-v2和LoRA模块外，还使用了Hu等人（2022a年）提出的其他组件。

实验设置：所有定量对比均在512×512分辨率下进行。这样的设定可以避免因不同基线模型的分辨率调整策略、宽高比分类方式或高分辨率采样方法带来的干扰因素，同时也能确保所有方法的边界相关指标，如Edge-F1和GradCorr，都在相同的空间网格上计算得出。

结论：在本研究中，我们提出了一种与DiT模型兼容的框架，可用于实现外观特征与几何结构之间的选择性、注重几何特性的耦合。在联合RGB图像-深度生成以及基于深度条件的图像合成任务中，该框架都能提升图像质量指标，即使在复杂的深度统计条件下，也能实现更一致的结构性与边界对齐效果。消融实验表明，所提出的耦合机制与注重几何特性的监督策略具有互补作用。除了整体性能指标之外，该框架还能为各项性能贡献提供清晰的分析依据。

作者贡献声明：Honggang Zhao：验证、监督、方法设计、研究分析、资金获取、正式分析、数据整理、概念构思。Yi-jun Yang：。Wei Zeng：。

利益冲突声明：作者声明不存在任何可能影响本文研究结果的已知财务利益或个人关系。

Honggang Zhao|Yi-jun Yang|Wei Zeng