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

TL;DR

This work tackles the inefficiency of deploying multi-task image generation with diffusion models on devices. It introduces Multi-Task Upcycling (MTU), which splits FFN blocks into $N$ small experts and uses a router guided by task embeddings to dynamically fuse their outputs, while freezing the backbone. By focusing on task-specific input processing and layer norms, MTU achieves multi-task I2I capabilities with near the same computational cost as single-task baselines and outperforms both single-task and existing multi-task approaches across image editing, super-resolution, and inpainting. The results demonstrate strong on-device applicability and pave the way for scalable, efficient multi-task generative vision systems.

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.

Figures (10)

• Figure 1: A chatbot showcasing a potential use case of Multi-Task Upcycling. Our approach efficiently upcycles pre-trained text-to-image models, enabling them to perform multiple image generation tasks using a single backbone.
• Figure 2: We analyze the deviation between fine-tuned weights $\theta_{f}^{\tau}$ and pre-trained initialization $\theta_{p}$ across different layers in the LDM (i.e., $\Phi^{\tau} = || \theta_{f}^{\tau} - \theta_{p} ||$) and rank them accordingly. We present the average rank of these deviations across all tasks. The x-axis represents layer depth, while the y-axis indicates the component type. FFN layers show the highest deviation, suggesting they specialize in adapting to downstream tasks.
• 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, we replace the FFN layer in the pre-trained model with a set of smaller experts, which are dynamically combined using a router mechanism.
• Figure 4: Qualitative comparison of MTU based on SDv1.5 (left) and SDXL (right) with corresponding single-task baselines for Image Editing (IE) brooks2022instructpix2pix, Super Resolution (SR) 2022CVPRLDM, and Inpainting (IP). yildirim2023instinpaint
• 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).