---
title: 'JointDiT: Diffusion Transformer for RGB-Depth Modeling'
url: https://www.emergentmind.com/topics/jointdit
type: topic
---

# JointDiT: Diffusion Transformer for RGB-Depth Modeling

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" [2505.00482]. 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 [2601.03112] [2606.13676].

## 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 [2505.00482].

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 [2505.00482].

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 [2505.00482].

## 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 [2505.00482].

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 [2505.00482].

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 [2505.00482].

## 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

$$
L_{JCFM}(\theta) = E_{t_x, t_y, q(x_1, y_1), p_{t_x,t_y}(x, y \mid x_1, y_1)} \left[ \left\| v_{\theta,t_x,t_y}(x, y) - v_{t_x,t_y}(x, y \mid x_1, y_1) \right\| \right].
$$

This is paired with a transformer that predicts branch-specific vector fields,
$\hat{v}^{(\text{rgb})}$ and $\hat{v}^{(\text{depth})}$, conditioned on both noisy modalities, their timesteps, text embeddings, and the outputs of joint cross-attention modules [2505.00482].

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

$$
G_x = \mathrm{Attn}(S_x) + W_x(t_x, t_y) \cdot \mathrm{JointAttn}(S_x, S_y),
$$

$$
G_y = \mathrm{Attn}(S_y) + W_y(t_x, t_y) \cdot \mathrm{JointAttn}(S_x, S_y),
$$

with

$$
W_x(t_x, t_y) = \mathrm{sigmoid}\left( a \cdot \left( \frac{t_y}{t_x + t_y} - \frac{1}{2} \right) \right),
$$

$$
W_y(t_x, t_y) = \mathrm{sigmoid}\left( a \cdot \left( \frac{t_x}{t_x + t_y} - \frac{1}{2} \right) \right),
$$

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 [2505.00482].

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

- with `50%` probability, $t_x \sim f(t)$ and $t_y \sim g(t)$ independently, with symmetric swapping in practice;
- with `50%` probability, a single $t \sim f(t)$ is sampled and set as $t_x = t_y = t$.

The distributions are

$$
f(t) = \frac{(\sigma(z))^S}{1 + (S - 1)\cdot(\sigma(z))^S}, \quad z \sim N(0,1),
$$

with `S = 3.1582` for $f(t)$, and the same form with `S = 0.25` for $g(t)$. The stated effect is to cover both the diagonal region $t_x \approx t_y$ relevant to joint generation and the off-diagonal extremes required for conditional tasks [2505.00482].

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 [2505.00482].

## 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 [2505.00482].

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 [2505.00482].

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 [2505.00482].

## 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 [2505.00482].

| 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`, $\delta_1$ `96.9` | Ours base |
| ScanNet depth estimation | AbsRel `6.6`, $\delta_1$ `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 [2505.00482].

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` [2505.00482].

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 [2505.00482].

## 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 [2505.00482] | A diffusion transformer modeling the joint distribution of RGB and depth | Modality-wise timesteps, adaptive scheduling weights, unbalanced timestep sampling |
| DiT-JSCC [2601.03112] | 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 [2606.13676] | 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 [2601.03112].

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 [2606.13676].

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 [2505.00482].

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 [2505.00482].

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 [2601.03112] [2606.13676]. 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.

Source: https://www.emergentmind.com/topics/jointdit