Jodi: Unification of Visual Generation and Understanding via Joint Modeling

Title: \thetable Comparison of training costs among unified models. We use much less data and resources.
ArXiv: 2505.19084v1
Version: v1
Sections: 4
Estimated tokens: 1.2k

## Contents
- 1 Experiment
  - \thesubsection Setup
  - \thesubsection Visual Generation and Understanding
  - \thesubsection Analysis

## 1 Experiment

### \thesubsection Setup

[-0.5ex]Training Details
We adopt Sana [xie2025sana] as our base model.
We train our model using the CAME-8bit optimizer [luo2023came] for 130K steps, with a learning rate of $4\times 10^{-5}$, a batch size of 32, and BF16 mixed-precision, which takes around 535 hours on 8 RTX A6000.
Since our dataset contains images with various aspect ratios, we use a ratio bucketing strategy [novelai] during training to prevent important contents from being cropped.
This also allows users to generate images with a wide range of aspect ratios during inference.
It is worth noting that we use significantly less data and computational resources than the other unified models, as shown in \Tabtab:compare-cost.
\linetitle[0ex]Sampling Details
We employ Flow-DPM-Solver [xie2025sana], a variant of DPM-Solver++ [lu2022dpm++] adapted for rectified flow.
The classifier-free guidance [ho2022classifier] scale is set to 4.5.
For joint generation and controllable generation, we use 20 sampling steps.
For image perception, we use 10 sampling steps since increasing the steps leads to little performance gain.

### \thesubsection Visual Generation and Understanding

[-0.5ex]Joint Generation
In \Figfig:vis-joint-generation, we illustrate the capability of our Jodi to simultaneously generate high-quality images of various aspect ratios along with corresponding labels, including depth, normal, albedo, edge, lineart, segmentation, and openpose.
The generated images and the generated labels are semantically consistent and spatially aligned, credited to the linear attention and domain-invariant positional embeddings.
Please refer to the appendix for more results.
\linetitle[0ex]Controllable Generation
To demonstrate Jodi’s performance in controllable generation, we first generate images using existing labels as input conditions and evaluate their fidelity using FID scores [heusel2017gans].
Then, to evaluate the faithfulness of the generated images to the input conditions, we re-extract the conditions from the generated images and compare them to the input conditions using LPIPS [zhang2018unreasonable].
As shown in \Tabtab:compare-control and \Figfig:compare-control, our Jodi achieves superior performance for all conditions compared to existing unified models as well as generation-only specialist models.

Figure: Figure \thefigure: Joint generation of images and labels across a wide range of aspect ratios.

Table: Table \thetable: Quantitative comparison of controllable generation.

### \thesubsection Analysis

[0ex]Effect of Domain-invariant Positional Embeddings
As described in \Sectsec:method-architecture, we introduce domain-invariant positional embeddings to encourage the spatial alignment across different visual domains.
To validate the effect, we compare our models trained for 10K steps with and without positional embeddings, by observing whether the joint generated images and labels are spatially aligned.
As shown in \Figfig:ablate-posemb, our model aligns the image domain and label domains significantly better with positional embeddings, whereas obvious misalignment is observed without positional embeddings.
\linetitle[0ex]Attention Map Visualization
To further investigate how the tokens from different visual domains align and interact with each other, we pick two query tokens from the image domain and visualize the corresponding attention maps in \Figfig:vis-attention.
As can be seen, most domains show strong activation at the same spatial location as the query token, demonstrating a good alignment between these domains.
Interestingly, attention maps of different domains also reveal their own unique structural patterns.
For example, the segmentation domain exhibits strong activation along semantic boundaries, and the openpose domain focuses more on the human figure.
\linetitle[0ex]Joint Consistency
In \Figfig:vis-joint-consistency, we illustrate the consistency of our unified model across joint generation, controllable generation, and image perception tasks.
We first perform joint generation based on the input prompt to produce samples covering all of the image and label domains.
According to each generated label, we then apply controllable generation to generate new images that comply with these labels.
Besides, we perform image perception on the image generated in the first step to detect all its labels.
As can be observed, three types of tasks produce visually consistent results.
\linetitle[0ex]Extension to New Domains
Our well-trained Jodi model can be readily extended to one or more new domains by appending the corresponding tokens to the existing ones.
\Figfig:vis-extend presents the joint generation results after fine-tuning the model on the doodle sketch domain [arar2025swiftsketch] as well as simultaneous fine-tuning on the pixel, irradiance, and canny domains.

Figure: Figure \thefigure: Effect of positional embeddings. Generated labels are overlaid on images for a better view.

Figure: Figure \thefigure: Visualization of attention map.

Figure: Figure \thefigure: Jodi shows consistency among joint generation, controllable generation, image perception.

Figure: Figure \thefigure: Joint generation results of our model extended to new domains.

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 6
presentation: 3
questions: * How well does the method scale to more detailed label domains (e.g., full-class/instance segmentation or dense keypoints) without redesign?

* What are the failure cases on real or diverse datasets, especially with human subjects? Although the limitations are discussed, the failure cases are not provided (neither in the main paper nor in the supplementary).

* Can the authors provide per-domain trade-offs -- do any tasks degrade notably compared to specialist methods?
 The results currently reported appear not to be on fair ground comparisons.
soundness: 4
strengths: * The role switch mechanism:  by randomly switching each domain between being generated, used as conditioning, or ignored, the model learns all the key distributions for both image generation and perception at once -- giving one network the flexibility to do many things.

* Paper is well motivated, and generally well written and easy to follow.

* The proposed method uses shared positional embeddings across domains (+ small role tags) so that the model knows which pixels align spatially, making it easier to keep different visual modalities consistent with each other.

* The proposed Joint generation could be very valuable in many downstream applications, such as artistic manipulation or asset usage in the traditional frameworks. 
  
* Experimental results validate the proposed method convincingly.
summary: This paper introduces Jodi, a diffusion-based model that tries to unify image generation and image understanding by jointly modeling images and multiple label modalities. With the proposed “Role-Switch” mechanism, the same model can generate images with labels, generate images conditioned on any combination of labels, or predict labels from an input image. It uses a linear diffusion transformer and domain-invariant positional embeddings to keep computation manageable and maintain consistency across domains. The authors also provide a new dataset covering 8 visual domains. Experiments show good performance on both generation and perception tasks and strong scalability. Overall, the idea is intuitive and nicely executed, and the results suggest that a unified approach like this is becoming useful.
weaknesses: * “Extensibility to more domains” is claimed, but it’s unclear how much effort is needed for an entirely new modality.

* Baselines for multimodal generation/understanding might not share the same supervision or label availability. 
This needs to be clarified in the paper. 

* Randomly switching roles during training might make optimization harder or convergence slower, especially as domains increase.

* Even with linear attention, jointly modeling 8+ domains could still be memory-intensive and slow for large resolutions.

* The “multiple label domains” are pseudo labels (depth, normals, etc.), so errors/noise in these domains may propagate through the unified model.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: 1. The perception results (especially Tables 2, 4) are lower than some SOTA unified models and specialist models significantly. Is this performance gap an unavoidable cost of the "jack of all trades" unification, or do the authors see a clear path for Jodi to actually surpass some specialist models in understanding tasks?
2. What was the sampling strategy for the [G], [C], and [X] roles? Were they sampled uniformly at random for all domains? Or was a curriculum used (e.g., more [C] roles early in training) to stabilize the learning of this complex joint distribution?
soundness: 3
strengths: 1. Principled and Elegant Framework: The core idea of unifying $p(x|y)$ and $p(y|x)$ by modeling the joint distribution $p(x, y)$ is statistically elegant. The "Role-Switch" mechanism is a clever and direct implementation of this principle, forcing the model to learn a wide range of conditional and marginal distributions within one architecture.
2. The authors made smart architectural choices. The use of a linear diffusion transformer (Sana) correctly identifies and solves the $\mathcal{O}(M^2)$ computational bottleneck of multi-domain ($M$) attention. The "masked linear attention" and "domain-invariant positional embeddings" are solid supporting contributions that are well-motivated and essential for the framework to function.
3. The paper introduces the "Joint-1.6M" dataset (200K images + 7 predicted labels) and aggregates 90K images with ground-truth labels. This is a valuable contribution to the community that will enable future research in joint visual modeling.
summary: This paper presents Jodi, a unified diffusion framework designed to jointly model the distribution of images ($x$) and multiple label domains ($y^n$), such as depth, normal, and segmentation. The goal is to unify visual generation and understanding, which are typically treated as separate tasks. The core technical contribution is a "Role-Switch" mechanism, where each domain is randomly assigned as a generation target [G], a condition [C], or ignored [X] during training. This principled approach, based on joint probability modeling, allows the single model to perform three distinct tasks: joint generation ($p(x, y, ...)$), controllable generation ($p(x|y, ...)$), and multi-label image perception ($p(y, ...|x)$). The framework is built on an efficient linear diffusion transformer to handle the computational load of many domains and is trained on a newly curated "Joint-1.6M" dataset.
weaknesses: The model's "unification" comes at a significant cost to specialist performance. Though perform better than omni-models on edge detection and normal estimation tasks, Jodi performs noticeably worse than SOTA specialist models on albedo estimation and depth estimation. For example, in depth estimation (Table 2), Jodi achieves 10.1 AbsRel on NYUv2, while the specialist Lotus-D achieves 5.1. In albedo estimation (Table 4), Jodi gets 15.5 PSNR, while the specialist RGB2X gets 20.6. The model excels at generality but does not "excel in... understanding tasks" as claimed.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to "weaknesses"
soundness: 3
strengths: I like this paper due to the following strengths:

(1) Flexible Control: The framework naturally supports complex, multi-modal conditioning, offering unparalleled flexibility for creative applications.

(2) Enhanced Generalization: By forcing the model to simultaneously learn the generative process  and the analytical structure, JODI is likely to develop a richer, more robust latent representation of visual concepts.
summary: The paper proposes Jodi, a unified diffusion model that can jointly generate images and their corresponding conditions (depth, mask, etc). To achieve this, Jodi considers a unified input format (role switching) and utilizes domain-invariant positional embeddings to optimize flow-based DiT model. Extensive experiments on various benchmarks demonstrate their methods efficacy in unified generation.
weaknesses: (1) Lack of T2I evaluation. As demonstrated in Fig. 11, the generated image is the bridge between text and other condisions. Therefore, I would like to see some comparison with T2I methods on GenEval or other T2I benchmarks.

(2) The proposed method achieves condition to image by modeling the joint distribution of various conditions. However, some generation tasks (e.g., depth to image [a]) could be evaluated in a more precise manner. The authors should compare, or at least, discuss with such related works. 

(3) Minor: In Tab. 1-openpose, controlnet seems to achieve better result than yours in terms of LPIPS.


[a] 3dis: Depth-driven decoupled instance synthesis for text-to-image generation. In ICLR'25.
[b] 3DIS-FLUX: simple and efficient multi-instance generation with DiT rendering. In Arxiv.

---

confidence: 3
contribution: 2
initial_rating: 2
presentation: 3
questions: - Why would somebody want to use this instead of just chaining a few off-the-shelf models to do the target tasks? I get that there's some elegance to having it all in one model, but the pretrained models are already so powerful, that I'm skeptical that somebody would use this approach in practice. Furthermore, the results show that the specialist models often beat Jodi (although Jodi is better than the other joint models).
- What is the methodological novelty?
- Have you considered more modern fast attention mechanisms? [1] Or why not only train a few conditions at a time to get around the quadratic scaling?

[1] Guo, Han, et al. "Log-linear attention." arXiv preprint arXiv:2506.04761 (2025).
soundness: 2
strengths: - I like the probabilistic perspective on joint modeling p(x, y) = p(x | y) p(y) = p(y | x) p(x). Although simple, it provides an intuitive motivation for why you would want to pursue such a method.
- Visual results clearly show that your method works.
- The paper is well-written and easy to follow.
summary: This paper proposes a joint modeling framework to train an all-in-one generative model that can do joint generation of images and caption, conditional image generation, and label prediction.
weaknesses: see questions


Decision: Reject