JointDiT: Diffusion Transformer for RGB-Depth Modeling

• JointDiT is a diffusion-transformer based framework that jointly models RGB images and depth maps to support flexible joint generation, depth estimation, and depth-conditioned image synthesis.
• It employs modality-wise timestep control and adaptive scheduling weights to transfer information between noisy and clean modalities, enhancing geometric consistency and structure.
• The architecture integrates a frozen Flux backbone with trainable LoRA adapters and joint connection modules, increasing parameters by 19.8% while achieving competitive metrics across multiple benchmarks.

JointDiT denotes a diffusion-transformer framework for modeling the joint distribution of RGB imagery and depth, introduced in "JointDiT: Enhancing RGB-Depth Joint Modeling with Diffusion Transformers" (Byung-Ki et al., 1 May 2025). Its central premise is that a single unified generative model of images and depth can capture their mutual dependence across all noise levels and thereby support multiple tasks—joint generation, depth estimation, and depth-conditioned image generation—by controlling modality-wise timesteps rather than switching architectures or training separate conditional pipelines. In later work, the term also appears in a broader, paradigm-like sense: DiT-JSCC uses it to describe a jointly optimized semantics-prioritized encoder and conditional DiT decoder for generative joint source–channel coding, while Modality Forcing is described as a JointDiT-style method for scalable joint image-depth generation with separate per-modality noise levels (Tan et al., 6 Jan 2026, Duisterhof et al., 11 Jun 2026).

1. Conceptual basis and task formulation

The original JointDiT is motivated by the idea that direct modeling of $p(x_{\text{rgb}}, x_{\text{depth}})$ is structurally different from modeling $p(x_{\text{rgb}} \mid x_{\text{depth}})$ or $p(x_{\text{depth}} \mid x_{\text{rgb}})$. In the formulation reported for JointDiT, joint modeling supports a single parameterization that naturally handles multiple combinatorial generation tasks: joint generation when both modalities are noisy, depth estimation when RGB is clean and depth is noisy, and depth-conditioned image generation when depth is clean and RGB is noisy (Byung-Ki et al., 1 May 2025).

The stated rationale is not only flexibility but also geometric consistency. Because both modalities co-evolve during training and sampling, JointDiT is reported to yield smoother and more structured 3D lifting by inverse projection to point clouds than single-image depth estimators, whose outputs can exhibit rough surfaces under monocular ambiguity. The same framework can operate with or without text conditioning and is demonstrated at $512 \times 512$ during training while generalizing to $1024 \times 1024$ at inference (Byung-Ki et al., 1 May 2025).

A recurrent misconception is to treat JointDiT as merely a conditional depth estimator built on a strong text-to-image prior. The original description is more specific: it presents joint distribution modeling itself as a replaceable alternative to conditional generation, with competitive or matched performance on conditional tasks while maintaining superior joint RGB-depth synthesis (Byung-Ki et al., 1 May 2025).

2. Architectural organization

JointDiT is built on Flux, described as a Rectified Flow Transformer with DiT-style transformer blocks and strong image priors. The backbone comprises Multi-Modal DiT (MM-DiT) and Parallel DiT (P-DiT) blocks. Architecturally, the model has parallel RGB and depth branches connected through joint connection modules that exchange information at the attention level (Byung-Ki et al., 1 May 2025).

The RGB branch is the pre-trained Flux model and remains frozen during JointDiT training. The depth branch is newly introduced and trained through LoRA adapters on the frozen Flux backbone, with LoRA rank $64$ in all experiments. Only the LoRA parameters and the joint connection modules are trainable. The model inserts a joint cross-attention module inside each DiT block; this exchanges queries between RGB and depth branches to encourage aligned structure formation across modalities (Byung-Ki et al., 1 May 2025).

JointDiT inherits Flux’s pixel-space operation rather than using a VAE latent space. Inputs are $512 \times 512$ RGB images and disparity maps normalized to $[0,1]$. Text conditioning is supported through Flux’s text modules, with a frozen text encoder supplying embeddings, but text is not mandatory. The reported LoRA insertions are:

• MM-DiT: img_mod.lin, img_attn.qkv, img_attn.proj, txt_mod.lin, txt_attn.qkv, txt_attn.proj
• P-DiT: linear1, modulation.lin
• Input stage: vector_in.in_layer, vector_in.out_layer, txt_in

The reported cost profile is explicit. Relative to the base Flux model, JointDiT adds +19.8% parameters due to LoRA and joint modules. For 20 sampling steps, sampling time increases from 1.6 s for Flux image-only generation to 4.7 s for JointDiT RGB-depth generation on the reported hardware (Byung-Ki et al., 1 May 2025).

3. Joint flow matching, adaptive scheduling, and timestep control

The learning problem is formulated in continuous time through flow matching rather than DDPM-style Gaussian-Markov diffusion. JointDiT extends Conditional Flow Matching to two modalities by assigning separate timesteps $t_x$ and $t_y$ to RGB and depth. The stated Joint Conditional Flow Matching objective is

$p(x_{\text{rgb}} \mid x_{\text{depth}})$0

This is paired with a transformer that predicts branch-specific vector fields, $p(x_{\text{rgb}} \mid x_{\text{depth}})$1 and $p(x_{\text{rgb}} \mid x_{\text{depth}})$2, conditioned on both noisy modalities, their timesteps, text embeddings, and the outputs of joint cross-attention modules (Byung-Ki et al., 1 May 2025).

A distinctive mechanism is the use of adaptive scheduling weights inside attention, not as loss weights. The branch outputs are defined as

$p(x_{\text{rgb}} \mid x_{\text{depth}})$3

$p(x_{\text{rgb}} \mid x_{\text{depth}})$4

with

$p(x_{\text{rgb}} \mid x_{\text{depth}})$5

$p(x_{\text{rgb}} \mid x_{\text{depth}})$6

and a = 3 in all experiments. The noisier branch therefore receives more cross-modality information from the cleaner branch, which is intended to stabilize structure formation across disparate noise levels (Byung-Ki et al., 1 May 2025).

The second key technique is unbalanced timestep sampling. The training scheme mixes two cases:

• with 50% probability, $p(x_{\text{rgb}} \mid x_{\text{depth}})$7 and $p(x_{\text{rgb}} \mid x_{\text{depth}})$8 independently, with symmetric swapping in practice;
• with 50% probability, a single $p(x_{\text{rgb}} \mid x_{\text{depth}})$9 is sampled and set as $p(x_{\text{depth}} \mid x_{\text{rgb}})$0.

The distributions are

$p(x_{\text{depth}} \mid x_{\text{rgb}})$1

with S = 3.1582 for $p(x_{\text{depth}} \mid x_{\text{rgb}})$2, and the same form with S = 0.25 for $p(x_{\text{depth}} \mid x_{\text{rgb}})$3. The stated effect is to cover both the diagonal region $p(x_{\text{depth}} \mid x_{\text{rgb}})$4 relevant to joint generation and the off-diagonal extremes required for conditional tasks (Byung-Ki et al., 1 May 2025).

Inference is then controlled entirely through modality-wise timestep endpoints. Joint generation sets both branches to noisy initial states; depth estimation fixes RGB clean and updates only depth; depth-conditioned image generation fixes depth clean and updates only RGB. Classifier-free guidance is enabled by 10% text drop during training and uses guidance scale γ = 3.5 at inference (Byung-Ki et al., 1 May 2025).

4. Training regime and data construction

The main training corpus for JointDiT consists of 50k RGB–text–depth triplets collected from an internal real-world video corpus. Frames are resized to keep aspect ratio and center-cropped to 512×512. Depth supervision is obtained from Depth-Anything v2, then converted to disparity and normalized to [0,1]. Text prompts are generated using LLaVA (Byung-Ki et al., 1 May 2025).

For depth-estimation verification, the model is additionally fine-tuned on +80k synthetic samples from Hypersim, Replica, IRS, and MatrixCity. Depth and disparity are unified; invalid regions are handled explicitly, including sky in MatrixCity and near-camera artifacts in Replica. Bias and scale are applied to match the statistics of Depth-Anything v2’s disparity in valid regions, and invalid regions are filled with Depth-Anything v2 predictions for consistency (Byung-Ki et al., 1 May 2025).

The reported optimization setup is compact but specific: 75k iterations, batch size 4, learning rate 1e−5, and a single NVIDIA H100, for an overall training time of ~3.5 days. LoRA rank is 64 throughout, and only LoRA adapters and joint connection modules are trained. Inference uses 20 denoising or integration steps. These details reinforce that the method is framed as a parameter-efficient augmentation of a frozen high-capacity generative prior rather than end-to-end retraining of the base transformer (Byung-Ki et al., 1 May 2025).

5. Empirical performance and ablation structure

JointDiT is reported to achieve strong joint RGB-depth generation, competitive depth estimation, and strong depth-conditioned image generation within one unified model. The reported quantitative results span generative fidelity, semantic-text alignment, human preference, and downstream depth metrics (Byung-Ki et al., 1 May 2025).

Evaluation setting	Reported result	Context
ImageNet 6K joint generation	FID 24.26, IS 37.81, CLIP 30.51	Supplementary joint generation
Pexel 6K joint generation	FID 19.87, IS 22.51, CLIP 30.71	Supplementary joint generation
MSCOCO 30K joint generation	FID 11.27, IS 34.35, CLIP 30.76	Supplementary joint generation
NYUv2 depth estimation	AbsRel 5.7, $p(x_{\text{depth}} \mid x_{\text{rgb}})$5 96.9	Ours base
ScanNet depth estimation	AbsRel 6.6, $p(x_{\text{depth}} \mid x_{\text{rgb}})$6 95.7	Ours base
OpenImages 6K depth-conditioned image generation	FID 12.62, AbsRel 6.99	Best reported in comparison

Human preference, measured by ImageReward with higher Rank1% preferred, is also reported in JointDiT’s favor: across ImageNet, Pexel, and MSCOCO, the method achieves the highest Rank1 share at approximately 42–43% and the lowest Rank3 share at approximately 26–27%, outperforming JointNet and LDM3D (Byung-Ki et al., 1 May 2025).

The depth-estimation numbers show a more differentiated pattern. Against joint generation baselines, JointDiT outperforms JointNet and UniCon on all listed datasets. Against generative depth estimators such as Marigold and GeoWizard, it is described as comparable on most datasets except ETH3D, with the paper attributing that case to training depths sourced from Depth-Anything v2. Fine-tuning further improves several datasets, including NYUv2 from AbsRel 5.7 to 5.0 and ScanNet from 6.6 to 5.6 (Byung-Ki et al., 1 May 2025).

The ablation structure is central to the paper’s claims. On ImageNet joint generation, removing both adaptive scheduling weights and unbalanced timestep sampling degrades performance to FID 30.88, IS 31.61, and CLIP 29.80. Using weights only yields FID 24.20, IS 37.04, and CLIP 30.37. Using both yields FID 24.26, IS 37.81, and CLIP 30.51, with the best reported IS and CLIP. Similar gains are reported on Pexel 6K and MSCOCO 30K, and depth-estimation and depth-conditioned-image-generation ablations likewise improve when the two techniques are combined (Byung-Ki et al., 1 May 2025).

6. Subsequent reinterpretations and related developments

Later papers reuse the JointDiT label or explicitly position themselves relative to it, but they do so in ways that broaden the original notion beyond RGB-depth flow matching. The following comparison captures the uses described in the supplied sources.

Work	What “JointDiT” denotes	Distinguishing feature
JointDiT (Byung-Ki et al., 1 May 2025)	A diffusion transformer modeling the joint distribution of RGB and depth	Modality-wise timesteps, adaptive scheduling weights, unbalanced timestep sampling
DiT-JSCC (Tan et al., 6 Jan 2026)	A jointly optimized transceiver with a semantics-prioritized representation encoder and a conditional DiT decoder	Coarse-to-fine semantic/detail conditioning for generative joint source–channel coding
Modality Forcing (Duisterhof et al., 11 Jun 2026)	A JointDiT-style single DiT with separate noise levels per modality	Pixel-space depth, sparse real-world depth training, per-modality decoders

In DiT-JSCC, “JointDiT” refers to a jointly optimized transceiver in which a semantics-prioritized representation encoder and a conditional Diffusion Transformer decoder are co-designed and trained end-to-end for generative joint source–channel coding. The encoder produces complementary semantic and detail conditions that guide a coarse-to-fine DiT during denoising. This framing is explicitly presented as a response to the mismatch between reconstruction-oriented encoders and generative decoders in extreme wireless regimes, including ultra-low bandwidth and low-SNR settings (Tan et al., 6 Jan 2026).

Modality Forcing, by contrast, describes itself as analogous in spirit to JointDiT because it uses a single DiT trained with separate timesteps or noise levels per modality. Its emphasis is different: depth is represented in pixel space, missing depth pixels are replaced by isotropic Gaussian noise, and per-modality decoders allow training on sparse, real-world depth. The paper states that this design yields state-of-the-art image-to-depth results versus prior joint image-depth generative models and reports an average 57% AbsRel reduction relative to those prior joint models (Duisterhof et al., 11 Jun 2026).

These later uses suggest that “JointDiT” has acquired a broader methodological meaning: a single or jointly optimized diffusion-transformer system in which modality-specific conditions or timesteps are integrated into a shared denoising process. That broader usage should not obscure the narrower original referent, which is the Flux-based RGB-depth joint model introduced in 2025.

7. Limitations, trade-offs, and open directions

The original JointDiT reports several limitations and trade-offs. Computationally, generating both RGB and depth requires more memory and time than image-only Flux: the reported sampling time rises from 1.6 s to 4.7 s for 20 steps, and the parameter count increases by +19.8%. The main training data also relies on pseudo-depth supervision from Depth-Anything v2, which the paper identifies as a source of limited cross-dataset generalization, especially on out-of-distribution domains such as ETH3D (Byung-Ki et al., 1 May 2025).

The model also inherits stylistic tendencies from Flux, which can affect FID relative to real-image datasets. This is presented as a trade-off rather than a failure of the joint modeling premise. The paper nonetheless argues that the global receptive field and strong image prior of diffusion transformers benefit dense prediction and image synthesis simultaneously, making the unified formulation attractive even when it does not dominate every specialized baseline on every dataset (Byung-Ki et al., 1 May 2025).

Subsequent work highlights complementary limitations. DiT-JSCC notes reliance on a pre-trained VFM, per-SNR models, and a caption-based Kolmogorov-complexity proxy for adaptive bandwidth allocation, while Modality Forcing notes that depth-to-image conditioning can be looser than JointDiT in some scenes and that metric-scale depth is not the target because evaluation is affine-invariant (Tan et al., 6 Jan 2026, Duisterhof et al., 11 Jun 2026). Taken together, these points indicate that the JointDiT line of work is less a settled architecture than a family of joint-modality DiT strategies whose effectiveness depends on how modality coupling, conditioning strength, data sparsity, and training objectives are matched to the task domain.