UniCon: Unidirectional Information Flow for Effective Control of Large-Scale Diffusion Models
Authors: Fanghua Yu
Jinjin Gu
Jinfan Hu
Zheyuan Li
Chao Dong
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion Models
Lowlevel Vision
Conditional Generation
Primary area: generative models

Abstract: We introduce UniCon, a novel architecture designed to enhance control and efficiency in training adapters for large-scale diffusion models. Unlike existing methods that rely on bidirectional interaction between the diffusion model and control adapter, UniCon implements a unidirectional flow from the diffusion network to the adapter, allowing the adapter alone to generate the final output. UniCon reduces computational demands by eliminating the need for the diffusion model to compute and store gradients during adapter training. Our results indicate that UniCon reduces GPU memory usage by one-third and increases training speed by 2.3 times, while maintaining the same adapter parameter size. Additionally, without requiring extra computational resources, UniCon enables the training of adapters with double the parameter volume of existing ControlNets. In a series of image conditional generation tasks, UniCon has demonstrated precise responsiveness to control inputs and exceptional generation capabilities.

Published as a conference paper at ICLR 2025
UNICON: UNIDIRECTIONAL INFORMATION FLOW FOR EFFECTIVE CONTROL OF LARGE-SCALE DIFFU- SION MODELS
Fanghua Yu', Jinjin Gu², Jinfan Hu', Zheyuan Li', Chao Dong!
Shenzhen Institutes of Advanced Technology. Chinese Academy of Sciences
"The University of Sydney *Shanghai AI Laboratory
fanghuayu968 amail.. com, { rjf.hul. 2y 113, chao dong @siat. ac. cn,
jinjin. guu sydney. edu.ad au
(a) ControlNet
(b) UniCon
(c) Given the same trainable parameters.
2X
2.3K
Input 2t
Condition
Input Zt
Condition
28.43
22.07
1X
Trainable
Diff. Block
è Cti. Block
Diff. Block
Cti. Block
Fixed
Diff. Block
A Ctl. Block
Diff. Block
Cti. Block
VRAM
Spcedup
Trahiina
FID Score
Diff. Block
Block
Block
Cti. Block
(c) With the same training resources.
0.58
0.203
Blocks
Diff. Block
Diff. Block
Cti. Block
1X
Require
Gradient
Diff. Block
Diff. Block
Cti. Block
Diff. Block
Block
Cti. Block
Trainable Parameters
Cortrolataity (SSIM)
Quality (Man-IGA)
Generation
2t-1
21 -
ControlNe
Figure I: This figure illustrates the schematic comparison between our proposed UniCon and ControlNct. In UniCon, information flows unidirectionally from the diffusion model the adapter network, which directly outputs the results. This design is highly computationally efficient as does not require computing and storing gradients for the diffusion model. (c) displays results generated from downsampled images, and (d) shows outcomes based on depth maps. UniCon achieves improved performance while utilizing fewer resources.
ABSTRACT
We introduce UniCon, novel architecture designed to enhance control and ficiency in training adapters for large-scale diffusion models. Unlike existing methods that rely on bidirectional interaction between the diffusion model and con- trol adapter, UniCon implements unidirectional flow from the diffusion network to the adapter, allowing the adapter alone to generate the final output. UniCon reduces computational demands by eliminating the need for the diffusion model to compute and store gradients during adapter training. Our results indicate that UniCon reduces GPU memory usage by one-third and increases training speed by 2.3 times, while maintaining the same adapter parameter size. Additionally, without requiring extra computational resources, UniCon enables the training of adapters with double the parameter volume of existing ControlNets. In a series of image conditional generation tasks, UniCon has demonstrated precise responsiveness control inputs and exceptional generation capabilities.
INTRODUCTION
Diffusion generative models (Ho et al. 2020 Song et al. 2020: Nichol & Dhariwal 2021 Rombach
et al., 2022), with their exceptional generative effects and diversity, have significantly impacted
computer vision fields, such as creative design (Anciukevićius et al. 2023 Mittal et al.
2021 Cao
al. 2024), image processing (Wang et al.] 2023b Lin et al. 2023 Yang al. 2023
Yu al.
2024. and personalized content generation (Zhang et al. 2023 Mou et al. 2024 Li et al 2024
*Corresponding author (c-mail:

---

Published as a conference paper at ICLR 2025
2023
Zhao et al. 2024) This is largely attributed to their ability precisely control complex layouts, poses, shapes, and image conditions using both textual and visual prompts. These functions typically require training
additional adapter controller network associated with the diffusion model.
enabling control of the output during the inference process (Zhang et al. 2023; Mou et al.. 2024
7avadski al. 2023 Zhao et al. 2024). The most successful representative is ControlNet Zhang
et al. (023) for the Stable Diffussin (SD) model (Rombach et al., (202). which includes trainable copy of the SD U-Net encoder part and zero convolution layers as connectors. To achieve higher
quality generation and more precise control, the parameter size of diffusion models has also been significantly increased (Podell et al.] 2023 Lin et al., 2024] Sauer et al. (002).. Starting from the original SD model (Rombach et al., 2022) with 0.86 billion parameters, the advanced SD3 (Esser et al. 2024) has scaled up to 8 billion parameters. The latest models (Peebles & Xiel [023] [aao et al.
2023 Esser et al. 2024) replaced the commonly used U-Net architecture with transformer- based models. The benefits of scaling up have not yet reached their limit. We may expect larger diffusion
models and advanced transformer architectures to appear in the future (Xie ct al.) [023] Han et al.)
2023 Luo et al., 2023] Chen et al. 2024a). Under this trend, the limitations of existing control
adapters for diffusion models have become increasingly apparent.
There are three main problems: First, as the parameters of the diffusion model increase, the size of its adapter also needs to be expanded accordingly. Existing methods not only calculate the gradients of the adapters during training, but also calculate and store the gradients of the diffusion model, placing significant strain on computational resources. This typically implies nearly double the additional overhead. Second, while existing adapter designs exert control by modifying features in intermediate layers, they still process these modified features using the original parameters of the diffusion model. The generative capabilities of models with fixed parameters are limited when only the intermediate features are modified. Third, the design of existing control adapters assumes that the diffusion models use an encoder-decoder architecture (i.e. U-Net), which is inadequate for transformer-bas based diffusion models due to their inability to separate encoder and decoder components. Since existing adapters primarily implement control within the encoder part, their response to control signals lacks pixel-level precision, particularly problematic in nigh-precision tasks such as image restoration based on low-quality images
YYu et al. 2024).
In this paper, we introduce a novel design specifically tailored for training control adapters for the next generation of large-scale diffusion models. This design facilitates further scaling up of diffusion models and their control. Unlike the existing methods (Zhang et al., 2023] Zavadski et al.] 2023) that use adapters to intervene during the forward inference of diffusion models in a bidirectional manner, our approach allows information to flow unidirectionally from the diffusion network to the adapter. without flowing backward. The final output image will be generated by the adapter, NOT the diffusion model. Our method is therefore called UniCon. Figure T iillustrates the design difference between UniCon and previous methods. The advantage of UniCon is that during training, the diffusion model only needs to perform forward propagation and does not need to compute or store gradients. Additionally, UniCon can utilize all parameters and architectures of the diffusion model, not just those pertaining to the encoder, making it suitable for diffusion models with large-scale transformer architecture. Moreover, since the output processing is handled by the adapter, our design enables more precise control over generation based on conditions. UniCon avoids micro-interventions in the diffusion model architecture, thus it could further enhance the versatility.
We test the UniCon architecture across variety of conditional generative tasks, validating it on both the SD U-Net diffusion model and the DiT diffusion model. Notably, when applying the UniCon adapter to the hansformer-based DiT diffusion model, our approach saves half of the video memory (VRAM) usage while achieving 2.3X increase in training speed, without increasing the adapter's parameter size. It also delivers superior FID scores and condition fidelity in tasks involving image generation from downsampled images, see Figure (c). By employing our method in the generation task according Canny edge, the parameter size of the adapter can be doubled without additional computational resources. UniCon can further improve performance, where our method not only enhances the quality of generation but also ensures optimal controllability.
2
RELATED WORK
Diffusson generative models. The diffusion models (Sohl-Dickstein et al.) 2013) have been ex- plored across variety of generative tasks, such as image-to-image translation (Zhao al. (022).
2

---

Published as a conference paper at ICLR 2025
text-to-image synthesis (Rombach al., 2022
Saharia et al. 2022a Avrahami et al. 2022 Jiang
iet al. 2022).
image restoration (Saharia et al.
2022b; Daniels et al. 2021 Kawar et al. [0022] Lin
et al. 2023
Yu et al.) 2024: Tao al.] 2024)
image editing (Meng et al.
Avrahham et al.
2022
2023 image inpainting (Nichol et al.
2021: Lugmayr et al. 2022). etc. Model scaling up
has been an important means to improve the capabi of these models, and lot of efforts have been made, e.g.ä series of works proposed by Stability AI, including SD-2.1 (Rombach et al. 2022).
SD-XL (Podell al. 2023), SD-Cascade (Pernias et al. (023), and SD3 (Esser et al. 2024). In order to further expand and improve the capacity of the model, the large-scale transformer architecture is
also introduced for diffusion models, including DiT (Peebles & Xie 2023). U-ViT (Bao et al.) 2023)
SD3 (Esser et al., 2024) and the PixArt family (Chen et al. 2023 2024b a).
Controlled generation of diffusion models. Pre-trained large-scale diffusion models serve as foundational models that can be fine-tuned for a variety of downstream tasks. The fine-tuning process should be designed to mitigate issues such as overfitting, model collapse, and catastrophic forgetting.
Various fine-tuning methods have been developed across multiple fields, including HyperNetworks
(Alaluf et al.] [2022] Dinh et al. 2022). additive learning (Rosenfeld & Tsotsos] 2018 Mallya
et al. 2018: Mallya & Lazebnik] 2018 Serra et al. 2018), and Low-Lowaan Adaptation (LoRA) (Hu
ie al. 2021 J. For diffusion models, the most effective and commonly used approach for enhancing
generation quality and controllability involves training an additional adapter network (Houlsby et al.
2019| (Stcckland & Murray [2019] Li et al. 2022] Mou et al. (0244] Ju et al.) [024; Mo et al. [023))
Zhang et al. (Zhang et al. (202)) first introduced ControlNet, an architecture designed to incorporate
spatial conditioning controls into large, pre- rre-trained diffusion models. The ControlNet is connected
with zero-initialized convolution layers that progressively grow the parameters from zero and ensure that no harmful noise could affect the fine-tuning. Building on this, ControlNet-XS (Denis Zavadski & Rother (023) explored different sizes and architectural designs of ControlNet to enhance control over the image generation process in stable diffusion-basee models. Additionally, T21-Adapter (Mou) jet al. 2024) focused on aligning internal knowledge within T21 models with external control signals, while keeping the original large T2I models unchanged. This approach allows for the training of various adapters under different conditions, such as text and image, to achieve detailed control and editing capabilities in the color and structure of the generated images. Nni-ControlNet (Zhao et al. 2024) integrates various local and global controls into a single model, using only two additional adapters for fine-tuning on pre-trained diffusion models, thus avoiding the extensive costs of training restoration methods (Wang et al.
from scratch. Furthermore, surpassing existing diffusion-
2023b Lin et al. 2023 Yang al. 2023), SUPIR (Yu et al. 2024) leverages the generative prior and the potential of large diffusion model SDXL (Podell et al. 2023) to enable photo-realistic restoration
of severely degraded images through textual prompts. This illustrates the significant potential of
diffusion models in enhancing image quality and fidelity in controlled generation tasks.
3 METHOD
Preliminary and motivation. A typical diffusion model involves two processes: forward process that gradually adds a small amount of noise to the image, and a corresponding backward denoising process that recovers the input image by gradually removing the noise. Given a pre-trained diffusion model H(-), it can either have a U-Net-like structure, as in SD, or a transformer structure, as in the DiT model. When performing the backward generation process, the diffusion model takes the last denoised result 2n and time step t as well as the text prompt as input, and predicts the next result of denoising generation - H(22, p). The existing adapter methods directly modify the features X, of the intermediate layer of H during the denoising generation process to modify the generation results according to conditions. Assume we have an adapter C(.), usually may take 2,, 4, and p as inputs. Additionally, it may also take Xn as inputs (Zaaddsk et al. 2023). The output of the adapter C is a series of residual modified values X. of the intermediate layers Xn. In ControlNet, X, will be directly added to the corresponding intermediate layer of H. using element-wise addition generate a new representation X₁- = X + X,, to implement conditional generation.
This intuitive paradigm aligns with the common practice of fine-tuning large pre-trained diffusion models, yet it introduces several significant issues:
Firstly, this approach directly impacts the diffusion model; therefore. training the adapter involves calculating the gradients of the modified diffusion model and subsequently backpropagating these gradients to the adapter's trainable parameters. To achieve rapid and stable convergence, the
3

---

Published as a conference paper at ICLR 2025
ppeapp
procppt
4prumpt
X.
PulciEenoed
Encoder Block
Triksalle
ZooFT
ZeroConwMLP
Creader Block
Troknable
Zeaii
ZeroConw/LLP
Blok
Tranable Copy
Foooder Flock
Trainasle
(c) ZeroFT Connector
ST Blnkk
Triisanee Enpyy
e
Nid Block
3 Tarrable Copy
ZwrFT
Trsinsble
Fixed
301 Back
Traraaile Lapp
Dnoocer Book
A Trilnalee Copy
ZXIFT
Zarol
Figure 2: The UniCon design for both DiT and SD U-Net. We omitted some blocks in the SD U-Net due to the space limit.
Umiooder Block
Irainable Capy
DIT Blook
A
Trainable Dopy
Zarolt
ZorFFT
Hepppjection
ZeroF
Gutressaon
UniCon Design for DIT
(6) UniCon Design for SD U-Net
2
adapter's initial parameters are set as a trainable copy of the diffusion model. This can lead to training costs for the adapter surpassing those of the diffusion model itself. For example, when training ControlNet for DiT, when the gradient of ControlNet itself occupies about 18GB VRAM, the gradient brought by the DiT diffusion model needs to occupy 16GB VRAM, and this cost can be optimized. As the parameter number and resource demands of diffusion models continue to grow, training large adapters for large-scale models poses significant engineering challenges.
Secondly, while modifying the features of intermediate layers can directly alter the output, this technique still depends on pre-trained parameters to process these changes. When parameters fixed, adapting the model to new data by only changing the input of each layer restricts its generative capabilities. Alternative methods such as LoRA (Hu et al. 202T) or full-parameter fine-tuning permit direct modifications of parameters but also come with limitations. LoRA is constrained by its scale, and full -arammerreetrr tuning risks losing previously learned generative capabilities.
Finally, the existing design of control adapters, which often assumes a U-Net-like encoder-decoder diffusion architecture (Zhang et al., 20033 Mou et al. 2024 Zavadski et al.) 2023 Zhao et al. 2024) and focuses primarily on the encoder (Zhang el al., 2023). 3. does not suit for transformer- based diffusion models. This mismatch arises because difficutt distinctly separate the encoder and decoder components in such models. Indiscriminately controlling the entire transformer model escalates resource demands due to the necessity of computing and storing the diffusion model's gradients. Furthermore, since existing adapters mainly focus on the encoder, their response to control signals lacks pixel-level precision, particularly problematic high-precision tasks such as generation based on low- low-quality images.
Unidirectional information flow design paradigm. We propose design paradigm named UniCon to address these issues. For clearer understanding, readers are encouraged to refer to Figure Our approach ensures that the pre-trained diffusion model H. is used only for inference, thereby eliminating the need to compute and store its gradients. Information flows unidirectionally from H to the UniCon adapter Cur., positioning the adapter trainable decoder that processes information from the diffusion model's layers. Specifically, denote X,, Th1, Th2. as the output from the intermediate layers of H. The UniCon adapter receives these intermediate features along with the inputs of the condition c, time step f, and prompt p. The output of the adapter under this step is the denoising result 2 = CuC L, p, Xn.) This design establishes unidirectional information flow path from + H((- + X₁, + Cuc(.) + 21-1, ensuring that information does not revert to the pre- pre-trained diffusion model. By mirroring the architecture and parameters of the diffusion model in the adapter and using zero-inin pero-initialized connector for integration, we can ensure stable training.
This design resolves issues found in existing approaches. Firstly, the diffusion model only participates in forward inference, eliminating the need for computing and storing its gradients, which significantly reduces computational costs. Secondly, all adapter parameters are trainable, and the adapter directly delivers processed outputs, bypassing the fixed parameters of the diffusion model. At the same time, the well-trained diffusion model remains unchanged, ensuring the generative capabilities are less likely to be forgotten during fine-tuning. This enhances the generative capabilities and leads to improved generation quality. Moreover, as there is no need to intervene in the diffusion model - only to extract the computed features from intermediate layers solution is highly adaptable to
4

---

Published as a conference paper at ICLR 2025
"'Enodder" "Decoder" Fouuud
"Decoder" Part-Fooused Design
Peplace "Decoder"
Sklp-Layer CentrolNet
Block
00000
Tnerable Aciepler
Block
Blacks
Gusdert
Figure 3: Schematic representation of the five different variants we covered in our ablation studies.
Ti
FiS
Canny
Replace Diffusion "Decoder" Part
"Decoder -Part UniCon Design
Downsampled
Replace Diffusion "Decoder" Part
Decoder-Paa UniCon Design
Figure 4: The comparison between decoder-part-focused UniCon and replace diffusion decoder part. The results indicate that if the complete pre-trained diffusion model is not preserved, there is significant decline in generative capabilities. These two models are shown in Figure (d) and (e).
different diffusion architectures, from U-Net to transformer model with ease. Finally, the adapter no longer relies solely on the encoder or parts of the diffusion model but instead produces the output directly, further improving the preservation of control signals in tasks requiring high precision.
The proposed UniCon adapter designs. this work, we use DiT and SD U-Net to validate our UniCon adapter paradigm. Nevertheless, note that the UniCon design paradigm is widely applicable and can be employed across various diffusion models. FigureZ a) and Figure show the structural details when UniCon is applied to DIT and SD U-Net. The design of the UniCon adapter, applicable to both DiT and U-Net models, starts with trainable copy that includes all parameters and architectures of the diffusion model. In order connect the diffusion model to the adapter, connector module also required. For each output 2i of a DiT block or U-Net block, it introduces information from 21k to the corresponding location in the adapter. Zero-initialized convolutional layers (ZeroConv) or fully connected layers (ZeroMLP) are simple and stable choices for U-Net and DIT. This approach prevents unwanted noise from affecting features in trainable replicas the beginning of training. We additionally propose here new connector called vero-initialization Feature Transform (ZeroFT), ZeroFT not only utilizes element-wise addition but also incorporates element- wise multiplication and a shortcut connection, enhancing the integration and transfer of information. This setup is detailed Figure](c). c)). The parameters both two layers within ZeroFT whether convolutional fully-connected layers initialized to zero. According to experiments, ZeroFT is more effective than ZeroMLP/Conv in generating effects and controlling performance.
4 EXPERIMENTS
4.1
EXPERIMENT SETTINGS
Tasks and datasets. We employ five different conditions for image generation for testing: canny edge (Canny 1986), depth maps (Ranftl al. 2021), OpenPose (Cao et al. 2019), 4x down- sampling (SR) (Wang et al., 2018). and blurring (sigma=2) followed by 4x downsampling (de- blur-downsampling (Kong et al. 2022 Yu al. 2024h. These conditions can broadly be classified into two categories: (1) tasks that require the model to generate high-level semantic content, such as OpenPose, depth maps, and canny edge, which refer high-level control generation tasks;
5

---

Published as a conference paper at ICLR 2025
ControlNet Size
Task
Controllability
FIDJ
Clip-Scoret
Task
Connector
Controllabllity
FTD
Clip-Scoref
0.7724
Type
Canny (SSIM)
Docoder
0.5131
59.32
0.7507
ZoroMLP
0.5343
55.22
0.7612
Skip-I.ayer
0,4981
49.78
0.7776
Canuy (SSIM)
ShareAttn
0.5236
56.22
0.7606
Fall
0.5053
50.17
0.7818
ZeroFTT
0.5426
52.31
0.7696
Encoder
34.82
26.43
0.7996
ZeroMLP
35.67
22,99
0.8013
SR (PSNR)
Decoder
34,85
25.84
0.8013
(PSNR)
SharcAtta
35.55
23.03
0.8012
Skip-Layer
35,49
24.99
0.8009
ZotoFT
35.64
22.07
0.8025
0,8026
(b) The results of different connector designs used in the UniCon Adapter.
(a) The results of copy different parts in ControlNet.
Task
Adupter
Unidirectional
Coninollahility
FID
Clip-IQAf
Generation Quality MAN-IQAT
MUSIQI
Text Consistency Clip-Scoref
Skip-Layer
x
01.4983
49.78
0.6629
0.1978
66.05
0.7776
0.5078
$6.93
0.7561
Canny (SSIM)
Decoder
x
0.5131
59.32
0.6047
0.1621
62.51
0.7507
0.5343
55.22
0.1780
64.27
0.7612
Full
5
x
0,5053
0.5458
46.71
50.17
0.6397
0.6877
0.1867
0,2029
64.70
66.45
0.7889
0.7818
Decoder
34.85
25.84
0.6979
0,2325
68,26
0.8013
35.59
23.55
0.7036
0.2358
68.61
0.5018
(PSNR)
Full
36.53
23.04
0,7212
0,2609
69.91
0.8026
37.34
20.34
0.7251
0.2831
69.99
0.8022
(c) The effect of the proposed unidirectional information flow design.
Table Ablation study of different adapter designs. The diffusion model used in these experiments is DiT transformer-baat -based model. 1 indicates the larger the better and indicates the lower the better.
and (2) tasks that emphasize local generation and require higher precision in control, such as 4x downsampling and blurrig--donnanppling which we call low-level control generation. The images for training are selected from the LAION dataset (Schuhmann et al. (022). We randomly select million images with resolutions higher than 512 x 512. We center-crop and resize these images to 512 x 512, and pair them with the original text annotations from the LAION dataset to form image-tex pairs. All images are pre- pre-processed to compute the various conditions required for training, except for the Pose Condition. For the Pose Condition, we filter out images where no human pose was detected or key poses accounted for less than 30% of the body. We also exclude images with more than three people, as the pose detection error rate is higher in such cases.
Implementation details. In our experiments, we employed the PixelArt-a diffusion model as a representative example of the DiT model (Peebles & Xie (023) and StableDiffusion-2.1 (Rombach) al. 2022) as a representative example of the U-Net model. The training of the ControlNets was conducted using IDDPM (Nichol & Dhariwal 2021), maintaining the same noise schedule as the diffusion models. All experiments were performed on four NVIDIA-RTX A6000 GPUs, employing the AdamW W optimizer with learning rate of 2 10-4, The experiments were conducted with total batch size of 64 and a total of 100,000 training steps.
Evaluation. We sample 1,000 images from the LAION Dataset for testing. Consistent with the training setting, the pose condition in the test underwent the same image selection criteria. There are no duplicate images between the training and testing sets. To assess the controllability of different methods, we evaluate the alignment between the conditions of the generated images and the ground-trutt conditions. Different metrics were chosen based on the type of condition: peak signal-to-noise ratio (PSNR) for low-level conditions; structural similarity (SSIM) (Wang et al. (2004) for canny edge conditions; mean square error (MSE) for depth conditions; and mean average precision (mAP) for object keypoint similarity in pose conditions (Zhao et al., 2024). In addition to assessing controllability. we evaluated the quality of generation using metrics such as FID (Heusel et al., (017)) Clip-IQA (Wang et al. 2023a), MAN-IQA Yang et al. 2022). and MUSIQ (Ke) et al. 2021). Furthermore, the Clip-Score (Hessel al. (021) was used to evaluate the consistency between generated images and their corresponding text prompts, ensuring an integrated and thorough assessment of both image quality and fidelity conditioned inputs.

---

Published as a conference paper at ICLR 2025
A Black Skm Stretch Cotton Shet.
beautilu
DoT-CoobrolNet
DIT-ControNNe
DIT-UriCon
traditional
includes
snug,, wegetables
sprong
Canny
SD-T21
DDCoontroiNet
SD-UniCon
Canny
DIT-ConbrlNet
DIT-UniCon
Ranuncuuus
Contemperary Foyer
Depth Map
SD-T21
SD-ControlNet
SDDUnnCon
Depth Map
DIT-ControlNet
DiT-UniCon
Figure 5: Comparison different methods. We present the qualitative comparisons ControlNet (Zang et al. 2023) and T2I (Mou et al., 2024) with both Stable Diffusion (SD) (Rombach et al.) 2022 and Diffusion Transformer (DiT) (Peebles & Xie 2023).
4.2 ABLATION STUDY
The effect of copying "encoders" and copying "decoders" on diffusion control. Using Control- Net an example, we explored the influence of different parts of diffusion models controlled generation. The design of ControlNet involves making a trainable copy of the network's "encoder". Based on the DiT architecture, a trainable copy of and controlling the "decoder" part (the latter half) the network, as shown in Figure 3)). The second variant controls the entire network, covering both the encoder and decoder parts. To keep the parameter number constant, we also adopt the third variant of skipping and deleting every other block, shown in Figure Bc), referred as the "Skip-Layer" design.
further designed three variants. The first variant involves making
We test these variants on canny and SR tasks. The results are displayed in Table la First, comparing different designs focusing on the encoder and decoder, the results indicate that while focusing on the encoder leads to better image quality, focusing on the decoder enhances controllability. This probably because the understanding and generation of images and control signals are primarily conducted in the first half of the network, whereas the precise control of generation details occurs in the latter half. Specifically, for tasks such as SR that demand precise local control, a decoder- focused design might be more suitable. Given the distinct advantages of both encoder and decoder in controlling generation, we further explored the effects utilizing both simultaneously. The results from the Skip-Layer variant confirmed this approach, demonstrating improved generative performance and controllability at the same parameter number. This proves that for DiT, distinguishing between the encoder and decoder is not effective, and should leverage the capabilities of different parts of the entire diffusion model. The Full version, building upon the Skip-Layer, further increased the parameter number, and its enhanced generative performance and controllability also affirmed the importance of adapter capacity in controlling diffusion models.
Unidirectional information flow v.s. changing feature directly. The method proposed this paper differs significantly from existing approaches by introducing the concept of unidirectional information flow. Maintaining the same adapter architecture and initialization, we opted to output using adapters rather than directly changing the features of the diffusion model. We test various adapter variants mentioned above, and the results are presented in Table I Our results show that employing the unidirectional information flow for output substantially enhances performance,

---

Published as a conference paper at ICLR 2025
(a) GPU Memory Usage of DMT
[b) Traihing VRAW Cost of U-Net
[c] Training Time of DIT
(d) Training Time U-Net
40
Optimizer
Oprimizer
1.75
Forward Backward
Forwand Backmand
Gradient
Activatton
30
Activation
Gradient
0.5
Weight
Weight
30
Memroyy Cont (CB)
1.25
25
246B
Training Time
0.6
Monory
0.4
13
0.25
0.00
0.0
Controltin- Rull
Controiter... Deroder
UniCan- full
UniCon Decader
T2iadapter Controllin- Rull
Contremiet Decnder
UniK Full
UniCom Decader
Controllet- Rill
Controltan Decuder Unioo Decader
UniCon Rull
T2jAdapter Cammsllett- Fud
Controteet Deennder UniCon Decader
UnlCon Full
Figure 6: Comparison of training VRAM usage and training time of different adapter designs.
X4 Downsampled
SD-T21
SD-ControlNet
SD-UniCon
Ground Truth
X4 Downsampled
SD-T2I
ID-ContriiNet
SD-UniCon
Ground Truth
X4 Downsampled
DIT-ControlNet
DIT-UniCon-Haif
DIT-UniCon
Ground Truth
Figure 7: Comparison of different methods on conditional generation with low-level control inputs.
improving controllability and generative quality in both high-level and low-level tasks, whether using the diffusion model's "decoder" part or the full model as the adapter. For the "Skip-Layer" adapter, the unidirectional information flow design did not lead to performance gains. This is intuitively due to the skip-layer design compromising the output capability of the copied diffusion model. While the Skip-Layer design enhances the effectiveness of ControlNet, is not suitable for UniCon. This ablation study demonstrates that simple implementation of unidirectional information flow can significantly enhance both control effectiveness and generation quality.
Training cost. A major advantage of UniCon is its computational efficiency. As gradients do not need be computed and stored for the diffusion model, UniCon significantly reduces VRAM usage and speeds up training. We compare VRAM occupation and training speed between UniCon and ControlNet a single-GPU setup without any acceleration libraries. All experiments maintain consistency with standard training in terms of model parameters, diffusion settings. All necessary feature maps, such as text tokens, latent condition images, and latent ground- -ruth images, were pre-computed. We select the same 100 training batches, each with size of 16, and all speed tests were performed on the same NVIDIA-RTE A6000 GPU. All parameters and computations were performed in 6 precision. We detail the VRAM consumption as follows: (1) Weight: VRAM
8

---

Published as a conference paper at ICLR 2025
Diffusion
Model
Task
Adapter
Metric
Controllability
Valuc
FIDL
Clip-IQA
Generation Quality MAN-IQA"
MUSIQT
Text Consistency Clip-Score"
Conny
UniCon
ControlNeet
SSIM
SSIMI
0.4748
0.5458
46,71
51.52
0.6439
0.6577
0.1861
0.2029
66.45
65.24
0.7724
0.7889
ControlNet
MSEL
Depth
UniCon
MSE
82.56
51.49
0.6514
0.2002
65.51
0.7785
DIT
Pose
UniCon
ControlNet
mAPt
MAPt
0.4627
58.62
57.85
0.6332
0.6600
0.1819
0.1929
65.00
0.7545
ControlNet
PSNRT
34.82
26.43
0.7147
0.2459
69.37
0,7996
SR
UniCon-half
PSNRT
35.64
22.07
0.2042
0:2675
69.51
0,8025
UniCon
PSNR1
37.34
20.34
0.7251
0.2831
69.99
0,8022
Deblur+
SR
UniCon-half
UniCon
ControlNet
PSNR1
PSNRT
PSNR1
37.63
38.33
25.12
29.18
0.0998
0.7199
0.2563
0.2701
0.25117
69.31
69.72
0.8014
0.8004
ControlNet
SSIMI
0.4895
49,80
0.6683
0.2215
67.19
0.8168
Canny
T21-Adapter
0.2266
0.8155
UniCon
SSIMT
0.5570
47,11
0.6704
0,2355
67.99
0.8189
ControlNet
MSEL
54.30
0.6828
0.2262
67.90
0,8202
SD U--Net
Depth
T21-Adapter
UniCon
MSE]
MSEJ
85.00
87.72
53.45
55.09
0.6807
0.6906
0.2262
0.233
67.85
68.12
0.8214
0.8209
ControlNet
PSNR1
31.56
30.19
0.7373
0.3266
70.94
0.8044
T21-Adapter
PSNRT
18.94
48.20
0.6822
0:2795
70.91
0,78 178112
SR
UniCon-half
PSNR1
UniCon
PSNRT
35.69
22.80
0.7442
0.3271
70.48
0.8027
Table 2: Comparisons of different controllable diffusion models. indicates the larger the better and J indicates the lower the better.
used just for loading the model. (2) Activation: VRAM used for network forward propagation without gradient computation. (3) Gradient: VRAM used to store gradients. (4) Optimizer: VRAM used for updating Iearnable parameters of the model. We also record the time costs for network forward propagation (FP) and backward propagation (BP). FP time accounted for the time spent the diffusion model forward step, and BP time included the total time for updating gradients and optimizer parameters. Finally, we report the peak VRAM and average time costs over 100 iterations, with results presented in Figure It can be seen that UniCon significantly reduces VRAM usage compared ControlNet, saving nearly half the storage required for gradients, which are a major component of VRAM consumption during training. In terms of training time, because it eliminates the need compute gradients for the diffusion model, the time spent on BP is also nearly halved.
The choice the connector. We explored various methods for connecting diffusion models with adapters and found that different connectors lead to different outcomes. Based on the DiT diffusion model, we test three connector designs, including Z. which consistent with the NontrolNet (Zhang et al. 2023: Li et al. 2024), ShareAttention (Esser et al. 2024 Zhang & Agrawala 2024), which is suited for Transformer attention layers, and ZeroFT. The results these designs are displayed in Table lb] The results indicate that our proposed ZeroFT connector offers superior control and generative performance for both high-level and low-level tasks.
Retaining the original diffuusinn model preserves the previously learned generative capabilities. During the development of UniCon, natural question emerged: Given that the outputs are already handled by the adapter, is it necessary to retain the pre-trained diffusion model, or can we just fine-tune part of it? We compared the UniCon structure shown in Figure d)) with structure that discards part the diffusion model as shown in Figure e). Figure 4 presents the results of this comparison. The findings indicate that if the complete pre-trained diffusion model is not preserved, there is a significant decline in generative capabilities. This demonstrates that although the diffusion model in UniCon used only for inference, the generative capabilities it has learned can still be effectively utilized by the UniCon adapter.
4.3 COMPARISON
We also compare UniCon with existing adapter designs. For the DiT (Peebles & Xie 2023) diffusion models, we primarily conduct a direct comparison between UniCon and ControlNet (Zhang et al. 2023). For the SD U-Net diffusion models, compared UniCon with both ControlNet (Zhang et al. 2023) and T2I-Adapter (Mou et al. 2024). These comparisons of controllable diffusion models were conducted across all datasets.

---

Published as a conference paper at ICLR 2025
N655
Low Quality Input
SUPIR
SUPIR-UniCon
Low Quality Input
SUPIR
SUPIR-UniCon
Figure 8: Results of the SUPIR-UniCOn showcasing UniCon's effectiveness in image restoration.
Table 2 shows the results. It be seen that our proposed UniCon outperforms the ControlNet and 121-Adapter in all tasks. In terms of controllability, UniCon excels in maintaining control signals. UniCon also outperforms existing methods in image quality metrics, especially FID scores. It should be noted that although the T2I-Adapter method better than UniCon in some image quality metrics, the control effect of the T2I method is not good. Combining our ablation studies and comparative experiments, we also found that the scale of the adapter impacts the final outcomes. UniCon- Half, with only half the parameters, performs notably worse than the full-parameter UniCon but still performs better than ControlNet with comparable parameter number Notably, even the full-parameter UniCon has lower training computational cost than ControlNet.
Figure [ shows some visual comparisons of high-level conditional generation tasks. As one can see, our proposed UniCon excels generating images with superior detail structure, such as more intricate facial features and fewer artifacts and tears. Additionally, UniCon surpasses other methods in precision control, as demonstrated in the comparison of the second row's first set. Although the sushi rolls are small in scale and complex in structure, UniCon accurately generates sushi rolls the correct location, faithfully following the control conditions, unlike other methods that fail to strictly adhere to these conditions, and their generative capabilities inferior to UniCon. Figure 7 further illustrates UniCon's performance in low-level control generation tasks, which demand high fidelity. Our method only maintains high-quality generation but also strictly adheres the controls of the input images, effectively generating even the smallest structures. UniCon demonstrates better results in tasks such as image restoration and super-resolution based on diffusion models, showcasing broad application prospects. More results can be found in the Appendix.
UniCon's 's efficient resource utilization, enhanced control, and high-fidelity performance make particularly well-suited for low-level vision and image processing tasks. SUPIR (Yu et al. (2024), the first image restoration model based on large-scale diffusion models, currently represents the state of the art in image restoration. However, SUPIR relies on ControlNet, and further scaling under this approach demands immense computational resources, making large-scale expansion impractical. UniCon effectively addresses this limitation. Building the SUPIR framework, trained new model using SD3 and UniCon. Some of the results, shown in Figure g highlight UniCon's potential for broader applications.
CONCLUSION
This paper presents UniCon, a novel approach tailored for controlling large-scale diffusion models. UniCon leverages unidirectional flow from the diffusion network to the adapter, simplifying the computational process by eliminating the need to compute and store gradients of the diffusion model. UniCon significantly reduces the VRAM requirements and enhances training speeds while maintaining high fidelity in generated images.
Note that the UniCon-Encon Encoder design is ineffective. UniCon operation requires information to flow unidirectionally from the diffusion model to the adapter. Duplicating the encoder leaves no adapter to process information in the decoder, compromising generation quality and potentially failing produce images.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGE
This work was supported by the National Natural Science Foundation of China (Grant No. 62276251) and the Joint Lab of CAS-HK.
REFERENCES
Yuval Alaluf, Omer Tov, Ron Mokady, Rinon Gal, and Amit Bermano. Hyperstyle: Stylegan inversion with hypemeetwrks for real image editing. In Proceedings of the IEEE/CVF conference on computer Vision and pattern recognition, pp. 18511-18521, 2022.
Titas Anciukevičius, Zexiang Xu, Matthew Fisher, Paul Henderson, Hakan Bilen, Niloy J. Mitra, and Paul Guerrero. Renderdiffusion: Image diffusion for 3d reconstruction, inpainting and generation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR). pp. 12608 12618, June 2023.
Omri Avrahami, Dani Lischinski, and Ohad Fried. Blended diffusion text-driven editing of natural images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18208-18216 18218, 2022.
Omri Avrahami, Ohad Fried, and Dani Lischinski. Blended latent diffusion. ACM Transactions on Graphics (TOG), 12(4):1-11, 2023.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backhone for diffusion models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 22669-226679, 2023.
John Canny. A computational approach to edge detection. IEEE Transactions pattern analysis and machine intelligence, (6):679-999) 698, 1986.
Yu Cao, Xiangqiao Meng, PY Mok, Tong- Yee Lee, Xueting Liu, and Ping Li. Animediffusion: Anime diffusion colorization. IEEE Transactions on Visualization and Computer Graphics. 2024.
Z. Cao, G. Hidalgo Martinez, T. Simon, S. Wei, and Y. A. Sheikh. Openpose: Realtime multi-person 2d pose estimation using part affinity fields. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2019.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthosis, 2023.
Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Weak-to-strowe training of diffusion transformer for 4k text-to-image generation, 2024a.
Junsong Chen, Yue Wu, Simian Luo, Enze Xie, Sayak Paul, Ping Luo. Hang Zhao, and Zhenguo Li. Pixart-6: Fast and controllable image generation with latent consistency models, 2024b.
Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances neural information processing systems, 34: 34:12955-12940 2021.
Friidrihr Feiden Denis Zavadski and Carsten Rother. Controlnet-xs: Designing efficient and effective architecture for controlling text-to-image diffusion models. 2023.
Tan M Dinh, Anh Tuan Tran, Rang Nguyen, and Binh-Son Hua. Hyperinverter: Improving stylegan inversion via hypernetwork. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 11389-11399 11398, 2022.
Patrick Esser, Sumith Kulal, Andreas Blautmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for nigh-resolution image synthesis. arXiv preprint 2403.2240300200 2024.
Ligong Han, Yinxiao Li, Han Zhang. Peyman Milanfar, Dimitris Metaxas, and Feng Yang. Svdiff:: Compact parameter space for diffusson fine-tuning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7323-7334, 2023.
Jack Hessel, Ari Holtzman, Maxwell Forhes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning. arXiv preprint arXiv:2104.087/18 2021.
11

---

Published as a conference paper at ICLR 2025
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXi::200.115998 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840 6851, 2020.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Ges- mundo, Mona Attariyan, and Sylvain Gelly. Parameter-effiient transfer learning for nlp. In International conference on machine learning. pp. 2790-2799. PMLR, 2019.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, arXiv preprint arXiv:2166.09663 2021.
Yuming Jiang, Shuai Yang. Haonan Qiu, Wayne Wu, Chen Change Loy, and Ziwei Liu. Text2human: Text-driven controllable human image generation. ACM Transactions on Graphics (TOG), (1): 2022.
Xuan Ju, Xian Liu, Xintao Wang, Yuxuan Bian, Ying Shan, and Qiang Xu. Brushnet: plug-and-play image inpainting model with decomposed dual-branch diffusion. arXiv preprint arXiv:2403.06977 2024.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances Neural Information Processing Systems, 35:23593-23606, 2022.
Junjie Ke, Qifei Wang, Yilin Wang, Peyman Milanfar, and Feng Yang. Musiq: Multi-scale image quality
transformer. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 5148-5157, 2021.
Xiangtao Kong, Xina Liu, Jinjin Gu, Qiao, and Chao Dong. Reflash dropout image super-resolution In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 6002-6012. 2022.
Ming Li, Taojiannan Yang. Huafeng Kuang, Jie Wu, Zhaoning Wang. Xuefeng Xiao, and Chen Chen, Control-
net++: Improving conditional controls with efficient consistency feedback. arXiv preprint arXiv:244.007987, 2024.
Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer backbones for object detection. In European Conference Computer Vision, pp. 280-28 296. Springer, 2022.
Yuheng Li, Haotian Liu, Qingyang Wu, Fangzhou Mu, Jianwei Yang. Jianfeng Gao, Chunyuan Li, and Yong Jac Lee. Gligen: Open-set grounded text-to-image generation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 22511-22521, 2023.
Shanchuan Lin, Anran Wang, and Xiao Yang. Sdxl-lightnin:: Progressive adversarial diffusion distillation. arXiv preprint arXiv:2402. 13999, 2024.
Xinqi Lin, Jingwen He, Ziyan Chen, Zhaoyang Lyu, Ben Fei, Bo Dai, Wanli Ouyang, Yu Qiao, and Chao Dong. Difibir: Towards blind image restoration with generative diffusion prior. arXiv preprint arXiv:2308. 15070. 2023.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Lue Van Gool. Repaint: Inpainting using denoising diffusion probabillitic models. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 111,,
Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu, Patrick von Platen, Apolinário Passos, Longbo Huang, Jian Li, and Hang Zhao. Lem-lora: A universal stable-diffusion acceleration module. arXiv preprint arXiv:2311.055. 2023.
Anun Mallya and Svetlana Lazebnik. Packnet: Adding multiple tasks to a single network by itcrative pruning. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pp. 7765 7773, 2018.
Arun Mallya, Dillon Davis, and Svetlana Lazebuik. Piggyback: Adapting a single network to multiple tasks by learning to mask weights. In Proceedings of the European conference on computer vision (ECCV), pp. 67-82, 2018.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. arXiv preprint arXtv:2008.0000 2021.
12

---

Published as a conference paper at ICLR 2025
Gautam Mittal, Jesse Engel. Curtis Hawthorne, and Ian Simon. Symbolic music generation with diffusion models. arXiv preprint arXiv:2103. 16091, 2021.
Sicheng Mo, Fangzhou Mu, Kuan Heng Lin, Yanli Liu, Bochen Guan, Yin Li, and Bolei Zhou. Freecontrol: Training-free free spatial control of any text-to-image diffusion model with any condition. arXiv preprint arXiv:2312.07586 2023.
Chong Mou, Xintao Wang. Liangbin Xie, Yanze Wu, Jian Zhang, Zhongang Qi, and Ying Shan. Tži-adapter: Learning adapters dig more controllable ability for text-to-image diffusion models. In Proceedings the AAAI Conference on Artificial Intelligence, volume 38, pp. 4296-4304, 2024.
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, llya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint urXi::2112.1777 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
William Peebles and Saining Xie. Scalable diffusion models with transformers. Proceedings of IEEE/CVF International Conference on Computer Vision, pp. 4205- 2023.
Bohao Peng. Jian Wang. Yuechen Zhang. Wenbo Li, Ming-MMnaaat Yang. and Jiaya Jia. Controlnext: Powerful and efficient control for image and vidco generation. arXiv preprint arXiv:2408.066 06070, 2024.
Pablo Pernias, Dominic Rampas, Mats L. Richter, Christopher J. Pal, and Mare Aubreville. Wuerstchen: An efficient architecture for large-scale text-to-image diffusion models, 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.0093 2023.
René Ranftl, Alexey Bachkovskiy, and Vladlen Koltun. Vision transformers for dense prediction. In Proceedings of the IEEE/CVP international conference computer vision, pp. 12179-12188, 2021.
Robin Rombach, Andreas Blatmmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10684-10695, 2022.
Amir Rosenfeld and John K Tsotsos. Incremental learning through deep adaptation. JEEE transactions on pattern analysis and machine intelligence, 42(3):651 663, 2018
Chitwan Saharia, William Chan. Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Chasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances neural information processing systems, 35:36479- 36494, 2022a.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE transactions on pattern analysis and machine intelligence, 45 (4):4713 4726, 2022b.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint arX0422311.17042, 2023.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, el al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35: 25278 25294, 2022.
Joan Serra, Didac Suris, Marius Miron, and Alexandros Karatzoglou. Overcoming catastrophic forgetting with hard attention the task. International conference on machine learning, pp. 4548 4557, PMLR, 2018
Jascha Sohl-Dickstee Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint 000.v22000.02502, 2020.
13

---

Published as a conference paper at ICLR 2025
Asa Cooper Stickland and lain Murray, Bert and pals: Projected attention layers for efficient adaptation in multi-task learning. International Conference on Machine Learning, pp. 5986-59669 PMLR, 2019.
Keda Tao, Jinjin Gu, Yulun Zhang. Xiucheng Wang, and Nan Cheng. Overcoming false illusions in real-world face restoration with multi moati-modal guided diffusion model. arXiv preprint Xiiv:2410.0416.. 2024.
Jianyi Wang. Kelvin CK Chan, and Chen Change Loy. Exploring clip for assessing the look and feel of images. In Proceedings of AAAI Conference Artificial Intelligence, volume 37. pp. 2555 2563, 2023a.
Jianyi Wang, Zongsheng Yue. Shangchen Zhou, Kelvin CK Chan, and Chen Change Loy. Exploiting diffusion prior for real-world image super-resolution. arXiv preprini arXiv:2305 230002205.00015. 2023b.
Xintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Yu Qiao, and Chen Change Loy. Esrgan: Enhanced resouttenolttion generative adversarial networks. Proceedings of the European conference on computer vision (ECCV) workshops, pp. 0-0, 2018.
Zhou Wang, Alan Bovik, Hamid R Sheikh, and Eero Simoncelli. Image quality assessment: from emor visibility to structural similarity. IEEE transactions on image processing, 13(4);600- 612, 2004.
Enze Xie, Lewei Yao, Han Shi. Zhili Liu. Daquan Zhou, Zhaoqiang Liu, Jiawei Li, and Zhenguo Li. Difffit: Unlocking transferability of large diffusion models via simple parameter-efficient fine-tuning, In Proceedings of the IEEE/CVF International Conference Computer Vision. pp- 4230-4239, 2023.
Sidi Yang. Tianhe Wu, Shuwei Shi, Shanshan Lao, Yuan Gong, Mingdeng Cao, Jiahao Wang, and Yujiu Yang. Maniqa: Multi nulti-dimension attention network for no-reference image quality assessment. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition. pp. 1191-1200, 2022.
Tao Yang, Peiran Ren, Xuansong Xie, and Lei Zhang. Pixel-aware stable diffusion for realistic image super- resolution and personalized stylization. arXiv preprint arXiv:2308. 244.::39...4499, 2023.
Fanghua Yu, Jinjin Gu, Zheyuan Li, Jinfan Hu, Xiangtao Kong. Xintao Wang, Jingwen He, Yu Qiao, and Chao Dong. Scaling up excellence: Practicing model scaling for photo-realistic image restoration in the wild. arXiv preprint arXiv:2401. 13627, 2024.
Denis Zavadski, Johann-Frieen Friedrich Feiden, and Carsten Rother. Controlnet-xs: Designing an efficient and effective architecture for controlling text-to-image diffusion models. arXiv preprint arXiv:2312.06573 2023.
Lvmin Zhang and Maneesh Agrawala. Transparent image layer diffusion using latent transparency. arXiv preprint arXiv:2402. 17113, 2024.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3836-3847, 2023.
Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu, Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations. Advances in Neural Information Processing Systems, 35:3609 3623, 2022.
Shihao Zhao, Dongdong Chen, Yen-Chun Chen, Jianmin Bao, Shaozhe Hao, Lu Yuan, and Kwan-Yer K Wong. Uni-controlnet: All-in-one control to text-to-image diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Published as a conference paper at ICLR 2025
A BROADER IMPACT
Controlled generation technology, as a pivotal innovation in the field of diffusion models, exerts significant impact actoss multiple sectors of society. In the creative industries, it enables artists and designers to realize complex visions with unprecedented precision and flexibility, fostering innovation in digital art, design, and multimedia content creation. In commercial applications. controlled generation technology enhances marketing strategies by offering more targeted and dynamic advertising visuals, effectively engaging consumers. Additionally, its influence extends to education and training, where it can revolutionize teaching methods and materials, especially in visually-dependent disciplines, by generating customized educational content and simulations.
Although system enables artists, designers, and content creators to realize their creative visions through precise control, is crucial to recognize the potential negative societal impacts that could arise from misuse abuse, similar to other AI models for image generation and editing. To address these issues, responsible deployment practices, ethical standards, and including special markers in generated images to increase transparency essential steps in achieving responsible use.
B LIMITATIONS
UniCon's 's training process eliminates the need for the gradient of the base generative model, signifi- cantly reducing computational overhead. However, UniCon does not decrease the network's 's parameter count, resulting limited speed improvement during sampling compared to DiT-ControlNet..
C
MORE DETAILS
Comparison Models. In our experiment with the stable diffusion v2.1 (SD), we employ the official versions of the ControlNet (Zhang et al.) 2023) and T21Adapter (Mou et al. 2024) models without any modifications. Since ControlNet is originally designed for U-Net, there no official version compatible with the DiT transformer- base diffusion model. Therefore, the DoT-ControlNet variant used in our study a reproduction by Art-rt (Chen et al. 2023). Following the design principles of ControlNet, the initial halves of the DiT blocks are copied and trained as the controller: The primary difference from the U-Net-based ControlNet is the lack of skip connections in DiT-ControlNet. This makes DiT-ControlNet unable detach the gradient of the "Encoder" part of the base model, significantly increasing peak VRAM usage.
Sampling scheme. We employ second-order DPM solver sampler to generate images, with sampling step of 24. In our Pix (Chen et al. (023) experiment, we incorporated null tokens as negative prompts and applied CFG (Ho & Salimans] 2022) scale of 4.5. Conversely, in the SD v2.1 experiment, the negative prompt included terms such as "blurring, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, low-res, over-smooth' and we used CFG scale of 7.5.
ControlNet
0.16
UniCon-half
Loss
0.14
0.12
0.1
30
35
Hours
Figure 9: Comparison of training curves between ControlNet and UniCon-hall for 100K iterations under Canny edge condition. Notice that ControlNet and UniCon have similar values at the start of training, but show differences after smoothing due UniCon's faster convergence.
15

---

Published as a conference paper at ICLR 2025
Task
Training Steps
Contrellability
Generation Quality MAN-IQAT
Text Consistency Clip-Scoret
Metric
Value
FID
Clip-IQAT
MUSIQU
50K
SSIM1
0.2974
70.40
0.4326
0.1167
50.96
0.7623
100K
SSIM
0.5458
46.71
0.6577
0.2029
66.45
0.7889
Canny
400K
SSIMT
0,5507
46.13
0,6633
0,2082
66.50
0,7933
22.86
0.6996
0.2401
0.7909
LQ
[00K
PSNRT
37.34
20.34
0.7251
0.2831
69.99
0.8022
200K
PSNR1
37.40
20.07
0.7172
0.2846
0,8014
PSNRT
19.94
0.7291
0.2819
0,8009
Table 3: Comparisons of different training steps for DiT-UniCon-Ful indicates the larger the better and J. indicates the lower the better. Bold represents the best performance.
DIT
U-Net
Forward
Forward Backward
3.5
Backward
3.0
TFloats
3
TFloats
2.0
2
1.5
1
1.0
ControlNet-Full
UniCon-Full
ControiNet- Decoder
Controltet-Decader
MniCon-Decoder
0.0
T21Adapter
ControiNet-Hall CottrolNet-Fulll
UniCon-Fuil
UniCon-Half
Figure 10: Comparison of (a) DiT and (b) U-Net forward and backward propagating TFloats with different adapters.
Metric calculations. We computed PSNR, SSIM, LPIPS, CLIP-IQA, MANIQA, and MUSIQ using the pyiqa library? 2 with PSNR specifically calculated on the Y channel. FID scores were determined via the clean-FID Additionally, MSE was caleulated in the RGB domain, with values ranging from O to 256. We adopt the mean average precision (mAP) based on object keypoint similarity (OKS) from (Cao al. 2019
Zhao et al.) 2024) as the metric for the pose condition.
Convergence tendency. trends with ControlNet. As shown in Figure 9 UniCon retains ControlNet's fast convergence. Despite reversing the direction of information flow, UniCon ensures effective initialization by using acro-initialized Connector Block and copying parameters from the pre-trained base model.
To highlight UniCon's training advantages, compared its convergence
Training Steps. We evaluated the performance of UniCon across varying training steps to identify the optimal step size. As presented in Table B UniCon fails converge at 50K steps for both LQ and Canny tasks. Performance improves substantially from 50K to 100K steps, while further increases beyond 100K yield only marginal gains. Hence, 100K steps is determined be the most cost-effective training step size.
Computational evaluation. As shown in Figure 10 we further analyzed the computational cost of each adaptor during forward and backward propagation. With the same number of learnable parameters, UniCon's 's total computational cost is lower than ControlNet's. In the UNet-Half setting, UniCon's FP TFloats slightly higher because it copies the Decoder, which has more parameters
hhttp://github.co http:///ghtbbccom//cccoon/o///// http:://gtthua.com/caaaaaaraarre
16

---

Published as a conference paper at ICLR 2025
Task
Adapter
Controllallity
Generation Quality MAN-IQAT
Text Consistency Clip-Scoret
Metrie
Value
FID
Clip-IQAT
MUSIQT
ControlNet
SSIMI
0.4895
49.80
0.6683
0.2215
67.19
0.8168
UniConntollee
SSIM
0.4996
49.95
0.6749
0,2348
67:95
0.8152
GLIGEN
Canny
121-Adapter
SSIM
0.3936
52.05
0.6783
0.2266
68.16
0.8155
ControlNeXi
SSIM
0.4568
53.03
0.6616
0.2153
66.68
0.8141
UniCon
SSIM
0.5570
47.11
0.6704
0,2335
67.99
0.8189
ControlNet
MSEL
85,70
54.30
0.6828
0.2262
67.90
0,8202
UniControlNet
MSEL
54,49
0.6846
0.2302
0.8217
Depth
GLIGEN
MSF
87.65
58.25
0.6851
0.2234
68.41
0.7964
121-Adapter
ControlNeXi
MSF
86.81
55.71
0.6799
0.2291
68.00
0.8198
UniCoon
MSE
85.00
53.45
0.6807
0.2262
67.85
0.8214
Table 4: Comparisons of different adapters for SD-UNet. 1 indicates the larger the better and indicates the lower the better. Bold represents the best performance.
Adaptor
Base Model
PSNR
LPIPS!
MUSIQ
CLIP-Scoret
FIDJ
StableSR
SD2.1
35.39
0.0839
69.57
0.7786
37.85
DiftBIR
SD2.1
34.81
0.0940
69.76
0,7844
32.85
PASD
SD2.1
36.17
0.0552
69.62
0,7925
25.61
UniCon
SD2.1
35.69
0.0530
70.48
0,8027
22.80
UniCon
PixArt-a
37.34
0.0453
69.99
0,8022
20.34
Table 5: Results different diffusion-basse image restoration models. t indicates the larger the better and J indicates the lower the better. Bold represents the best performance.
than ControlNet's s copied Encoder. While T21-Adaptor has the lowest FP TFloats, its control capability significantly weaker than UniCon's, ss and its performance degrades noticeably in low-level tasks, as shown in Table
D PSEUDO CODES
Training Cost Experiments. In Algorithm 1 we present a pseudo code to evaluate the single-round training The sampling methods for training noise and timestep remain consistent with those used during training. We subsequently report the peak memory usage and average time based on one hundred rounds of training cost evaluation.
Core Difference between ControlNet and UniCon. Algorithm and Algorithm illustrate the pseudo codes for the forward passes of ControlNet and UniCon, respectively. Unlike ControlNet, UniCon avoids storing the base model's gradients during forward propagation and truncates all gradient connections from the base model to the controller. Consequently, the base model gradients are not utilized during backpropagation. These modifications significantly reduce UniCon's memory overhead and training time compared to ControlNet.
E
MORE COMPARISON RESULTS
More Adapters on Highlevel Tasks. To comprehensively compare UniCon's performance, we evaluated it against other high-level control methods the UNet model, including UniControl- Net (Zhao et al. 2024), GLIGEN (Li al. 2023), and ControNeXt (Peng et al. 2024).. As shown in UniCon achieved the best controllability and FID scores both canny and depth tasks. NR metrics confirm that UniCon maintains strong control capabilities without notably compromising
Table 4
image quality.
Considering that several diffusion based methods have been applied to low-level tasks, we compared UniCon with related approaches in image restoration (Wang
More Adapters Lowlevel Tasks.
al. 20235: Lin et al. 2023 Yang al. 2023 As shown in Table with SD2.1 as the base
17

---

Published as a conference paper at ICLR 2025
Condition Consistency
Image Quality MAN-IQA-
Text Consistency Clip-Scoret
Diversity
Adapter
PSNR1
SSIM
L.PPSS]
Clip-IQAt
MUSIQU
FIDI
ControlNet
34,82
0.9352
0.0650
0.7147
0.2459
69.37
0.7996
26.43
XontrolNeXt
34,31
0.9329
0.0679
0,7171
0.2534
69.36
0.7997
26.28
UniCon-Half
35.64
0,9512
0.0475
0.7042
0.2675
69.51
0.8025
22.07
UniCon-Haldt
36.41
0,9512
0.0462
0.7189
0,2782
69.66
0.8008
21.34
Table 6: Comparisons of different adapters for PixArt-a x4 Super- Resour-Rnoolutin task. f represents
the use of the pre-trained ControlNeXt model as the base model. f indicates the larger the better and I indicates the lower the better. Bold represents the best performance.
model, UniCon achieves the highest image quality, with fidelity significantly surpassing StableSR and DifTBIR, and comparable to PASD. Furthermore, using PixArt-a the base model, UniCon's fidelity improves significantly.
Discussion of ControlNeXt. Since ControlNeXt introduces only a small number of additional parameters during training (Peng et al. 2024), it offers a significant advantage in inference speed compared to other ControlNet variants. However, as shown in Table 6 XontrolNe control capabil- ities are slightly below ControlNet and significantly behind UniCon While UniCon requires longer inference time, it delivers better performance in scenarios with stricter consistency requirements. Additionally, we found that UniCon and ControlNeXt offer complementary benefits. Using Control- NeXt as UniCon's 's pre-trained model significantly boosts UniCon's 's performance, achieving much higher fidelity and quality than ControlNet with the same number of parameters. This two-stage training process can be viewed as a course-to-fing fo-ee approach with high efficiency.
Visual Comparison. Due to space limitations, we provide more visual comparison results in Figure [5] Figure an Figure Massive visual comparisons prove our method not only maintains high-quality generation but also strictly adheres to the controls of the input images.
F FULL ABLATION STUDY RESULTS
Due to space limitations in the main text, we only present some results in the main text. We present Figure
the complete ablation experimental results in Table 7 Figure 11 Figure
12
13 and Figure 14
The Effect of Copying "Encoders" and "Decoders" on Diffusion Control. Visual comparisons of different adapter structures in Figure 1T reveal that the "Encoders" model exhibits insufficient controllability. In the canny task, this inadequacy manifests as the addition of objects misaligned with the canny lines, such as the giraffe's 's neck in Figure Ia).. In the SR task, it results in color shifts and artifacts, such as the overall darkness and the sky with messy spots in Figure 11 b). Conversely, while copying "Decoders" improves controllability, it compromises the quality of the generated images, as evidenced by the coloring errors on the face in Figure and the messy textures on vegetables and human faces in Figure (b). This aligns with the conclusion that focusing on encoders leads to better generation effects, whereas focusing on decoders enhances controllability. As shown in Table 7 the "Skip-Layer" type achieves the highest controllability and image quality among the three copying schemes with the same parameter count. The "Full" copying structure achieves higher metrics and better visual effects by adding parameters on top of the "Skip-Laaerr" structure, as demonstrated in Figure T a), where the "Full" type maintains facial features intact from multiple angles.
Unidirectional Information Flow vs. Changing Features Directly. As indicated in Table 7 using adapters as outputs under the "Skip-Layer" structure negatively impacts the Canny task and causes collapses in the SR task. This likely results from error accumulation in the adapters during initialization with the "Skip-Layer" copying scheme. These errors are less influential when the base diffusion model is used as the output but affect training stability when adapters directly output. Beyond the "Skip-Layer" type, unidirectional information flow steadily improves the model's controllability. In the canny task, unidirectional information flow produces results more consistent with the conditioned lines, as seen in Figure 12 a): the "Full" case curtains and the "Decoder" case 1 water surface. In the SR task, it enhances fidelity, exemplified by the Chinese string in Figure [2]b) In conclusion, Figure 12 shows that unidirectional information flow significantly reduces training overhead while maintaining overall image quality without degradation.
18

---

Published as a conference paper at ICLR 2025
The Choice of Connector. As demonstrated in Table using ZeroFT as the connector in both the Canny and SR tasks maintains the highest controllability and image quality metrics. Employing ZeroFT enables the network to effectively manage relationships between objects in complex scenes. With ZeroFT, the network can generate a more natural half-length photo of the woman in Figure 13(a) a) and accurately process the margin between the hollow chandelier and wall in Figure 13
The Influence of Attes-Attention Blocks in Adapters. Previous work, such as BrushNet Ju al., 2024). has discussed that removing coss-attention blocks connected the text prompt in the Adapter module can enhance the fidelity of ControlNet. We investigated whether this approach would also be effective for UniCon. Therefore, we tested the impact of removing cross-attention blocks in the canny and SR tasks in complete ablation study. As shown in Table 7 removing cross-attention blocks decreases all metrics the canny task but improves them in the SR task. in
Visually, dropping cross- attention causes incorrect color matching in the canny task, Figure 14 a). However, in the SR task, it prevents the generation of cluttered high-frequency textures, illustrated in Figure b)).. We speculate that this because the SR condition already provides sufficient information, making the text prompt less critical, while excessive text guidance can cause CFG to produce messy high-frequency textures. In the canny task, the network relies more the text prompt to supplement the information that the condition alone cannot provide. Consequently, we removed the pross-attention module adapters for the SR and Blur+SR tasks.
19

---

Published as a conference paper at ICLR 2025
Watercolor painting of deer.
Blonde Pixie, Disconnected.
Canny
Encoder
Decoder
Skip-Layer
Full
(b)
Blond girl sitting near red car.
A dining table.
Downsampled
Encoder
Decoder
Skip-Layer
Full
Ground Truth
Figure 11: Comparison of different variants of UniCon on conditional generation with (a) Canny and (b) downsampled images.
(e)
Skip-Layer
Decoder
Full
Tea plondation Comerrn highlands.
Moose and Denoii after sunnise from mirror pond.
While Kanihie fankher, adler whhrrrors blssererver
Design of modern living room
Soothing bedtooms that tokle inspraatoon from the clouds:
Ebe Sansitone Morning, Muao Rock Mesales, Landicone
Canny
W/
W/
(b)
Decoder
Full
Row green pingham lnen stinny tie.
girl with curly hair wearing headeiece.
Shanpha Dicarylaned
ehurch
RRBS
RESOR
2013
2018
2018
unan
2013 UPRB
x4 Downsampled
W/O
W/
Ground Truth
x4 Downsampler led
W/O
W/
Ground Truth
Figure 12: Comparison of different variants of UniCon with or without unidirectional flow conditional generation with (a) Canny and (b) downsampled images.
20

---

Published as a conference paper at ICLR 2025
(a)
Rain squall near sunset, Grand Canyon NP.
A woman in pale pink dress.
Canny
ZeroConv
ShareAttn
ZeroFT
(b)
old man in black suite holding camera portrait Italy black and white.
modern living room accessories furniture.
x4 Downsampled
ZeroConv
ShareAttn
ZeroFT
Ground Truth
Figure 13: Comparison of UniCon using different connectors on conditional generation with (a) Canny and (b) downsampled images.
(a) Tamako Love Story, Kyoto Animation.
(b)
Concrete Wall Designs Striking Bedrooms.
Ponoramic View of Kawaguchiko.
Outdoor Patio with Wicker Rattan Set.
Canny
W/O
W/
Downsampled
W/O
w/
Ground Truth
Figure 14: Comparison of UniCon adopting cross attention blocks in Adapter with (a) Canny and (b) downsampled images.
21

---

Published as a conference paper at ICLR 2025
x4 Downsampled
DiT-ControlNet
DiT-UniCon-Half
DiT-UniCon
Ground Truth
x4 Downsampled
DiT ControlNet
DiT-UniCon-Half
DIT-UniCon
Ground Truth
Downsampled
DiT-ControlNet
DiT-UniCon-Half
DiT-UniCon
Ground Truth
Downsample
CTT-ControlNet
DiT-UniCon-Half
DIT-UniCon
Ground Truth
Blur Downsample
DIT-ControlNet
DIT-UniCon-Half
DiT-UniCon
Ground Truth
Blur & Downsample
DiT-ContralNet
DiT-UniCon-Half
DiT-UniCon
Ground Truth
Blur & Downsample
DIT-ControlNet
DiT-UniCon-Half
DiT-UniCon
Ground Truth
Figure 15: Qualitative comparison of DiT with ControlNet (Zhang et al. 2023) and UniCon dealing with downsampled image condition.
22

---

Published as a conference paper at ICLR 2025
White Wolf With Blue Eyes.
man suit was standing astronomy exhibition
Canny
DIT- ControlNet
DiT-UniCon
Canny
DiT ControINet
DiT- DiT-UniCon
with mountalns
bird flight, captured in mid-air.
Canny
DiT-ControlNet
DIT-UniCon
Canny
DiT-ControlNet
DIT-UniCon
Two birds standing on branches.
A bear is in the gross.
Depth
DIT Controleet
DIT-UniCon
Depth
DiT-ControlNet
DIT-UniCon
set with variety of tableware oranges.
bridge on the lake, autumn.
Depth
DiT-ControlNet
DIT-UniCon
Depth
DIT-ControlNet
DIT-UniCon
man and womon together, both them wearing sunglosses.
A person performing push-up exercise.
DiT-ControlNet Corvo Attano, the protagonist of "Dishonored."
Pose
DIT-UniCon
Pose
DiT-ControlNet
DiT-UniCon
A woman wearing gorment with a shiny, metalicc texture.
Pose
DIT-ControlNet
DIT-UniCon
Pose
DIT-ControlNet
DIT-UniCon
Figure 16:
Qualitative comparison DiT with ControlNet (Zhang et al. 2023) and UniCon dealing with canny, depth, and pose image conditions.
23

---

Published as a conference paper at ICLR 2025
collection of macarons.
Canny
SD-ControlNet Woman and dog on the grass.
SD-T21
SD-UniCon
Canny
SD-T2I
SD-ControlNet
SD-UniCon
lotus flower full bloom.
Depth
SD-ControlNet A modern living room with contemporary design.
SD-T2I
SD-UniCon
Depth
SD-T2I
DD-ControlNet
SD-UniCon
x4 Downsampled
SD-T2I
SD-ControINet
SD-UniCon
Ground Truth
x4 Downsampled
SD-T2I
SD-ControiNet
SD-UniCon
Ground Truth
Figure 17: Qualitative comparison of SD with T2I-Adapter Moou et al. 2024), ControlNet (Zhang
al. 2023 and UniCon dealing with canny, depth, and downsampled image conditions.
24

---

Published as a conference paper at ICLR 2025
Table 7: Results of full ablation study. 1 indicates the larger the better and J. indicates the lower the better. Bold and underline represent best and second best performance.
Text Consistency
CLIP CLIP-ScoO Score"
0.7724
0.7776
0.7561
0,7507
0,7612
0.7520
0.7606
0.7696
0.7818
0.7889
0.7996
0.8009
0.8013
0.8018
0.8013
0.8012
0.8025
0.8026
0,8022
MUSIQ1
65.24
66.05
63.66
62.51
64.27
63.75
66.32
66.45
69.37
70.00
68.26
68.61
69.23
69.24
69.51
69.91
69. 99.99
62.
64.
Generation Quality MAN-IQAT
0.1861
0.1978
0.1737
0.1621
0.1780
0.1638
0.1723
0.1906
0.1867
0.2029
0.2459
0.2541
0.2325
0.2358
0.2499
0.2481
0.2675
0.2609
0,2831
CLIP-IQAT
Failed
0.6439
0.6629
0.6224
0.6047
0.6347
0.6177
0.6154
0.6563
0.6397
0.6577
0.7147
0.7229
0.6979
0.7036
0.7012
0.7078
0.7042
0.7212
0.7251
FID4
51.52
49.78
56.93
59.32
55.22
56.16
56.22
52.31
50.17
46,71
26.43
24.99
25.84
23.55
22.99
23.03
22.07
23.04
20.34
LPIPS
0.3280
0.3033
0.3016
0.2953
0.2714
0.2798
0.2778
0.2658
0.3070
0.2538
0.0650
0.0616
0.0589
0.0519
0.0502
0.0510
0.0475
0.0522
0.0453
Controllabiitt SSIM
0.4748
0.4983
0.5078
0.5131
0.5343
0.5260
0.5236
0.5426
0.5053
0.5458
0.9352
0.9385
0.9417
0.9478
0.9495
0.9486
0.9512
0.9475
0.9542
PSNR
34.82
35.49
34.85
35.59
35.67
35.55
35.64
36.53
37.34
Connector Type
ZeroMLP
ZeroMLP
ZeroMLP
ZeroMLP
ZetOMPP
ZeroMLP
ShareAttn
ZeroFT
ZeroFT
ZeroFT
ZeroLP
ZeroMLP
ZeroMLP
ZeroMLP
ZeroMLP
ZeroMLP
ShareAttu
ZeroFT
ZeroPT
ZeroFT
Controller CrossAtt
7
X
Unidiroctional
X
X
Adapter
Encoder
Layer
Decoder
Encoder
Skip- Layer
Decoder
Full
Full
Task
Canny
DS
25

---

Published as a conference paper at ICLR 2025
Algorithm 1 Pseudo Code for One Round Training Cost Evaluation
Require: sampled noise (eta), noised input (xt), timestep (t), prompt (y), condition (c)
1: device = torch.device())
2: model mnooel..o(device
3: WEIGHT BEEETTYEERRYY = torch.cuda.ma. we cemmor_alllccate((deeic))
4: with torch.no_grraae
pred eta = model(xt, c)
MSE Losspped eta, eta)
5: ACTIVATION SEEAATIIOOMMOOYY---
6: model(t y,
l MSs=MMME_Lsssppeee
eta)
bass.backward()
7: GRADIENT MEMORY torch.cuda. memory wwwommmmmlloaaaoooeaee allocated((eviic))
8: fp_start_tingg = time time()
9: pred_eta _tt_etta mxdel(xt, y. c)
loss MSF=_SLSssss LNSF_Loossppeed_ett eta)
TPMME time.time() fp start_ time
11: bp_start time = time. time() time.time()
optimizer.zero_greee loss. bass.backward))
optimizer.ste(()
13: BP 8P_TIME timettime) bp_start_tiine
14: OPTIMIZER _EMORRYYEERRY = wow.hccuwaw.eeneoeoooooeeeeeeeee
return * MEMORY, _TIME
Algorithm 2 Pseudo Code for DiT-ControlNet Forward Pass
Require: noised input (xt), timestep (t), prompt embedding (y), condition (c)
I: base modelimagee embedding(xt)
2: t_ t_emb base timestpp embedding(t)
3: controller.image (aargge__meedii(cc
4: t_emb_ cemb_cond- panatroller.imeseeeeeeeeeeeeeeee
5: = x
6: for base_block, control_block in zip(base_moodel.llos controller.blocks) do
7: x base mnse__locc(x, (_eb))
8: с control_ block(c, emb con__
9: x x control block.comnector(c)
10: end for
pred base_ sasee_mdel.proy proj__ou(x))
12: return pred
Algorithm 3 Pseudo Code for DiT-UniCon Forward Pass
Require:
noised input (xt), timestep (t), prompt embedding (y), condition (c)
1: with torch.no. torch.no_gradd. _groo()ra
X base model.image powddellaaaeeebmeddingiiiee embeeddingg(x))ddetacch)))
f embb base model.timestep powdee/ineeseppembbdiiiieg embedding(()
2: = controller.image hatprolleerinne_eemedennnnggg
3: [_emb_cond =
4: C x
5: for base
block. base_block. control block in hpp(base_model.bleee controller.blocs)) do
6: with torch.no ::
base block(x, t detacc))
7: c= blockol cemb_condl
8:
control_ caairoo__blccc.cconettrri
9: end for
pred = controller.rojj out(c)
11: return pred
26


Reviews:
[{"review_id": "piK7ZbfGwD", "reviewer": "Reviewer_Mp3R", "summary": "The paper presents UniCon, a novel architecture aimed at enhancing control and efficiency in training adapters for large-scale diffusion models, such as the Diffusion Transformer (DiT). Unlike existing methods that involve bidirectional communication between the control adapter and diffusion model, UniCon employs a unidirectional information flow, where information flows only from the diffusion model to the adapter. This enables the adapter to generate the final output independently, reducing computational demands and memory requirements by eliminating the need for the diffusion model to compute and store gradients during adapter training.", "questions": "The sentence in line 155 \"Given a ...\" is broken.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "UniCon’s key strength is demonstrating that gradients through the base diffusion model are unnecessary for effective control, achieving this with a unidirectional flow design. This architecture enables the adapter to independently generate outputs, significantly reducing memory usage and computational costs by omitting gradient calculations for the base model. The work is substantiated with an ablation study for different architectural alternatives.", "weaknesses": "I would have liked to see an insightful reasoning, why the unidirectional flow apparently increases the quality of the output. Would the results converge after longer learning times? Is the effect comparable to a \"low\" rank regularization (somewhat far fetched)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "0o5vIM4QlR", "reviewer": "Reviewer_5dEj", "summary": "The paper proposed a new method called UniCon, using unidirectional flow for adapting diffusion models. Validation had been made on both the SD U-Net diffusion model and the DiT diffusion model, on several datasets. It turns out that UniCon reached a good result of memory reduction and training speeding up while maintaining the same adapter parameter size and generative capabilities.", "questions": "1. Can you provide a schematic or computational rule of forward computation and backpropagation during training to compare ControlNet and UniCon? Just looking at Figure1 and Figure2, readers may not understand why UniCon does not require computing and storing gradients for the diffusion model while ControlNet requires. As in Line 187: \"For example, when training ControlNet for DiT, when the gradient of ControlNet itself occupies about 18GB VRAM, the gradient brought by the DiT diffusion model needs to occupy 16GB VRAM.\" \n2. As for the choice of the connector (Line 459) in Table 1b, it seems that each ZeroFT contains $\\textbf{2 layers}$ of ZeroMLP (joined by multiplication and addition). Is the ZeroMLP compared here consisting of $\\textbf{single-layer}$ structure as the one in the original article? Is it unfair? More rigorous experiment results may be needed to illustrate ZeroFT's effectiveness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The motivation makes sense: for a faster and more efficient training of controlled generation of diffusion models, towards further scaling up (good for training of scaling up).\n2. The evaluation is clear and the experiment results are abundant.\n3. Results of unidirectional flow are promising, making a good contribution to the controlled generation of diffusion models and further work.\n4. The appendix section of the paper is thorough and contains quality information.", "weaknesses": "1. A more detailed introduction could be made of diffusion models for readers' better following. Along with a detailed explanation of ControlNet.\n2. The speeding up and memory reduction are mainly for training. But during inference, will UniCon be similar even worse in latency and memory due to the more complex structure (e.g., more Control Blocks compared to the vanilla one on UNet structure, and each block seems to contain more layers of MLP)? This also seems to explain why UniCon performs better because the block structure is more complex.\n3. It seems that the page number of Page1, the header and the edge of Page2 can be incorrectly hyperlinked.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "qNrzt7dAgJ", "reviewer": "Reviewer_r7et", "summary": "This paper introduces UniCon, a novel approach for effective and efficient control in large-scale diffusion models using unidirectional information flow. The presentation is overall quite good, with clear and informative figures and tables, although there is some room for improvement in the writing accuracy. The experiments appear comprehensive and solid. I look forward to seeing feedback from other reviewers to ensure that I haven’t missed any important aspects.", "questions": "* Will the code and weights be released to facilitate broader adoption and further development by the research community?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper is well-structured and easy to follow. Visual representations in figures and tables effectively support the text.\n\n2. The \"Unidirectional information flow\" concept is novel, ensuring conditional inputs do not interfere with the inherent data structure, a significant advantage.\n\n3. The experiments conducted are comprehensive, validating the proposed method’s efficiency and performance across different tasks and models.", "weaknesses": "1. A minor typographical error was noted: \"PixeArt\" should be corrected to \"PixArt\" (Line 122).\n2. The statement that current designs may not be suitable for transformer-based diffusion models (Line 198-205) could be more thoroughly justified, as it might not hold universally.\n3. There is an error in Table 1(a), where an incorrect metric has been highlighted and needs correction.\n4. I'm afraid I have to disagree with the conclusion drawn in lines 372-376, as the phenomenon might be attributed to the decoder being a more significant part of the model. In fine-tuned models, there is often a trade-off between image quality and controllability—enhanced controllability may lead to a reduction in image quality. \n5. Additionally, \"generation effect\" does not seem to be the most appropriate term for this context and could be replaced with a more accurate expression.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "imTwfvgFKH", "reviewer": "Reviewer_1fb4", "summary": "This paper presents UniCon, tailored for both UNet-based and DiT-based large-scale diffusion models. The proposed unidirectional flow method is the core method, allowing the adapter alone to generate the final output to reduce VRAM memory consumption. Experiments demonstrate that UniCon achieves good generative performance with lower resource consumption.", "questions": "Please see the weaknesses. \n\nIf I misunderstand any aspects, I welcome the authors to clarify and discuss.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The writing of this paper is good and easy to read.\n\n2. UniCon offers a new approach to integrating adaptors in DiT-based diffusion models.\n\n3. The visualization and quantitative results from the experiments validate the effectiveness of the method.", "weaknesses": "1. UniCon designs two different methods for incorporating adaptors in UNet-based and DiT-based diffusion models, as shown in Figure 2. The differences between these methods are significant; for instance, UniCon for UNet adds many concatenation operations and skip connections, while the DiT-based model employs a different serial structure.\n   - a) What specific motivations led the authors to design these two distinct structures?\n   - b) If the method needs to adapt to other frameworks, such as Mamba-based diffusion or XLSTM-diffusion, would the adaptor structure need to be redesigned? This raises concerns about the generalizability of UniCon.\n   \n2. One of the core contributions of this paper is the introduction of the ‘unidirectional’ concept. Unlike ControlNet, which still uses the original diffusion model as output and therefore needs to store diffusion gradients, UniCon modifies this by freezing the original diffusion and using the adaptor as output, thereby reducing VRAM consumption. This makes UniCon fundamentally similar to the LoRA adaptor.\n   - a) How does the performance of UniCon compare to LoRA methods?\n   - b) Has the UniCon-encoder version been tested? This could be considered a unidirectional ControlNet, and I believe it is an essential comparison to make.\n\n3. I appreciate the extensive ablation studies conducted by the authors on different structures (in Figure 3) to validate the framework's effectiveness; however, the results are not surprising. The full version which has the most parameters, performs the best, which is expected. More importantly, what motivation or significant proof do these ablation studies provide for the design of this method (please do not conclude that the full version is better than other non-full structures)? If none, these results would render the paper more of an empirical study.\n\n4. For both the main experiments and the VRAM comparison experiments, please include comparisons with ControlNet-Unidirectional-Encoder, as this is a core experiment to validate the effectiveness of the unidirectional method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)