Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities
Authors: Ruchika Chavhan
Abhinav Mehrotra
Malcolm Chadwick
Alberto Gil Couto Pimentel Ramos
Luca Morreale
Mehdi Noroozi
Sourav Bhattacharya
Venue: ICML 2025 poster (2025)
Keywords: multi-task learning
upcycling
diffusion models
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Multi-task Learning in Diffusion Models

Abstract: Text-to-image synthesis has witnessed remarkable advancements in recent years. Many attempts have been made to adopt text-to-image models to support multiple tasks. However, existing approaches typically require resource-intensive re-training or additional parameters to accommodate for the new tasks, which makes the model inefficient for on-device deployment. We propose *Multi-Task Upcycling* (MTU), a simple yet effective recipe that extends the capabilities of a pre-trained text-to-image diffusion model to support a variety of image-to-image generation tasks. MTU replaces Feed-Forward Network (FFN) layers in the diffusion model with smaller FFNs, referred to as *experts*, and combines them with a dynamic routing mechanism. To the best of our knowledge, MTU is the first multi-task diffusion modeling approach that seamlessly blends multi-tasking with on-device compatibility, by mitigating the issue of parameter inflation. We show that the performance of MTU is on par with the single-task fine-tuned diffusion models across several tasks including *image editing, super-resolution*, and *inpainting*, while maintaining similar latency and computational load (GFLOPs) as the single-task fine-tuned models.

Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities
Ruchika Chavhan'
Abhinav Mehrotra'
Malcolm Chadwick' Alberto Gil Couto Pimentel Ramos Mehdi Noroozi
Luca Morreale'
Sourav Bhattacharya'
Abstract
Esser et al., 2024) at the forefront. Thanks to the open- sourcing efforts, developers have the opportunity to fine- tune them for a variety of creative use-cases. The growing demand for generative Al applications has also contributed to the requirement of deploying the state-of-the-ar (SOTA) models on personal and edge devices to address data-prive aati--prvaay and the cost of cloud hosting.
Text-to-image synthesis has witnessed remarkable advancements in recent years. Many attempts have been made to adopt text-to-image models to support multiple tasks. However, existing ap- proaches typically require resource-intensive re- training or additional parameters to accommodate for the new tasks, which makes the model inef- ficient for on-device deployment. We propose Multi-Task Upcycling (MTU), a simple yet effec- tive recipe that extends the capabilities of a pre- trained text-to-image diffusion model to support a variety of image-to to--tto-iiii to-image generation tasks. MTU replaces Teed-Forwazd Network (FFN) layers in the diffusion model with smaller FFNs, referred to experts, and combines them with a dynamic routing mechanism. To the best of our knowl- edge, MTU is the first multi-task diffusion model- ing approach that seamlessly blends multi-tasking with on-device compatibility, by mitigating the issue of parameter inflation. We show that the per- formance of MTU is on par with the single-task fine-tuned diffusion models across several tasks including image editing, super-resolttion, and in- painting, while maintaining similar latency and computational load (GFLOPs) as the single-task fine-tuned models.
Great effort has been made to optimize these foundation models for edge deployment by making them smaller, faster, and resource-efficien (Zhao et al., 2025; Castells et al., 2024; Zhang et al., 2024b). These optimization strategies include distillation of models to reduce size (Xiang et al., 2024; Tang et al., 2023; Fang et al., 2023). reduction of the frequency of model calls (Salimans & Ho, 2022; Meng et al., 2023; Kang et al., 2025; Zhu et al., 2025; Yin et al., 2024; Noroozi et al., 2024), reduction of on-device memory and latency requirements through quantization (Li et al., 2023; He et al., 2023; Wang et al., 2024), and removal of computationally intensive operations (Zhao et al., 2023).
As SOTA diffusion models have demonstrated capabilities to support various use cases through fine-tuning, there is also a growing interest to develop single model to per- form multiple image-to-image (12I) tasks, like image edit- ing (Brooks et al., 2023), super-resolution (Moser et al., 2024), in/outgpainting (Corneanu et al., 2024; Wasserman et al., 2024). However, incorporating multiple tasks presents significant challenges. Some approaches adopt universal modeling (Ye & Xu, 2024; Zhang et al., 2024a; Bao et al., 2023) to learn joint probability distribution that unifies multiple modalities within common diffusion space. Other approaches rely on designing models with specialized com- ponents tailored to specific modalities or tasks (Xu et al., 2023; Tang et al., 2024). This allows the diffusion space to vary while partitioning the computational graph based on the task, offering modularity and flexibility. However, these methods significantly increase the model size and parameter count, making them computationally inefficient. Moreover, these approaches face scalability issue as the computational requirements grow significantly with the increase in the number of tasks and modalities. Thus, a significant gap remains efficiently adapting diffusion models to multi- ple tasks, while ensuring they are suitable for on-device deployment.
1. Introduction
Text-to-image (T2I) generation with diffusion models is rapidly gaining traction across diverse applications, with foundational models such as DALLE2 (Ramesh et al., 2022), MidJourney, Stable Diffusion (Rombach et al., 2022; Zhang et al., 2023; Podell et al., 2023; Stability AI, 2023; Lin et al., 2024), and Diffusion Transformers (DiT) (Peebles & Xie, 2022; Gao et al., 2024; Zhuo et al., 2024; Xie et al., 2024;
'Samsung Al Center, Cambridge. Correspondence to: Ruchika Chavhan <12.chavan @amsunsaamausa...
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Upcycling Text-to- Image Diffusion Models for Multi-Tass Capabilittess
woman walking down red cool
Make & enowy day
Rlemous the woman
Enhance the resotution of the image
612
512
512
Make the banznass
Enihence resclution Image
plctur of bawls with bananas It
green
remove
512 512
x 1024
1024
Figure 1. chatbot showcasing a potential use case Multi Task Upeycling, Our approach efficiently upcycles pre-trained text-to-image modcls, enabling them to perform multiple image generation tasks using a single backbone.
To bridge the gap, we introduce the concept of Multi-Task Upeycling (MTU) of T2I diffusion models. MTU transforms single-task T21 models into image-generation generalists, ie., a single model is capable of handling multiple tasks. The concept of upcycling is widely used in the field of Large Language Models (LLMs) transform dense pre-trained models into sparse Mixture-of-Experte (MoE) (He et al., 2024; Komatsuzaki et al., 2023; Jiang et al., 2025), but it has not yet been explored for diffusion models. Unlike pre- vious multi-task diffusion models, our approach avoids the need to significantly increase model parameters to support additional new tasks. Instead, it only requires retraining a few components with new multi-task data, while ensuring that the model remains both efficient and scalable for multi- task learning. MTU is particularly effective in the following two main scenarios: (i) extending on-device T2I diffu- sion model to support multiple tasks without increasing the computational requirement, (ii) obtaining a multi-task dif- fusion model with an existing T2I model acting as strong prior. We demonstrate potential use case of our proposed approach as a chatbot, illustrated in Figure 1, where a single model seamlessly handles multiple tasks such image edit- ing, inpainting, and super-resolution based on user requests.
literature, we convert a single FFN layer in pre-trained mod- els into number of smaller FFN experts. These experts are fine-tuned simultancously on multiple tasks while keeping the rest of the model frozen. A router mechanism dynam- ically learns to combine the outputs of these individual experts, guided by task-specific embeddings. We evaluate our method on latent diffusion models, such SDvI.5 and SDXL, across various imagett--image tasks, including im- age editing, super-resolution, and inpainting. Our results demonstrate that we can develop models that are iso-FLOP. i.e., having the same FI.OPs as their pre-trained counter- parts, while achieving performance comparable to single- task models. Our upcycled SDXL achieves FID of 3.9 on T21 generation, while maintaining the same computational cost of 1.54 TFLOPs as the pre-trained model.
2. Related Work
Multi-task Diffusson Models: There are two primary ap- proaches in this area. The first universal modeling (Ye & Xu, 2024; Zhang et al., 2024a; Bao et al., 2023; Chen et al., 2024), which aims to learn joint probability distribu- tion unifying multiple modalities within a shared diffusion space. For instance, instead of learning p(imageltext as in single-task diffusion models, it learns p(text, image), the joint probability distribution for text and image. Text-to- image generation can then performed by marginalizing the joint distribution. However, a key drawback of this
Our method is inspired by the empirical observation that when a T2I model is fine-tuned for a new task, the parame- Network (FFN) layers undergo a ters in the compared the other layers of diffusion significant shift, model. Building on this insight and inspired by upcycling
2

---

Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities
approach is its scalability as the entire model needs to be re-trained when new tasks are added. The second approach involves designing models with specialized components tailored to specific modalities or tasks (Xu et al., 2023; Tang et al., 2024). This enables the diffusion space to vary across tasks, while partitioning the computational graph for modularity and flexibility. Despite these advantages, this method significantly increases model size and parameter count, reducing efficiency. For both approaches, computa- tional requirements scale rapidly with the number of tasks and modalities, resulting in scalability challenges.
models, followed by a detailed outline of our multi-task upcycling approach.
3. Preliminaries
Latent Diffusion Models: Latent Diffusion Models (LDMs) (Rombach et al., 2022; Zhang et al., 2023) are based on Diffusion Models (DMs) (Ho al., 2020: Song et al., 2021), that learn to reverse a forward Markov pro- cess in which noise is incrementally added to input im- ages over multiple time steps 1 E [0,7]. We denote an RGB image by Xg E R where H and W corre- spond to the height and width of the image. An encoder & transforms the input image Xy into a latent representation Zo E Rexhxw where h and w represent the height and width of the downscaled encoded image, and e indicates the number of channels in latent space. During training, noisy latent 2, at time t is obtained from a real image's latent Zg by z, vazzo + v1 a4c, where c ~ NNO, I) and ar is a parameter that gradually decays over time. A denoiser fe(.) is then trained to predict the noise added to zi conditioned on the input text embedding cr. This en- ables the reconstruction of zu by subtracting the predicted noise from 24- To achieve this, the denoiser is trained to predict the noise by stochastically minimizing the objective: L(z, cr) = .x |||c - fo (z,, er, t]]-- A decoder D, then maps the denoised z. back to the pixel space. During inference, given a text prompt Cr, a noisy latent embed- ding 2r is sampled and iteratively denoised over T steps to produce zo. which is decoded into the final image. Typi- cally, the encoder and decoder are derived from pre-trained autoencoder that remains frozen during training.
Upeycling has been explored in various studies as a method for transform- ing trained dense models into Mixture-of- -oxxure-of-EEperts (MoF) frameworks (He et al., 2024; Komatsuzaki et al., 2023; Jiang et al., 2025). The concept of upeycling arose from the chal- lenges of training sparse MoE models from scratch, as such training is highly unstable and sensitive to hyperparameters. Upcycling offers a solution by starting with a pre-trained dense model, which is often readily available online, and transforming into an MoE model to enhance performance and capacity. At the core of any upcycling method, dis- cussed in (He et al., 2024; Komatsuzaki et al., 2023; Jiang et al., 2025), lies expert architecture design, initialization techniques, and routing strategies within the MoE layer.
Sparse Upeycling of Pre-trained Models:
The Feed Forward Networks (FFNs) within pre-trained LLMs, which are two-layer MLPs with hidden dimension dan are replaced by MoE layers. An MoE layer comprises N FFN experts denoted by {E1, Eg.... En} } and a router that learns assign tokens in the input appropriate perts. Let denote the hidden dimension of experts as dexpent- In sparse upcycling methods, the parameters of these experts are initialized using those of the pre-trained model. Various routing mechanisms have been utilized in different works (Lewis et al., 2021; Clark et al., 2022), with the Ex- pert Choice router (Zhou et al., 2022) and standard topK routing (Shazeer et al., 2017) being the most commonly used methods for computing the output of the MoE layer.
Fine-tuning LDMs for Image-to-Image Generation Tasks: The objective of image-to-image (12I) generation tasks is to transform an input image e₁ into a target image Curget based on an edit prompt cr. In 12I literature (Brooks et al., 2022), CT and C₁ are commonly referred as the text and image conditions, respectively. The target image and input image are encoded by an encoder & to obtain the latent representations Zunget and Ze, respectively. To train a diffusion model, noise is added to Zearget to obtain ze by VaiZuaga + vI VI- aic, where e ~ .N I) and a is a parameter that is scheduled similar to T2I diffusion models. A denoiser fo is then trained predict the noise added to a noisy input latent Zearget given the image con- dition ze and the text instruction cr. To achieve this, the image condition ze is concatenated with the noisy latents Zr, and the resulting tensor is provided as input to the de- noiser. To accommodate the additional channels introduced by the image conditioning, the first convolutional layer is modified to include extra input channels, while the rest of the architecture remains unchanged. The training process involves minimizing the following latent diffusion objective:
Many studies have proposed that increasing the number of experts a token is routed to, while simultaneously reducing the dimensions of each expert such that dexpent < drrx, can be a more efficient approach. This model is referred to as a fine- fine-grained MoE model (Krajewski et al., 2024) and the ra- tio G d deespa is termed as granularity. Reducing the dimensions of the experts decreases the F1.OPs per expert, which in turn permits an increase in the topK (the number of experts a token is routed to) proportional to the reduction in expert size, all while maintaining the overall FLOPs count.
We draw inspiration from the concept of fine-grained MoE models in the LLM literature and adapt it for multi-tasking in diffusion models. In the following sections, we provide the motivation for focussing the FFN blocks of diffusion
3

---

Upcycling Text-to-mmaee Diffusion Models for Multi-Task Capabilities
SDv1.5
SDXL
FFN
FFN
SAK
SAK
Parameter
5A-V
SA-Oul
SA-Oul
CA-Q
CA-
CAK
CAK
33
CAV-
CAV
11
13
15
26
28
30
32
34
36
3B
40
Layer
Layer
Figure analyze the deviation between fine-tuned weights 07 and pre-trained initializaaion Oy across different layers LDM (i.e., O' = ||0 0zll) and rank them accordingly. We present the average rank of these deviations across all tasks. The x-axis represents layer depth, while y-axis indicates component type. layers show the highest deviation, suggesting they specialize in adapting to downstream tasks.
Image Editing (IE)
Super Resolution (SR)
Inpainting (IP)
Feed-Forward Network (Fx) blocks. Within the SA and CA blocks, we further consider the Q. K, V, and output layer matrices. We rank $ &C and & for each layer across all tasks and then compute the average rank across I2I tasks. Figure 2 illustrates layer-wise average ranks for I2I tasks. Higher ranks signify greater deviation from the pre-trained initialization. Our experiments show that FFN layers deviate the most from the pre-trained weights to adapt to downstream 121 tasks. This is also demonstrated Fig- ure in the appendix, which presents measured deviation values.
I-T Dir Sim f
LPIPS
1-1 Dir Sim
SA
17.6
25.0
42.9
CA
16.9
27.5
40.1
Table Quantitative comparison of fine-tume finettunig different compo- nents (SA, CA, and FFNs) of the diffusion model over image-lo- image generation tasks. Fine-tuning FFNs leads to better perfor- mance image-to-image tasks.
L(2, Cr. 2e) Bxx.er a.. [lle- So (ze: Cr, Ze, II].. Pre- trained text-to-image models, such as Stable Diffusion, are commonly used as initialization to leverage their extensive -ia-tge-mmgee generation capabilities.
To further validate this observation, fine-tune each com- ponent of the attention block separately and compare their performance. Details on the metrics reported are in Section 6. As shown in Table 1, fine-tur fine-tuning the FFN layers in diffu- sion models consistently yields better results compared to tuning other components. In other words, FFNs specialize in solving downstream task, while the other parameters in the attention block learn more general features. In this paper, we build on these findings, and propose approach for Multi-Task Upcycling for Diffusson Models.
4. Motivation
Previous studies (Loshchilov & Hutter, 2017; Kirkpatrick et al., 2017) in the area of transfer learning and domain adaptation have shown that parameters undergoing signif- icant changes during fine-tuning are more relevant to the specific task, while those with minimal changes are either already well-aligned with the task due to pre-training or less critical for fine-tuning. Inspired from these findings, identify subset of parameters within the diffusion mod- els that show the highest deviations from the pre-trained initialization. In this experiment, we consider two Latent Diffusion Models, namely SDv1.5 and SDXL. We fine-tune them on three tasks: (i) image editing, (ii) super-resolution, and (iii) inpainting. See the appendix for further details on the dataset and training procedure in Section A.2.
5. Methodology: Multi-task Upeycling for Diffusion Models
In MTU, we segment the FFN component into experts such that their weighted combination solves specific tasks without the need of extra parameters. MTU comprises four key steps: (i) split the pre-trained model's FFNs into smaller FFN experts, (ii) design a router to dynamically combine the outputs of these experts, (iii) define spek-specific input processing layers, and (iv) design the loss function train the FFN experts and the router. Note that only the FFN experts, the router, and the task-specific input processing layers are trainable parameters. An overview of our method is presented in Figure 3.
Let fp be the weights of a pre-trained diffusion model, and 07 the fine-tuned weights for specific task т. The distance between the fine-tuned weights and the initialization is com- puted as t = ||05- - l, where || || is the Frobenius norm. For each task, we define $ for all LDM compo- nents: Self-Antention (b3a), Cross-Attentio (.))), and
Expert architecture: We replace each FFN block in the pre-

---

Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities
a)
Task ID
Task embedding
b)
Task wileeoodni
Prompt
Router
Attention
Layer Norms
Experts
E
D
NTO, 1)
LN
LN
LN4
XL
Task specific Input Layers
Task
Figure 3. (a) Overview: We introduce Multi-task Upcycling (MTU), a method for transforming a pre-trained text-to-image model to support multiple tasks. (b) In MTU, replace the FFN layer in the pre-trained model with set of smaller experts, which are dynamically combined using router mechanism.
trained denoiser architecture with N experts ( (Fi,... F ). each with dimension dexpent where / is the layer index. To ensure that the upcycled model maintains a parameter count similar to the pre-trained model, we constrain d'spent dNrNNN, where d'nN denotes the dimension of 1-th FFN block in the pre-trained model.
across tasks, and thus the task- layer normalization facilitates the learning of these distinct distributions in the upcycled model. We denote the set of all task specifcc con- volution layers as L.
Task-specific input layers: As outlined in Section single- task image-to-image models share the same architecture, dif- fering only in the input convolution layer, which accommo- dates additional channels from image conditioning. In our upcycled architecture, we introduce separate task-specific input convolution layers, denoted by v to handle the varying conditioning distributions for different tasks. We denote all the set of task-specific convolution layers as C.
Upcycling methods typically initialize the parameters of the experts from the pre-trained model. However, in our case, since the size of the FFN in the pre-trained model differs from the upcycled model, copying the dense MLP weights to the upcycled model structure becomes non-trivial. To address this, we segment the dense layer into N shards along the appropriate dimension and copy each shard into the corresponding expert. We denote the set of parameters in all the experts within the entire model as 0g.
Multi-task Loss:
For task т, let = represent the encoded images, zj the noisy latents at time step t. We feed 27 and the text prompt oj to the denoiser f. and task it to predict the noise e. Note that in the case of image conditioning, =f is the concatenation of the noisy latents and the encoded context image o7 along the channel dimension.
Router: Let T denote the set of tasks. For each task T, define learnable task embedding e, Rd We add router to layer 7 as: g(. 0%) {[++)) - RR where each expert is assigned weight based on e,. The output of the FFN block for an input x, corresponding to task т is computed as follows:
To reduce the computational burden, we pre-compute the expert weights using the task identifiers c,. We use these weights to combine the expert outputs as shown in Equa- tion 1. While the shared parameters of the model are kept frozen, the task-specific layer norms F and input layers Me, experts 0g and routers 0R are trained to optimize a multi-task objective with the loss defined as:
wi softmax(g(eri: 0'));;
EH (I+) = Σ wt (x+),
where wi represents the weight assigned to the i-th expert, and F, ( is the output of the i-th expert for input x, in layer of the model. A key advantage of this formulation that the task-specific weights can be pre-calculated with minimal computational overhead using the task identifier.
L= == - Eo..toofooof (V (2)
TET
6. Experimental Settings
Model Architectures: We evaluate our method on two Stable Diffusion- based models SDv1.5 (Rombach et al., 2022) and SDXL (Podell et al., 2023) consisting of 860M and 2.6B parameters in the denoiser component respectively. In both models, the denoiser is UNet composed of trans- former blocks with Self-Attention (SA). Cross-Attention (CA), Feed-Forward Networks (FFNs), and residual blocks.
Task-specific Layer Norms: Our empirical observation suggests that incorporating task-specific layer normaliza- tion step before each expert significantly improves model performance. The arrangement of these layer normalization steps is illustrated Figure 3(b). Our experiments show that FFN layers exhibit significantly different distributions
5

---

Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities
Mulli-task
Model
TFLOPs
Parameters
Text-to-Inaee (T20)
Image Editing (IB)
Super Resolution (SR)
Inpainting (IP)
4
Direction Similarity
LPIPS
Dircctional Similarity
T21 (Rombach et ol., 2022)
12.9
SD
IE (Brooks al. 2023)
15.4
(Rombach at. 2022)
0.67
860M
IP (Yikdirim ot al,, 2023)
46.5
(Xu et 2023)
0.87
10.1
14.2
Unidifluser (Bao et al., 2023)
0.83
952M
7.4
MTU (Ours)
0.68
869M
7.2
17.2
24.8
44.0
T2I (Podell al., 2023)
4.1
SDXL
IE (Brooks at al., 2023)
SR
1.53
2,6
17.3
26.9
43.2
MTU (Ours)
1.54
2.6
3,9
20.1
26.5
44.2
Table 2. Quantitative comparison of Multi-task Upcycling (MTU) against single-task and multi-task baselines. We consider = 4 for SDv1.5 and N 1 for SDXL. MTU consistently surpasses baselines while preserving computational efficiency.
For multi-task upcyeling, we consider all transformer blocks in SDv1.5 (16 blocks) and SDXL (70 blocks).
rate of le-4. During sampling, we perform denoising for 20 and 50 iterations for multi-task SDv1.5 and SDXL models respectively. More details are presented in Section A.2 in the appendix.
Each router network g(- (%)) is implemented as two-layer MLP with ReLU activation. As shown in Equation 1, we then apply a softmax function over the router's s predictions to obtain the weights assigned to each expert.
Metrics: We used the following metrics that are commonly considered for evaluating models train for specific tasks.
Downstream Image Synthesis Tasks:
We consider four tasks in our study, including Text-to-Image (T2I) genera- tion, Image Editing (IE), Super Resolution (SR), and Image Inpainting (IP). Since the MTU model is initialized with pre-trained T2I model, we include T2I as one of the tasks to ensure the multi-task model maintains its text-to-image generation capability.
T2I: We report Frechet Inception Distance (FID) on the test set of the COCO captions dataset. Lower values indicated better images.
SR: We measure Learned Perceptual Image Patch Similar- ity (LPIPS) (Zhang et al., 2018) between generated and ground truth images, where lower values indicate more similarity.
Datasets: Since each task requires different data configura- tions, we use the following datasets to train the MTU model.
Image Editing: We report Image-Text (I-T) Directional Similarity (Brooks et al., 2023), which quantifies how well the change in text captions aligns with correspond- ing edits. Let Impui and Tedited represent CLIP- extt-extracted features of input and edited images respectively (Rad- ford et al., 2021). Similarly, let Tedited and Timput de- note the CLIr-extracted text features for the input and edited descriptions. I-T Directional Similarity is defined as S(Taiiee T impuss Tedited Impu)). ). where S is the cosine similarity.
T2I: We use the COCO Captions dataset (Lin et al., 2015), a large collection of image-text pairs.
Image Editing: We use the dataset introduced in (Brooks
al., 2023), which provides input and target images along with corresponding editing instructions.
Super Resolution: We use the Real-ESRGAN dataset (Wang al., 2021), which consists of high-resolution im- ages. We generate corresponding low-resolution images by applying degradations and downscaling them by half. Image Inpainting: We use the dataset from (Yildirim et al., 2023), which provides multi-modal inpainting dataset designed for object removal based on text prompts.
Inpainting: We report Image-Imas Image (I-I) Directional Similarity, which measures alignment with the ground truth. Given CLIP features corresponding to ground truth image Igt, I-I Directional Similarity is defined as S(T at T Imput Todiedd - /mnpat) Higher values indicate bet- ter similarity to the ground truth.
More details can be found in Section A. in the Appendix.
We freeze all other layers and train only the FFN experts, routers, and task-specific layers, as de- scribed in Section 5. This results in training 158M pa- rameters for SDv1.5 1.5 and 1.5B parameters for SDXL. Both models are trained on 8x A100 GPUs for 100 epochs, with a batch size of 16 per GPU and image resolution of 512 x 512. SDXL is optimized using Adam W with a learning rate of 5e-5, while SDv1.5 is trained using Adam with learning
Training Details:
7. Results
Comparison with single-task and multi-task models We compare MTU against both single-task and multi-task baselines. For single-task baselines, we prioritize open- source models based on SDv1.5 or SDXL whenever avail- able. If no open- source models are available, we fine-tune the pre-trained T2I model on the 121 task following the
6

---

Upcycling Text-to- Diffusion Models for Multi- Capabilities
Input
Ground Truth
Single-Task
MIUU (Curs)
Input image
Ground Truth
Single-Takk
MTU (Ours)
Make night
IE- Mari
desert
IP. removve the boy
PP
Figure 4. Qualitative comparison of MTU based SDv1.5 (leff) and SDXL (right) with corresponding single-task baselines for Image Editing (IE) (Brooks et al., 2023), Super Resolution (SR) (Rombach et al., 2022), and Inpainting (IP). (Yildirim et al., 2023)
Super
IB
SR
IP
67
Eieen
I-T Dir Sim f
LPIPS
I-I Dir Sim f
LORA
8.3
43.3
31.4
Image Editing
Inpaintin
63
SDv1.5
IA3
11.2
39.3
34.2
Full - FT
15.4
38.0
46.5
MTU (N = 4)
17.2
24.8
44,0
Layer
Layer
LORA
13.1
31.8
28,9
SDXL
IA3 Full
8.4
51.0
30.6
Figure 5. Analysis for expert selection by the router. We show the expert weight distribution assigned by the router for SDv1.5 with four experts (N= = 4).
17.3
26.9
43.2
MTU (N = 1)
20.1
26.5
44.2
Table 3. Comparison with PEFT methods. MTU significantly out-
methodology outlined by (Brooks et al., 2023) in Section 3. Specifically, we consider the image editing SDv1.5 1.5 and SDXL models fine-tuned by (Brooks et al., 2023). For Super-Resolution (SR) and Inpainting, we could only com- pare SDv based models from (Rombach et al., 2022) and (Yildirim et al., 2023), respectively. Since no equivalent open-source versions are available for SDXL., fine-tuned model locally for both SR and Inpainting.
performs PEFT methods applied on FFN layers.
comparable to single-task models. Qualitative results are presented in Figure and Section A.3 in the appendix.
Comparison with PEFT methods
Parameter-efffcient fine-tuning (PEFT) methods are widely used to adapt pre-trained models to different tasks. Methods like LoRA (Hu et al., 2022) and IA3 (Liu et al., 2022) are lightweight, requiring only a small fraction of parameters to be fine-tuned within task-specific adapters. Since our analysis in Section 4 highlights that FFNs play a crucial role in enabling support for mmage-to-image tasks, we compare MTU with LoRA and IA3 for tuning the FFN layers. Note that applying IA3 only FFN blocks did not yield good results, we apply IA3 to the entire block including self- attention, cross- cross-attention and FFN layers. By comparing our approach with these PEFT methods, we evaluate how effectively our method performs in comparison to efficient fine-tuning techniques that primarily target FFN layers. Ta- ble 3 compares MTU with PEFT methods such as LoRA and IA3. For SDXL, method outperforms IA3 by approxi- mately 6 absolute points in image editing and 15 absolute points in super-resolution (SR). Overall, MTU consistently outperforms LoRA and IA3 across all tasks, demonstrating better multi-tase support compared to task sask-specific adapters.
Additionally, we use the Versatile Diffusion (VD) (Xu et al., 2023) and UniDiffuser (Bao et al., 2023) models as multi- task model baselines, as they are, to the best of our knowl- edge, the only models comparable to the MTU model. Note that these multi-task models were trained under entirely dif- ferent settings and do not cover all the tasks considered in this paper. Therefore, for a fair comparison, we evaluate UniDiffuser solely on the T2I task and Versatile Diffusion on the T21 and Image Editing tasks.
Table shows that our method consistently outperforms both single-task and multi-task baselines. Specifically, MTU based on SDv1.5 achieves scores of 17.2 in Image Edit- (SR), while single-task baselines lag by 3 and 14 points, respectively. Additionally, MTU based on SDv1.5 outperforms existing multi-task base- lines, further demonstrating its effectiveness. We also report the compute requirements for each model in GFLOPs, show- ing that our approach maintains computational efficiency
ing and 24.8 in Super-Resolune
Analysing router assignment for tasks
7

---

Upcycling Text-to- Text-to-Image Diffusion Models for Multi-Task Capabilities
# #experts N top-k
T21
IE
SR
IP
varying number of experts under this constraint. Our find- ings show that for both SDv1.5 and SDXI., increasing the experts to higher number (smaller experts) significantly degrades performance. For SDv1.5, the optimal number of experts is N = 4, while for SDXL., dividing the model into experts is suboptimal, as tuning the task-specific layer norms and FFN blocks alone provide sufficient multi-task support. Given that MTU for SDXL works best with a sin- gle expert, we conducted additional experiment where FENs were frozen, and only the layer norms preceding them were fine-tuned. However, this model failed to converge, implying that FFNs play a key role in learning 12I tasks. An alternative is to increase experts to or 16, but select only the top 4 per task, keeping active parameters equal to the original FFN. Table (Top) shows that this strategy underperforms compared using all experts directly.
FID -
I-T Dir Sim
LPIPS
I-1 Dir Sim
7.3
17.6
25.3
44.0
SDv1.
8.4
9.0
30.9
43.2
7.2
44.0
8
8.0
10.2
27.8
38.7
7.9
8.9
25,8
39.3
16
8.5
9.5
27.2
36.5
16
4
8.2
8.6
26.5
38.6
20.1
44.2
SDXL
2
10.5
10.4
30.4
39.9
12.3
11.8
30.5
38.6
12.8
12.4
31.6
34.5
8
4
13.6
12.8
31.5
32.7
31.7
16
18.2
10.9
33.2
31.9
Mexperts N
4 Params
T21
IF
SR
IP
FID
1-T Dir Sin
LPIPS
[-I Dir Sim
2,68
3.9
20.1
26.5
44.2
G=1
30,0
76.3
SDXLL (2.68)
5.38
3.8
19.1
25.8
49.3
0.25
2.68
12.3
113
30.5
38.5
Gx1
0.5
2.68
10.5
10.4
30.4
30.9
0.5
3.58
11.3
12.3
28.6
40.9
For (ii), we lift the parameter constraint from (i) and vary G from 1 N to I in Table 4 (Bottom). Increasing the number of experts enhances performance across most tasks, particu- larly super-resolution (SR) and image inpainting (IP). For instance, in the case of G = (experts matching the pre- trained FFN size), the MTU model with N 2 N 4 outperforms the N = 1 baseline (Table 2). When G < 1, the configuration with N 4 consistently outperforms N = 2 for the same value of G.
3:58
6.3
198
26.6
44.5
G>1
30.4
5.3
13.7
18.9
30.3
36.7
Table 4. Top: Performance of MTU on SDv1.5 and SDXL with varying number of experts N and parameter constraint by setting G= Bottom: Performance of MTU models (SDXL) without parameter constraint. We allow model size to scale with the number of experts to show that multi-task performance is linked to both number of experts and the capacity of each expert G.
For a fixed number of experts, the configuration with G = (expert size equal to the pre-trained FFN) consistently outperforms both G < (smaller experts) and G > 1 (larger experts), with G Interestingly, Table 4 (top) shows the opposite trend for SDv1.5, where smaller experts lead to better performance. This suggests that multi-task performance is influenced by both expert capacity (width) and model depth. In shallower models like SDv1.5, reducing expert width is advantageous, while in deeper architectures such as SDXL, maintaining the original expert dimensions yields better outcomes. This is likely because modifying expert size in deep network, without proportionally adjusting the rest of the architecture, disrupts training stability and degrades overall performance.
We analyze how the router distributes weights across the set of experts for each task to determine which experts special- ize in specific tasks. Figure 5 illustrates the expert weight distribution assigned by the router for SDv1.5 with N = 4. Our findings show that Text-to-Image Image (T2I) generation and Image Editing share three experts: ET. Eq. and E7 Mean- while, FI specializes in Super-Resolution (SR), and Ff is assigned greater importance for Image Inpainting. The weight distribution across experts is similar for T2I (trained on COCO) and Image Editing, as both tasks require strong compp--foloowig capabilities. In contrast, Inpainting pri- marily relies on object removal prompts, while SR operates without any textual conditioning, necessitating the use of different experts for these tasks.
yielding the weakest results.
In summary, multi-task performance depends on three in- terrelated factors: expert capacity, expert count, and model depth. With fixed expert size, increasing the number of experts generally boosts performance. Shallow architec- tures like SDv1.5 benefit from smaller experts, while deeper models such as SDXL perform best when experts match the pre-trained FFN size. This occurs because, in very deep networks, changing expert dimensions without adjusting the surrounding layers can destabilize training and degrade quality. Section 8, we provide practical guideline for se- lecting the optimal expert size and count based on available compute resources.
MTU with varying the number of experts and sizes We perform two ablation studies (i) Determine the optimal number of experts N while constraining the total parameter count of the MTU model as described in Section 5. (ii) Re- move this constraint by varying N and expert size dexpen to evaluate how the number of experts and their capacity affect multi-task performance. We denote the ratio of expert's 's dimension in MTU to the FFN dimension of the original model as G dexpen/ / /dmn-
1 such that the hidden dimension of For (i), we set -= N each expert decreases as we increase the number of experts. Table 4 (Top) presents the performance of our method with
Exploring task interference within MTU framework
8

---

Upcycling Text-to-Image Diffusion Models for Multi-Task Capabiliites
Tasks
T21
IE
SR
IP
generation model. Until now, PEFT methods are often the first choice for practitioners seeking to enable multi-task support due to their computational efficiency. However, our experiments with task-specific PEFT methods show that while they are lightweight, our MTU approach consistently outperforms them across all tasks while maintaining the same computational budget. We also demonstrate that MTU models can be tailored to any compute budget by adjusting expert sizes, making MTU an ideal solution for resource- constrained multi-task applications. Based on our findings, we propose an improved recipe for enabling multi-task ca- pabilities in pre-trained models based on a fixed compute budget. We recommend determining your total parameter budget first.
FID
I-T Dir Sin
LPIPS J
Dir Sim 1
T21
12.9
IE
15.4
SR
29.3
46.5
T2I, IP
17.9
30.0
SR. IP
36.2
30.2
T=2
IE, IP
20.6
52.6
T2L. SR
13.4
22.3
IE, SR
17.1
24.6
T21,
6.9
18.5
T21, JE. SR
7.0
18.0
22.7
SR.
16.8
25.1
35.8
T21. IE, IP
13.5
17.1
45.2
T 4 T2I, IE, SR, IP
7.2
17.2
24.8
44.0
Table 5. Exploring Task interference by training the MTU model across all combinations of the four tasks.
Start by adding task-specific input convolution layers to process additional image conditioning. Introduce task- specific layer norms before FFNs, and fine-tune the model without splitting any FFN into smaller experts (i.e., keeping single expert). As shown in Table 4, this simple approach performs well for SDXL but is less effective for SDv1.5.
Table 5 illustrates how different tasks interact and interfere with each other within the MTU framework. We train the MTU model across all combinations of the four tasks to ana- lyze how the inclusion of one task impacts the performance of another. We denote by I the number of tasks included in each experiment (with T = for single-task training and == 4 when all tasks are included). For intermediate values of T, we explicitly specify which tasks are involved.
a single FFN underperforms, conduct an ablation by varying expert count until the total parameters fit your budget. Under tighter limits, increase expert count and proportionally shrink each expert's hidden size. If resources allow larger model, keep each expert at the pre-trained FFN size and add more experts for optimal multi-task performance.
From T = 2 experiments, we first observe that image edit- ing (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 T21 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 IP dataset (InstructPix2pix) includes editing instructions that also involve object removal. We observe that T21 and IE are highly compatible, as training them to- gether improves EE performance-an an effect also highlighted in Figure 5, where these tasks select the same experts. For T= adding IP to SR and IE hurts SR performance, and training T2I, IE, and IP together also reduces T2I despite strong IP-IE compatibility. In the full T = 4 setting, T21, IE, and SR improve over their single-task baselines, while IP drops by 2 points.
9. Conclusions
In this work we introduced Multi-task Upeycling, simple yet effective approach to enhance pre-trained text-to-image models, such as SDv1.5 and SDXL, to support multiple image editing tasks. Unlike previous approaches, MTU is the first multi-task diffusion modeling framework that seamlessly integrates multi-task learning with on-device compatibility, ensuring efficiency without compromising performance. Our idea is based on an empirical observation that parameters in FFN layers in diffusion models deviate the most during task-specific fine-tuning. We use this ob- servation to propose splitting of existing FFN layer into smaller FFN experts, which are then combined with router network. We conduct an extensive evaluation across Text- to-Image, Image Editing, ing tasks, demonstrating superior performance compared to both single-task and multi-task baselines. We believe that our approach will open up new avenues of research in the rapidly evolving area of image synthesis and will con- tinue helping the efforts in making multi- vision models efficient for oe-iivice deployment.
Our results highlight notable task interference, and we rec- ommend that future research adapt proven multi-task learn- ing techniques, such as gradient conflict mitigation (Zhang et al., 2024c), dynamic loss balancing (Navon et al., 2022) to alleviate interference in diffusion models.
and Inpaint-
8. Discussions
In this section, we provide recommendations for upeycling a pre-trained text-to-image model into a multi-task image

---

Comeanu, C., Gadde, R., and Martinez, A. M. Latentpaint: Image inpainting in latent space with diffusion models. In 2024 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pp. 4331- 2024. doi: 10. (109//ACV57700.222244 Upcycling Text-to-mmage Diffusion Models for Multi-Task Capabilities
Impact Statement
2024. URL httpp:://ag semant wwe.mmtiiicoolle....... csschllrr org/ http:/cuuu::262479950080
This paper presents work in the field of Multi-Task im- age generation models. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Fang, G., Ma, X., and Wang, X. Structural pruning for diffusion models. In Advances in Neural Information Processing Systems, 2023.
Gao, P., Zhuo, L., Liu, C., Du, R., Luo, X., Qiu, L., Zhang, Y., et al. Lumina- 12x: Transforming text into any modal- ity, resolution, and duration via flow-based large diffusion transformers. arXiv preprint arXiv:2405.05945 2024.
References
Bao, F., Nie, S., Xue, K., Li, C., Pu, S., Wang, Y., Yue, G., Cao, Y., Su, H., and Zhu, J. One transformer fits all distributions in multi-modal diffusion at scale. In International Conference on Machine Learning, pp. 1692- 1717. PMLR, 2023.
He, E., Khattar, A., Prenger, R., Korthikanti, V., Yan, Z., Liu, T., Fan, S., Aithal, A., Shoeybi, M., and Catanzaro, B. Upcycling large language models into mixture of experts, 2024.
Brooks, T., Holynski, A., and Efros, A. A. Instruct- Pix2Pix: Learning to follow image editing instructions. In arXiv:2211.09800 2022.
He, Y.. Liu, L., Liu, J., Wu, W., Zhou, H., and Zhuang, B. PTQD: Accurate post-training quantization for dif- fusion models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: / //openreview net wettt//orummiidd/l1lp0e
Brooks, T., Holynski, A., and Efros, A. A. Instructpix2pix: Learning follow image editing instructions. In CVPR, 2023.
Castells, T., Song, H.-K., Piao, T., Choi, S., Kim, B.-K., Yim, H., Lee, C., Kim, J. G., and Kim, T.-H. Edgefusion: On-device text-to-image generation. CVPR24 First Work- shop on Efficient and On-Device Generation (EDGE), 2024.
Ho, and Salimans, T. Classifier-free diffusion guidance. NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2022.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. arXiv preprint arxiv:2006. 1239, 2020.
Chen, C., Ding, H., Sisman, B., Xu, Y., Xie, O., Yao, B. Z Tran, S. D., and Zeng, B. Diffusion models for multi-task generative modeling. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: //openreview net w cdvosidsii
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL https: open openreviie ew net forum? denzeeeeeeeee..
Clark, A., Casas, D. d. I., Guy, A., Mensch, A., Paganini, M., Hoffmann, J., Damoc, B., Hechtman, B., Cai, T., Borgeaud, S., Driessche, G. v. d., Rutherford, E., Henni- gan, T. Johnson, M., Millican, K., Cassirer, A., Jones, C., Buchatskaya, E.,, Budden, D., Sifre, L., Osindero, S., Vinyals, O., Rae, J., Elsen, E., Kavukcuoglu, K., and Simonyan, K. Unified scaling laws for routed language models. In Proceedings of the 39th International Confer- ence on Machine Learning. PMLR, 2022.
Hudson, D. A. and Manning, C. D. Gqa: A new dataset for real-world visual reasoning and compositional question answering. In 2019 IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2019.
Jiang, W., Lu, Y., Lin, H., Han, X., and Sun, L. Improved sparse upeycling for instruction tuning. In Proceedings of the 31st International Conference on Computational Linguistics, 2025.
Kang, M., Zhang, R., Barnes, C., Paris, S., Kwak, S., Park, J., Shechtman, E., Zhu, J.-Y., and Park, T. Distilling diffusion models into conditional gans. In European Conference on Computer Vision, pp. 428-447. Springer, 2025.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Muller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, Z., Lacey, K., Goodwin, A., Marek, Y., and Rom- bach, R. Scaling rectified flow transformers for high- resolution image synthesis. ArXiv, abs/2403.03206,
Kirkpatrick, J., Pascanu, R., Rabinowitz, N., Veness, J., Des- jardins, G., Rusu, A. A., Milan, K., Quan, J., Ramalho, T., trabska-Barwingg
A., Hassabis, D., Clopath, C., Kumaran, D., and Hadsell, R. Overcoming catastrophic forgetting in neural networks. Proceedings of the Na- tional Academy of Sciences, 2017.
10

---

Upcycling Texae Diffusion Models for Multi-Task Capabilities
Komatsuzaki, A., Puigcerver, J., Lee-Thorp. J., Ruiz, C. R., Mustafa, B., Ainslie, J., Tay, Y., Dehghani, M., and Houlsby, N. Sparse upcycling: Training mixture-of- experts from dense checkpoints, 2023.
Moser, B. B., Shanbhag, A. S., Raue, F., Frolov, S., Pala- cio, S., and Dengel, A. Diffusion models, image super- resolution, and everything: A survey. IEEE Transac- tions on Neural Networks and Learning Systems, Pp. 1-21, 2024. ISSN 2162-2388.
doi: 1001 3476671. URL http: //tx doi. www....or/....... ANNLS.20244 .00022.3476555 3476671.
Krajewski, J., Ludziejewski, J., Adamezewski, K., Pi' oro,o, M., Krutul, M., Antoniak, S., Ciebiera, K., Kr'ol, K., Odrzyg'o'zd'z T., Sankowski, P., Cy- gan, M., and Jaszczur, S. Scaling laws for fine- grained mixture of experts. ArXiv, abs/2402.07871, 2024. URL 267626982.
Navon, A., Shamsian, A., Achituve, I., Maron, H., Kawaguchi, K., Chechik, G., and Fetaya, E. Multi-task learning bargaining game, 2022.
wtt.pppp.smttttccsoo.ooooommoeee org/CozpusID /ttp//oppusI::22766662..33
Noroozi, M., Hadji, I., Martinez, B., Bulat, A., and Tz- imiropoulos, G. You only need one step: Fast super- resolution with stable diffusion via scale distillation, 2024. URL https: / arxiv wwweeivivo..rrreee bargg/bbs/2411. 17258.
Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettle- moyer, L. Base layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning. PMLR, 2021.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In arXiv:2212.097448 2022.
Li, H., Kadav, A., Durdanovic, I., Samet, H., and Graf, H. Pruning filters for efficient convnets. In In- ternational Conference on Learning Representations, 2017. URL
Podell, D., English, Z., Lacey, K., Blattmann, A.. Dockhorn, T., Müller, J., Penna, J., and Rombach, R. SDXL: Im- proving latent diffusion models for high-resolution image synthesis. In arXiv:2307.01952. 2023.
id=rJqFGTslq.
Li, X., Liu, Y., Lian, L., Yang, H., Dong, Z., Kang, D., Zhang, S., and Keutzer, K. Q-diffusion: Quantizing diffusion models. In Proceedings of the IEEE/CVF In- ternational Conference on Computer Vision (ICCV), pp. 17535-17545, October 2023.
Radford, A., Kim, J. W., Hallacy, C.,, Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision. arXiv, 2021.
Lin, S., Wang, A., and Yang, X. Progressive adversarial diffusion distillation. In 20Xiv22402.13929 2024.
SDXL-Lightning:
Ramesh, A., Prafulla, D., Alex, N., Casey, C., and Mark, Hierarchical text-conditional image generation with clip latents. 2022.
Lin, T.-Y., Maire, M., Belongie, S., Bourdev, L., Girshick, R., Hays, J., Perona, P., Ramanan, D., Zitnick, C. L., and Dollár, P. Microsoft coco: Common objects in context. arXiv preprint arXiv:2405.1 /9237, 2015.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, 2022.
Liu, H., Tam, D., Mohammed, M., Mohta, J., Huang, T., Bansal, M., and Raffel, C. Few-shot garameter-efficeen fine-tuning is better and cheaper than in-context learn- ing. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances Neural Information Processing Systems, 2022. URL hitps: /ttp://///eeeeeeeeeeeeeeeee / forum? d= JaPd.
Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. In International Confer- ence Learning Representations, 2022. URL https: / //openreview.o wttp://eorrrmmii............oope
Shazeer, N. M., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. V., Hinton, G. E., and Dean, J. Outrageously large neu- ral networks: The sparsely-gate gated mixture-of-experts layer. ArXiv, Abs/1701.06536 2017. URL ht /api. s semanticscholar org/ Corpus ID: 12462234.
Loshchilov. I. and Hutter, F. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017. URL https:// openreview net www.m/orummmmm Si 9Scxx.
Song, J., Meng. C., and Ermon, S. Denoising diffusion implicit models. In International Conference Learning Representationss 2021. URL /tttt::e///peeevviiee.
Meng, C., Rombach, R., Gao, R., Kingma, D., Ermon, S., Ho, J., and Salimans, T. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Stability AI. SDXL-Turbo. co/ co/stabilityaisskkll 2023.

---

Upcycling Text-to-Image Diffusion Models for Multi- Capabilities
Tang, S., Wang, X., Chen, H., Guan, C., Tang, Y., and zhu, W. Lightweight diffusion models with distillation-baser block neural architecture search, 2023. URL https: //rxivv... org/ abs/ /212. 04950.
of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 3886-38477, October 2023.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang. O. The unreasonable effectiveness of deep features as a perceptual metric. In CVPR, 2018.
Tang, Z., Yang. Z., Zhu, C., Zeng, M., and Bansal, M. Any- to-any generation via composable diffusion. Advances in Neural Information Processing Systems, 36, 2024.
Zhang, Y., Zhang, J., Yue, S., Lu, W., Ren, J., and Shen, X. Mobile generative ai: Opportunities and challenges. IEEE Wireless Communications, 31(4):58 64, 2024b.
Wang, H., Shang, Y., Yuan, Z., Wu, J., and Yan, Y. Quest: Low-bit diffusion model quantization via efficient selec- tive finetuning, 2024.
Zhang. 7., Shen, J., Cao, C., Dai, G., Zhou, S., Zhang, Q. Zhang, S., and Shutova, E. Proactive gradient con- flict mitigation in multi-task learning: A sparse training perspective. 2024c.
Wang. X., Xie, L.,, Dong, C., and Shan, Y. Real-esrgan: Training real-world blind super-resolution with pure syn- thetic data. Proceedings of the IEEE/CVF international conference on computer vision, pp. 1905-1991 2021.
Zhao, Y., Xu, Y., Xiao, Z., and Hou, T. MobileDiffusion: Subsecond texa-to-image generation on mobile devices. In arXi::231 2023.
Wasserman, N., Rotstein, N., Ganz, R., and Kimmel, R. Paint by inpaint: Learning to add image objects by re- moving them first, 2024.
Zhao, Y., Xu, Y., Xiao, Z., Jia, H., and Hou, T. Mobilediffu- sion: Instant text-to-image generation on mobile devices. In European Conference on Computer Vision (ECCV), pp. 225-242. Springer, 2025.
Xiang, Q., Zhang, M., Shang, Y., Wu, J., Yan, Y., and Nie, L. Dkdm: Data-free knowledge distillation for diffusion models with any architecture, 2024. URL https:// arxiv.
Zhou, Y.-Q., Lei. T., Liu, H.C.,, Du N., Huang. Y., Zhao, V., Dai, A. M., Chen, Z., Le, Q. V. and Laudon, J. Mixture-of-expers with expert choice routing. ArXiv, abs/2202.09368 2022. URL https: //api. www..smmatiicccoolee... org /CorpusII 247011948.
.rw. abs/ 03550.
Xie, E., Chen, J., Chen, J., Cai, H., Tang, H., Lin, Y., Zhang, Z. Li, M., Zhu, L., Lu, Y., and Han, S. Sana: Effi cient high-resolution image synthesis with linear diffu- sion transformer, 2024. URL ptattss://rxivv.. abs 2410 10629.
Zhu, Y., Liu, X., and Liu, Q. Slimflow: Training smaller one- step diffusion models with rectified flow. In European Conference on Computer Vision, pp. 342-359. Springer, 2025.
Xu, X., Wang, Z., Zhang, G., Wang. K., and Shi, H. Versatile diffusion: Text, images and variations all in one diffusion model. pp. 7754-7765, 2023.
Ye, H. and Xu, D. Diffusiommtl: Learning multi-task de- noising diffusion model from partially annotated data. In CVPR, 2024.
Zhuo, L.,, Du, R., Han, X., Li, Y., Liu, D., Huang, R., Liu, W., et al. Lumina-next: Making lumina-12x stronger and faster with nixt-dit. arXiv preprint 2406 2024.
Yildirim, A. B., Baday, V., Erdem, E., Erdem, A., and Dundar, A. Inst-inpaint: Instructing to remove objects with diffusion models, 2023.
Yin, T., Gharbi, M., Zhang, R., Shechtman, E., Durand, F., Freeman, W. T., and Park, T. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 6613-6623, 2024.
Zhang, J., Ye, H., Li, X., Wang, W., and Xu, Multi- task label discovery via hierarchical task tokens for partially annotated dense predictions. arXiv preprint arXiv:2411. 18823, 2024a.
Zhang, L., Rao, A., and Agrawala, M. Adding conditional control text-to-image diffusion models. In Proceedings
12

---

Upcycling Text-to-mmage Diffusion Models for Multi-Task Capabilities
A. Appendix
A.1. Dataset details
In this section, we provide details on the tasks and datasets used for training MTU models. Our study incorporates the following datasets, with the exact data splits outlined in Table 6.
Text-to-Image Generation (T2I): We utilize the COCO Captions dataset (Lin al., 2015), large-scale collection of image-text pairs. Each image in this dataset is accompanied by five captions, with one randomly selected during training.
Image Editing: We use the dataset introduced in (Brooks et al., 2023), which includes input-target image pairs along with corresponding editing instructions. Each input image has 4-5 target variations for a given edit instruction, with one randomly selected during training.
Super Resolution: We use the Real-ESRGAN dataset (Wang et al., 2021), which consists of high-resolution images. We generate corresponding low-resolution images by applying degradations like Poisson and Gaussian blur and downscaling them by half. For SR, we imput an empty string to the model.
Image Inpainting: We utilize the dataset from (Yildirim et al., 2023), multi-modal inpainting dataset designed for object removal based on text prompts. Built on the GQA dataset (Hudson & Manning. 2019), it leverages scene graphs to generate paired training data using state-of of-thh--rttt instance segmentation and inpainting techniques.
A.2. Training and Inference Details
In this section, we present training and inference details for both single-task and MTU models.
Training details: Both single-task and MTU models trained on 8x A100 GPUs for 100 epochs, with a batch size of 16 per GPU and image resolution of 512 x 512. SDXL is optimized using Adam with learning rate of 5e-5, while SDv1.5 is trained using Adam with a learning rate of le-4. For SDXL, we find that using a weight decay of 0.01 helps stabilize training.
Inference details: During sampling, we perform denoising for 20 iterations in multi-task SDv1.5 and 50 iterations in SDXL. For text-to-Image (T2I) generation, Image Editing, and Inpainting, we apply Classifier-Fres Guidance (CFG) (Ho & Salimans, 2022). However, for Super-Resolution (SR), CFG is used, it only processes an empty string as input. use guidance scale of 7.5 for SDv1.5 and 5.0 for SDXL. For Image Editing and Inpainting, we follow the CFG strategy from (Brooks et al., 2023), which employs dual guidance scales-one one for image and another for text. For SDv1.5, we use an image guidance scale of 1.6 and a text guidance scale of 7.5 for Image Editing, and 1.5 and 4.0 for Inpainting, respectively. For SDXL., we set the image guidance scale to 1.5 and the text guidance scale to 10.0 for Image Editing, while for Inpainting, we use 1.5 for image and 4.0 for text. For more details on the formulation of CFG, we direct the readers (Brooks et al., 2023).
For T2I generation,
Task T2I Image Editing Super Resolution Inpainting
Dataset
Train
Val
Test
COCO Captions (Li et al., 2017)
118287
5000
5000
InstructPix2pix (Brooks al., 2023)
281709
31301
2000
Real- ESRGAN(Wang al., 2021)
23744
100
100
GQA-Inpaint (Yildirim et al., 2023)
90089
10009
5553
Table 6. Training, validation and test splits of the datasets used in training MTU
A.3. Qualitative Results
We provide more qualitative results for each of the tasks considered in the paper from Figures 7, 8, 9, and 10.
13

---

Upcycling Text-to-Image Diffusion Models for Multi-Task Capabilities
Image Editing
Super Resolution
Inpainting
5A-0
540
58.K-
SA-
56-5
SDV1.5 1.5
SA-V
5A
5A.V
CA-K
CH-V
CA-
CAV
CA-
CAV
30
CArOut
CA-Out
CA-Out
13
33
15
11
15
11
15
25
FFN-
FR
FRN
Paramaler
56-0
56:0
SDXL
V
SA.
30
S4-Out
$6.
54.0ur
CA-0
CA.
CA.
CA-Out
CA-O
CA/
Layer
Layer
Layer
Figure 6. Distances between fine-tuned weights 8j and pre-trained initialization for SDv1.5 (top) and SDXL (bottom) across Image Editing, Super-Resolutin and Inpainting, Here x-axis corresponds to the Layer index and the y-axis corresponds to the distance between the parameters. For all tasks, FFNs exhibit the highest deviation from initialization, highlighting their crucial role in adapting downstream tasks. In SDXL, we focus only on the middle layers, where this deviation is most pronounced.
photography starry (Single
trails.
ngle
bewutiful bonsei
astronaut riding
COCO iSingle Tuski
MTU (Ours)
(single-task)
Realistic restaurant, paris, elegant, lights
Kearu Reeves portruit photo an wsite warrior (Ours)
COCO (SingleTan Tasik
MTU (Durs)
COCO Tagl)
Shuds
Figure 7. Qualitative results of MTU based SDv1.5 (first two rows) and SDXL (bottom row) for Text-to- Imnaee Generation.
14

---

Upcyeling Text-to-Image Diffusion Models for Multi-Tass Capabilities
SDv1.5
mustache he
make her green
Input Image
iSingle
were made gold
make photo
Input Image
(Single-Thss))
painting. Piz2Pie
ghost
image
Pa2ffix (SingleTTask)
SDXL.
make hor vampire
maki it winter scone
(Single-hhski
make lock a painting Ground
Tur into painting
image
Pie2Pix (Single-Tack)
image
Ground
a diorama
beach dlevett
Input
Ground Truth Fix2Pxx (Single/Tack)
Ground Truth
Fiz (Single-Tasi)
Figure 8. Qualitative results of MTU based on SDv1.5 and SDXL for Image Editing. In both models, our approach produces high-quality images with superior prompt adherence and faithful edits.
15

---

Upcyeling Text-to-Image Diffusion Models for Multi-Tann Capabilities
SDv1.5
Ground
ISingle-Taask
Input Ground Truch
Tasil
(Durst
Ground Truth
(SingleTask)
(Ours)
maago
Binghelnat
image
SDXL
Input
Ground Thuth
5R (Single-Tasic)
MTU [Guns
Ground Truth
SR [Single-Bakk)
MTU (Dum)
Input image
Ground Truth
(SingleTasii)
MTU (Qurs)
Input image
Ground Tuuh
(Snglleee)
MTU (Curs)
Figure 9. Qualitative results of MTU based SDv1.5 and SDXL for Super Resolution. Our method effectively restores high resolution images from low-resolution inputs that have been corrupted by image degradations, producing clearer and more detailed outputs.
16

---

Upcycling Text-to-Image Diffusion Models for Multi-Tass Capabilities
1.5
Input Image
Ground Tiuch
IP (SingleTTskk)
MTU (Ours)
Imputimage
Ground Truth
IP (SingleTask)
MIU (Ours)
the
the
Input image
Ground Truth
1iiggre Tark)
MTU (Ours)
Input Image
Ground Truth
(Single Task)
NTU (Curs)
the bottle the lelt ohe IP (iingle-Task)
the brown the center (Single-Bisk
Input image
Ground Thuth
MTU (Ours)
Input image
Ground Truth
NTU (Gurs)
SDXL
Groand Einngle
Groand That
Giogl
Cingle
Gound iStree
remove the silver mouse
remove the large woman
Input Image
Ground Truth
1Singic Task)
MTU (Ours!
Input Image
Ground Tuth
(Single Task)
MTU (Curs)
Figure 10. Qualitative results of MTU based SDv1.5 and SDXL for Image Inpainting, Our method demonstrates better object removal, generating clean inpainted images.
17


Reviews:
[{"review_id": "aelBOqo00o", "reviewer": "Reviewer_iNBi", "summary": "The paper proposes Multi-Task Upcycling (MTU), an extension of pretrained text-to-image models for multi-task on-device deployment. The authors first investigate the differences between fine-tuned weights and pretrained initial weights across different layers in LDM. Based on this observation, they split the Feed Forward Networks (FFNs) into smaller blocks and add components for task-specific processing. MTU achieves state-of-the-art performance for SD 1.5 and SDXL in downstream image synthesis tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We have answered your questions below. \n\n**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.\n\n 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. \n\n\n|       | FFN   | attn-Q  | attn-K  | attn.-V | attn-Out | attn2-Q | attn2-K | attn2-V | attn2-Out |\n|----|-------|---------|---------|---------|----------|---------|---------|---------|----------|\n| SR | **1.032** | 8.7e-3  | 9.5e-3  | 9.6e-3  | 0.344    | 7.4e-3  | 4.5e-3  | 4.6e-3  | 0.320    |\n| IE  | **0.856** | 6.7e-3  | 7.6e-3  | 5.1e-3  | 0.318    | 0.01    | 8.3e-3  | 4.5e-3  | 0.320    |\n\nTable A: Mode of the deviations between the fine-tuned and pre-trained model for SDv3.5.\n\n**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.\n\n**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.\n\n**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).\n\n[1] FreeFlux: Understanding and Exploiting Layer-Specific Roles in RoPE-Based MMDiT for Versatile Image Editing, Wei et al.", "claims_and_evidence": "1. The proposed multi-task learning method achieves better performance while preserving TFLOPs and parameters compared to other baselines.\n\n2. Inst-Inpaint can only remove objects, and the authors follow the dataset to train MTU. However, MTU has a lower I-I Directional Similarity than vanilla Inst-Inpaint in SD 1.5. Although the performance of other tasks is improved, inpainting is not. The authors must analyze why this happens.", "methods_and_evaluation_criteria": "Benchmark and evaluation metrics are reasonable.", "theoretical_claims": "This paper lacks any proof or theoretical claims.", "experimental_designs_or_analyses": "The experimental designs are valid.", "relation_to_broader_scientific_literature": "The proposed model is more applicable to real-world on-device applications since its number of parameters is almost the same as that of the base diffusion models, yet it can handle multiple tasks with a single model.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strengths**\n\n1. The analysis of each block (FFN, CA, SA) in LDMs is interesting.\n\n**Weaknesses**\n\n1. If the base model changes (e.g., LCM, SD3.5, Flux), the weight distribution of FFNs could differ.", "other_comments_or_suggestions": "1. It would be better to report the performance of the base model in Tab. 1.\n\n2. Please add qualitative comparisons for text-to-image generation in the main paper. Currently, the text-to-image generation results are only presented quantitatively.", "questions_for_authors": "1. The number of training samples is 118K for T2I, 281K for image editing, 23K for super-resolution, and 90K for image inpainting. Doesn’t this data imbalance affect the training results?\n\n2. How is the training batch composed? Does each training batch consist of different tasks or the same task? And what is its effect?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "biFwBwacQw", "reviewer": "Reviewer_zrKU", "summary": "The paper aims to achieve multi-tasking ability for a pre-trained T2I model, via a \"lightweight\" approach - MTU.\nThe paper starts from the insight that Feed-Forward Networks (FFNs) receive the most significant change when finetuned to a new task for a pre-trained T2I model. \nMTU splits original FFNs into smaller FFN experts and uses a dynamic router mechanism to adapt to new tasks.\n\n## update after rebuttal\nMy concerns are resolved after the rebuttal. I keep my original rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n1. **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. \n\n\n2. **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.\"\n\nPlease let us know if you have any more questions. Thank you.", "claims_and_evidence": "One of the most important conclusions is authors claim that FFNs in the pre-trained T2I models receive the most dramatic change when adapted to downstream tasks. The authors demonstrate this insight by visualizing the deviations from previous weights in Fig. 2, supporting that FFNs are specializing in adapting to new tasks.", "methods_and_evaluation_criteria": "The methods are intuitive based on the insights in Table 1 and Figure 2. Due to the specialization of FFNs, MTU uses them as expert and use a router as a kind of modulation block to better adapt to different new tasks. \n\nFor evaluation criteria, authors evaluate three different tasks, using FID, LPIPS, Directional similarity as metrics.  However, for tasks like Super-Resolution, only LPIPS is evaluated while some traditional metrics such as PSNR and SSIM are not considered. \n\nIt might be helpful to add a small user study to demonstrate the superiority of MTU.", "theoretical_claims": "No theoretical claims are found in the paper.", "experimental_designs_or_analyses": "The authors compare in different tasks to demonstrate the performance of MTU and its versatility. MTU achieves the best in most of the tasks. Given that MTU is the only approach handling 4 different tasks, it is promising the proposed techniques are generalizable to different downstream tasks. \n\nThe authors also show an analysis of expert in Figure 5, which is helpful to understand the mechanism of router assignment.", "relation_to_broader_scientific_literature": "N/A.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "Weaknesses:\n- Somehow all the image results shown in the paper are in a low quality. For example, in Figure 2, both cases show a low quality (e.g., blurry background, over-saturated color style) in the first place. Is it because of the pre-trained T2I models' poor performance? So in Figure 3 (wolf) and IE in Figure 4. This becomes a concern that the improvement brought by MTU may be trivial.", "other_comments_or_suggestions": "- In Fig. 3, it might be helpful to label each component, e.g., \"Input image\", \"Output image\", \"Encoder\", \"Decoder\". Also might be more clear if authors can exactly show what the prompt is in Fig. 3. \n- In Fig. 1, it might be helpful to indicate which T2I model is used here.", "questions_for_authors": "Please see my previous sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SpxmbTwwwE", "reviewer": "Reviewer_ptLD", "summary": "This paper suggests that upcycling the text-to-image (T2I) diffusion model when adapting it on the multi-task learning. In detail, when fine-tuning T2I diffusion mode with multiple tasks (e.g. SR, Inpainting, T2I generation, Image editing), they argue to utilize Mixture-of-Experts (MoE). For the motivation, they observe that the fine-tuning diffusion model on a specific task leads to significant changes in FFN parameters compared to the other modules such as attention layers. Thus, they divide FFN into multiple experts with a task-specific layer normalization. Experimental results show that the proposed Multi-Task Upcycling (MTU) achieves a better performance compared to other PEFT methods. Notably, they also show that the model w/ MTU surpasses a single expert model fine-tuned on a single task.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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}$.\n\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.\n\n**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.\n\nTable A: Ablation over number of experts for SDXL\n\n|          | Single-task |   | $N=1$   |   | $N=2$   | $N=2$   |   | $N=4$    | $N=4$    | $N=4$   |\n|----------|-------------|---|---------|---|---------|---------|---|----------|----------|---------|\n|          |             |   | $G=1$   |   | $G=0.5$ | $G=1$   |   | $G=0.25$ | $G=0.5$  | $G=1$   |\n| #Params  | 2.6B        |   | 2.6B    |   | 2.6B    | 3.5B    |   | 2.6B     | 3.5B     | 5.2B    |\n| T2I   (FID $\\downarrow$)   | 4.1         |   | 3.9     |   | 10.5    | 3.8     |   | 12.3     | 11.3     | 3.8     |\n| IE    (I-T Dir Sim $\\uparrow$)   | 17.3        |   | 20.1    |   | 10.4    | 20.0    |   | 11.8     | 12.3     | 19.1    |\n| SR  (LPIPS $\\downarrow$ )   | 26.9        |   | 26.5    |   | 30.4    | 26.3    |   | 30.5     | 28.6     | 25.8    |\n| IP   (I-I Dir Sim $\\uparrow$)    | 43.2        |   | 44.2    |   | 39.9    | 46.9    |   | 38.6     | 40.9     | 49.3    |\n\n**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.\n\n**Why do we need MTU?:** \n - 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.\n- 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. \n- 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.\n\nWe hope this addresses all of the reviewer's concerns, and we thank you for your feedback led to interesting experiments.", "claims_and_evidence": "- They claim that when applying multi-task learning to the diffusion model (i.e. training diffusion model on several image generation tasks), we need to use multi-task upcycling. \n\n- As evidence, they show the experiment that the parameters of FFN change mostly when fine-tuning the diffusion model to a specific task, and achieving the highest performance compared to the models only training with SA or CA. \n\n- Experimental results also support their claim, as diffusion model trained on multiple tasks with/ MTU surpasses a single expert model on a specific task.\n\n- One major concern is the experiment in Table. 4. It shows a single expert is enough in most cases, thus it degrades the significance of the author's claim, the importance of MoE architecture (instead, a just task-specific normalization is sufficient for multi-task learning in DM). Also, 2 experts in SDv1.5 reports much worse results compared to (1, 4) expert models, and I wonder why 2 experts have such worse performance.", "methods_and_evaluation_criteria": "- The proposed methods seem to make sense. They introduce a task-specific normalization and task-specific MoE and task-specific input layers. Input layer should be adaptive as the input image would differ as they are trained on different tasks. Also, as they show the FFN mostly changes during fine-tuning the diffusion model, introducing the task-specific normalization and MoE also seems to work well.\n\n- Evaluation metrics are common metrics for each task. For example, they use FID for text-to-image image generation and LPIPS for SR task. These are the most common metrics to measure the generative and perceptual SR capability.", "theoretical_claims": "There are no theoretical claims in this paper.", "experimental_designs_or_analyses": "- The main table (Table. 2) shows the superiority of the proposed method by suggesting the performance enhancement of the diffusion model trained with multi-tasks surpasses the performance of a task-specific model. \n\n- Figure 5. shows that the different experts are activated when the different types of tasks are given. \n\n- Table 3. shows the significance of the proposed method over the typical PEFT methods such as LoRA.", "relation_to_broader_scientific_literature": "The proposed method suggests the technique to enhance the fine-tuning of the diffusion model on multiple tasks. It would be extended to the other tasks like ControlNet, which has multiple types of inputs as conditions.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "As I mentioned, the major issue for me is that the significance of MoE.\nIn most cases, a single expert is sufficient to achieve the best performance.\nThen, why do we need to use MoE instead of just fine-tuning the FFN layer with a task-specific normalization?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pDaiRg4OQK", "reviewer": "Reviewer_FqBN", "summary": "This paper proposes Multi-Task Upcycling (MTU), a method to extend pre-trained text-to-image diffusion models for multiple image-to-image generation tasks without significantly increasing computational complexity or model parameters. The key idea is to replace the Feed-Forward Network (FFN) layers with smaller FFN experts combined through a dynamic routing mechanism guided by task-specific embeddings. This approach enables the model to support tasks like image editing, super-resolution, and inpainting, while maintaining comparable performance to single-task fine-tuned models with similar computational costs. The MTU method is inspired by Mixture-of-Experts (MoE) models commonly used in language models and adapts them for diffusion models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your insightful feedback. We answer your concerns below.\n\n**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.\n\nWe 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.\n\n**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.)\n\n| Tasks       | T2I       | IE          | SR       | IP          |\n|-------------|-----------|-------------|----------|-------------|\n|             | FID ↓     | I-T Dir Sim ↑ | LPIPS ↓ | I-I Dir Sim ↑ |\n| Single-task | 12.9      | 15.4        | 29.3     | 46.5        |\n| T2I-IP      | 17.9      | --          | --       | 30.0        |\n| T2I-IE      | 6.9       | 18.5        | --       | --          |\n| T2I-SR      | 13.4      | --          | 22.3     | --          |\n| SR-IP       | --        | --          | 36.2     | 30.2        |\n| SR-IE       | --        | 17.1        | 24.6     | --          |\n| IE-IP       | --        | 20.6        | --       | 52.6        |\n| T2I-IE-SR   | 7.0       | 18.0        | 22.7     | --          |\n| T2I-IE-IP   | 13.5      | 17.1        | --       | 45.2        |\n| IE-SR-IP    | --        | 16.8        | 25.1     | 35.8        |\n| All      | 7.2       | 17.2        | 24.8     | 44.0        |\n\nOur summary of task interference from the above table - \n\n- 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.\n- 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.\n- 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.\n\nWe 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.", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The method is well-motivated and evaluation criteria make sense for the image-to-image generation tasks. However, some qualitative examples could better highlight failure cases, especially in challenging scenarios.", "theoretical_claims": "No significant novel theoretical claims are made. The routing mechanism and expert combination follow the concept of fine-grained MoE models in the LLM literature and adapt it for multi-tasking in diffusion models.", "experimental_designs_or_analyses": "The experimental design is comprehensive, with ablation studies on the number of experts, routing mechanisms, and comparisons to PEFT methods. However, the paper could benefit from analyzing the optimal number of experts (N=4 for SDv1.5 and N=1 for SDXL).", "relation_to_broader_scientific_literature": "The paper is well-positioned in the multi-task diffusion models and Mixture-of-Experts (MoE) literature, and promotes the research on multi-task diffusion modeling methods.", "essential_references_not_discussed": "This paper includes the work related to the key contributions.", "other_strengths_and_weaknesses": "Strengths:\n1. The MTU framework presents a simple yet effective design that minimizes parameter inflation while preserving the model’s multi-task capabilities.\n2. The method is well-suited for on-device deployment, addressing an important challenge of diffusion models in resource-constrained environments.\n3. The paper includes comprehensive ablation studies.\n\nWeaknesses:\n1. The paper provides limited analysis of failure cases, which would help better understand the limitations of the proposed method across diverse tasks.\n2. The optimal number of experts is not rigorously analyzed. For example, the observation that MTU with 4 experts outperforms MTU with 1 expert on SDv1.5 lacks a clear explanation.\n3. The significant performance degradation of MTU on SDXL with 4 experts compared to 1 expert raises concerns about scalability. Further analysis is needed to clarify why increasing the number of experts negatively impacts performance on larger models.\n4. The paper does not explore task interference effects, which could potentially affect the model’s performance when handling multiple tasks simultaneously. Including an analysis of how the presence of certain tasks impacts others would strengthen the evaluation.", "other_comments_or_suggestions": "No other comments or suggestions.", "questions_for_authors": "See Other Strengths And Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper receives positive ratings of (3, 3, 3, 3). The rebuttal has resolved most of the reviewers' concerns, and the AC finds no reason to overturn the decisions of the reviewers. An acceptance is recommended.
Decision: Accept (poster)