Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities

We thank you for your insightful feedback. We answer your concerns below.

Concerns over scalability of experts in SDXL: We acknowledge your concern regarding the performance drop observed when increasing the number of experts in SDXL —a point also raised by Reviewer ptLD. In our rebuttal for Reviewer ptLD, we demonstrate that multi-task performance is linked to the width (i.e., expert capacity) and depth of the model. Specifically, reducing the width of the experts in shallow models like SDv1.5 works well, however in the case of deeper models like SDXL, larger experts perform well.

We point out that in Tables 1 and 4 of the main paper, we maintained the overall model size by reducing the hidden dimensions of the experts as their number increased. However, when this constraint is removed (as shown in Table A of ptLD's rebuttal), increasing the number of experts in our MoE setup significantly enhances multi-task performance while remaining more efficient than deploying separate models for each task. These results address your concerns regarding scalability and align with our goal of avoiding parameter inflation as the number of tasks grows.

Exploring task interference: We thank the reviewer for raising this important point. In the table below, we present an analysis of how different tasks interact by training various combinations of the four tasks considered in our paper. For this analysis, we trained SDv1.5 with $N=4, G=0.25$ (where $G$ is the ratio of an expert's dimension in MTU to the FFN dimension of the original model) across all task combinations, and we also report the performance of single-task models trained without interference. (Note: In the main paper, the performance of the super-resolution (SR) model is reported for an open-source version that uses image captions to generate high-resolution images. For a fair comparison and to analyze task interference, we trained an SR model with an empty string as the text condition and report its performance accordingly. Please look at Section A.2 in the paper.)

Our summary of task interference from the above table -

• First, we observe that image editing (IE) and image inpainting (IP) are highly compatible tasks, with performance increasing from 44.0 in our MTU model to 52.6 when combined. In contrast, super-resolution (SR) and text-to-image (T2I) appear to be less compatible with IP, likely because T2I and SR require generating new objects, whereas IP focuses on removing objects. Notably, the compatibility between IP and IE may stem from the fact that the IE dataset (InstructPix2pix) includes editing instructions that also involve object removal.
• We observe that T2I and IE are highly compatible, as training them together improves IE performance—an effect also highlighted in Figure 5 in the main paper, where these tasks select the same experts. While SR and IE may be compatible when trained together, adding IP into the mix leads to a significant drop in SR performance. Moreover, when T2I, IE, and IP are trained concurrently, T2I performance declines; however, the strong compatibility between IP and IE ensures that their performance remains high.
• In our multi-task model, where all tasks are trained jointly, we observe that the performance of the T2I, IE, and SR tasks improves compared to their single-task baselines, while performance decreases for IP by 2 points. This indicates that some degree of task interference may be prevalent.

We appreciate the reviewer for highlighting this issue. However, we view reducing task interference as an important direction for future work and will address this as a limitation in the revised version of the paper.

Thank you for your insightful review. We acknowledge that your concerns mainly pertain to Table 4, where we present an ablation study on the number of experts for both SDv1.5 and SDXL. We understand the reviewer is concerned regarding the significance of our method, given that in many cases a single expert ($N=1$) appears to perform adequately. We address your concerns below.

Experimental setup in Table 4: Our paper aims to maintain a parameter count similar to that of the original pre-trained model. To achieve this, we reduce the hidden dimension of each expert as the number of experts increases, ensuring that the ratio of an expert's dimension in MTU to the FFN dimension of the original model, remains $G = \frac{1}{N}$.

In Table 4 of the paper, we conducted an ablation study over $N$ (and consequently $G$) to determine the optimal configuration for both SDv1.5 and SDXL. Our results indicate that for SDv1.5, the best configuration is $N=4, G=0.25$. In contrast, for SDXL, the optimal performance is achieved with $N=1, G=1$, suggesting that retaining the original FFN size, with the addition of task-specific layer norms, is sufficient for handling multiple tasks in SDXL. Following this, we conduct an ablation over the size of the experts by removing the parameter constraint in Table 4.

Ablation over size of the experts: We hypothesize that for a given model, multi-task performance depends more on the size of the FFNs than on the number of experts. To test this, we removed the parameter constraint and performed an ablation study by varying $G$ from $\frac{1}{N}$ to 1. Table A below presents these results for SDXL, revealing that performance on tasks, particularly IE, SR, and IP, improves as the expert dimension increases. For instance, when $N=2$ with $G=1$ and $N=4$ with $G=1$, the performance in SR and IP improves significantly compared to the model reported in the paper $N=1, G=1$. These findings support our hypothesis. We also see that for a fixed value of $G$, increasing the number of experts improves performance for image-image tasks. Therefore, the MoE architecture becomes especially significant when scaling the model size for SDXL.

Table A: Ablation over number of experts for SDXL

Analysing the ablation over number of experts for SDv1.5 (Table 4): For SDv1.5, we determined that the optimal configuration is $N=4$ and $G=0.25$. Models using $N=1, G=1$ occasionally produced unusual artifacts in SR, and while we do not fully understand the cause, the $N=4, G=0.25$ setup consistently avoided these issues. Additionally, training became unstable with $N=8$ or $N=16$ due to the excessive reduction in expert size. The $N=2, G=0.5$ configuration also led to instability, likely because the experts conflicted across different tasks. Although $N=1, G=1$ may face similar task conflicts, its larger expert capacity helps mitigate this problem (similar to SDXL). In contrast, with $N=4, G=0.25$, tasks can select non-conflicting experts (see Figure 5), resulting in stable training.

Why do we need MTU?:

• Our results suggest that multi-task performance is linked to the width (i.e., expert capacity) and depth of the model. Specifically, reducing the width of the experts in shallow models like SDv1.5 works well, however in the case of deeper models like SDXL, larger experts perform well.
• However, the optimal expert size and count depend on your compute. If you are limited by resources $N=1, G=1$ may provide good results due to large expert capacity; however, if you can accommodate a larger model, an MTU setup with $N=4, G=1$ will lead to even better performance.
• Importantly, the success of the $N=1$ setup does not diminish the value of the MoE approach—in fact, our findings show that as the model scale increases, the benefits of MoE become even more pronounced, while still outperforming the alternative of using separate models for all tasks that would require significantly more compute.

We hope this addresses all of the reviewer's concerns, and we thank you for your feedback led to interesting experiments.

Thank you for your positive review, and for recognizing that MTU is the only approach so far to handle 4 different downstream tasks. We address your concerns below.

1. Quality of image generation: We believe our MTU model produces images of comparable quality to the base model. While Figure 3 does show some instances of over-saturation in the inpainting (IP) images (left columns, bottom row) for MTU model, the same over-saturation can also be seen in single-task models. The appendix includes several additional examples where this is not an issue. We point to Figure 7 in the appendix, where we show qualitative results of text-to-image generation for SDv1.5 and SDXL, where the quality of the images matches that of the base model. Images selected in the main paper and the appendix were selected at random from the test set.

2. Clarifications about Figures 1 and 3: In Figure 1, the top row displays results from SDv1.5, while the bottom row features outputs from SDXL. In Figure 3, the image editing prompt used is "turn into a polar bear."

Please let us know if you have any more questions. Thank you.

Thank you for your review. We have answered your questions below.

Weight distributions of FFNs in SD3.5 for image-to-image tasks: Thank you for your comment. To the best of our knowledge, general image-to-image generation using MMDiT-based models has not been thoroughly explored. Although there is some work in image editing [1], this work was released as open source only after our submission, and its applicability to other image-to-image tasks remains unclear. To address your question, we designed a baseline approach that adds separate positional embeddings to the source and target images, concatenates them along the feature dimension, and processes the extra channels with a linear layer while fine-tuning the entire model. While this baseline produces decent results, we have not included them in our paper, as image-to-image generation in MMDiTs is a research topic in its own right.

To understand the distribution of weights within our proposed baseline, we fine-tuned SD3.5 on super-resolution (SR) and image editing (IE) tasks, then measured the distance between the fine-tuned and pre-trained models for the FFNs and various attention layers. The table below reports the mode of the distance across all layers, revealing that the FFNs deviate most from the pre- trained initialization. While we see some outlier layers in some attention layers, the deviation is consistently high in FFNs. This suggests that these layers should be split into experts for MTU.

Table A: Mode of the deviations between the fine-tuned and pre-trained model for SDv3.5.

Training data and batches: Each training batch incorporates all tasks, and we minimize the sum of all losses in each iteration. We opted against random task selection per batch to avoid increased training time. During training, we construct the dataloader using the smallest dataset (super-resolution) and then sample subsets of equal size from the other datasets, shuffling these subsets each epoch. On average, this approach maintains balance across datasets.

Inpainting performance: Thank you for pointing this out. Reviewer FqBN asked about our analysis of task interference, and we addressed this issue in our response to Reviewer FqBN. In summary, we analyze model performance by evaluating different combinations of tasks to understand how they interfere with one another. Our findings indicate that inpainting (IP) is less compatible with text-to-image generation (T2I) and super resolution (SR), and it is highly compatible with image editing (IE). This incompatibility may contribute to the observed performance drop in IP. We view reducing task interference as an important direction for future work and will address it in the revised version of the paper.

Qualitative results for text-to-image generation: We will add the results for the text-to-image generation task in the main paper. These results are currently presented in the appendix (Figure 7).

[1] FreeFlux: Understanding and Exploiting Layer-Specific Roles in RoPE-Based MMDiT for Versatile Image Editing, Wei et al.