Simple Drop-in LoRA Conditioning on Attention Layers Will Improve Your Diffusion Model
Authors: Joo Young Choi
Jaesung R. Park
Inkyu Park
Jaewoong Cho
Albert No
Ernest K. Ryu
Venue: Accepted by TMLR ()

Abstract: Current state-of-the-art diffusion models employ U-Net architectures containing convolutional and (qkv) self-attention layers. The U-Net processes images while being conditioned on the time embedding input for each sampling step and the class or caption embedding input corresponding to the desired conditional generation. Such conditioning involves scale-and-shift operations to the convolutional layers but does not directly affect the attention layers. While these standard architectural choices are certainly effective, not conditioning the attention layers feels arbitrary and potentially suboptimal. In this work, we show that simply adding LoRA conditioning to the attention layers without changing or tuning the other parts of the U-Net architecture improves the image generation quality. For example, a drop-in addition of LoRA conditioning to EDM diffusion model yields FID scores of 1.91/1.75 for unconditional and class-conditional CIFAR-10 generation, improving upon the baseline of 1.97/1.79.

Published in Transactions on Machine Learning Research (09/2024)
Simple Drop-in LoRA Conditioning on Attention Layers Will Improve Your Diffusion Model
Joo Young Choi
thiinblover@smo..c..
Deportment of Mathemationl Sciences Scoul National Universily
Jacsung R. Park
ryanpark/7@smm..cc
Department Mathematical Sciences Seoul National University
Inkyu Park
inkyuparkk@kratton.com
KRAPTON
Jacwoong Cho
juchoo@kraftnn.com
KRAPTON
Albert No
albertno@yonseiic oc.
Department of Artificial Intelligence Yonsei University
Ernest K. Ryu Deportment of Mothematics
eryu@math. ucla edu
University of California, Los Angeles
Reviewed on OpenReview:
Res Block
Attn Block
Res Block
Attn Block
Scale-Shift
LoRA
Scale-Shift
LoRA
Res Block
Attn Block
Res Block
Attn Block
Scale-Shift
LoRA
Scale-Shift
LoRA
Block
Res Block
Scale-Shift
LoRA
Scale-Shift
The standard U-Net architecture for diffusion models conditions convolutional layers in residual blocks with scale-and-shift but does not condition attention blocks. Simply adding LoRA conditioning on attention layers improves the image generation quality.
Figure 1:
Abstract
Current state-of-the-art diffusion models employ U-Net architectures containing convolutional and (qkv) alff-tttentinn layers. The U-Net processes images while being conditioned on the time embedding input for each sampling step and the class or caption embedding input corresponding to the desired conditional generation. Such conditioning involves scale-and-shift operations to the convolutional layers but does not directly affect the attention layers. While these standard architectural choices are certainly effective, not conditioning the attention layers feels arbitrary and potentially suboptimal. In this work, we show that simply adding LoRA conditioning to the atlention layers without changing or tuning the other parts of the U-Net architecture improves the image generation quality. For example, a drop-in addition of LoRA conditioning EDM diffusion model yields FID scores of 1.91/1.75 for unconditional and consiconniiiinaa CIFAR-10 generation, improving upon the bascline of 1.97 1.97/1.70 1.79.. Code is available at httpss//giin whu.. /tt: nk over dii fusion.nn.. fusion_lore

---

Published in Transactions on Machine Learning Research (09/2024)
Dataset
Model
Type
NFE
# basis
rank
FID
# Params
bascline
4000
3.69
52546438
IDDPM
only LoRA
4001
11
4
3.64
47591440
with LoRA
4001
11
4
3.37
54880528
CIFAR-10 (uncond.)
baseline
35
1.97
55733891
EDM (vp)
only LoRA
35
18
1.99
53411675
with LoRA
35
18
1.96/1.91
1.96/
57745499
baseline
4000
3.38
52551558
IDDMM
only LoRA
4001
(11, 10)
2.91
48513040
with LoRA
4001
(11, 10)
3.12
55807248
CIFAR-10 (cond.)
baseline
35
1.79
55735299
EDM (vp)
only LoRA
35
18
4
1.82
53413083
with LoRA
35
18
4
1.75
57746907
baseline
4000
19.2
121063942
ImageNet64 (uncond.)
IDDM
only LoRA
4001
11
18.2
113602960
with LoRA
4001
18.1
139278224
baseline
79
2.39
61805571
FFHQ64 (uncond.)
EDM (vp)
only LoRA
79
20
2.46
58935795
with LoRA
79
20
2.37
2.37/2.28
63941631
Table 1: Image generation results. We compare three different conditionings for each setting: (baseline) conditioning only convolutional layers residual blocks with scalc-and-shift; (only LoRA) conditioning only atlention layers using LoRA conditioning and not conditioning convolutional layers; (with LoRA) conditioning both convolutional layers and attention layers with scale-and-shiff and LoRA conditioning, respectively. For unconditional CIFAR-10 and FFHQ64 sampling using EDM with LoRA, we also report the FID score obtained by initializing the base model with pre-trained weights.
Introduction
In recent years, diffusion models have led phenomenal advancements in image generation. Many cutting- edge diffusion models leverage U-Net architectures their backbone, consisting of convolutional and (gkv)
self-attention layers (Dhariwal & Nichol 2021 Kim et al. 2023 Saharia al. 2022 Rombach ct al. 2022 Podell et al.] 2024). In these models, the U-Net architeclure. based score network is conditioned on the time, onr/or, class, text embedding (Ho & Salimans 2021 using scale-and-shif operations applied to the convolutional layers in the so-called residual blocks. Notably, however, the attention layers are not directly alfected by the conditioning, and the rationale behind not extending conditioning to attention layers remains unclear. This gap suggests a need for in-depth studies searching for effective conditioning methods for attention layers and assessing their impact on performance.
Meanwhile, low-rat lowk adaptation (LoRA) has become the standard approach for crrameter-effcieen fine-tuning of large language models (L.MM) (Hu et al. 2022p. With LoRA, one trains low-rank updates that are added to frozen pre-trained dense weights in the attention layers of LLMs. The consistent effectiveness of LoRA for LLMs suggests that LoRA may be generally compatible with attention layers used in different architectures
and for different tasks (Chen et al. 2022 Pan et al. 2022 Lin et al. 2023 Gong et al. 2024).
In this work, we introduce novel method for effectively conditioning the attention layers in the U-Net architectures of diffusion models by jointly training multiple LoRA adapters along with the base model. We call these LoRAA adapters TimeLoRA and ClassLoRA for discrete-time settings, and Unified Compositional LoRA (UC-LoRA) for continuous sigual-to-ratio (SNR) settings. Simply adding these LoRA adapters in a drop-in fashion without modifying or luning the confiuurtiinss of original model brings consistent enhancement in FID scores across several popular models applied to CIFAR-10 10, FFHQ 64x64, and ImageNet datasets. particular, adding LooA-connitiionig to the EDM model (Karras et al. 2022) yiclds improved FID scores of 1.75, 1.91, 2.28 for class-conditional CIFAR-10. unconditional CIFAR-10, and FFHQ 64x64 datasets, respectively, outperforming the baseline scores of 1.79, 1.97, 2.39. Moreover, find that LoRA conditioning
2

---

Published in Transactions on Machine Learning Research (09/2024)
Input
Conditioning
Input
Conditioning
Group Norm
Linear
Group Norm
Am
SiLU
Bm
B
LoRA
MLP
Convolution
QKV
scacee
# (t)
mm (t)
Scale-Shift
LoRA
TimeLoRA and ClassLoRA
Group Norm
Dot Product
SiLU
Convolution
LoRA
MLP
M.P
Projection
scale
Byn
LoRA
Altn. Block
oig
LoRA
U-Net block
LoRA conditioning attention block
Unified compositional LoRA
Figure 2: Conditioning of U-Net Block: (left) scale-and-shift conditioning on the convolutional block (middle)
LoRA conditioning on the attention block (right) top: TimeLoRA and ClassLoRA for the discrete-ting setting, bottom:
unified composition LoRA for the contimous-SSNR setting.
by itself is powerful enough to perform effectively. Our experiments show that only conditioning the attention layers using LoRA adapters (without the conditioning convolutional layers with scale-and-shill) achieves comparable FID scores compared to the baseline scale-and-shift conditioning (without LoRA).
Contribution. Our experiments show that using LoRA to condition time and class information on attention layers effective across various models and datasets, including nano diffusion (Lelarge al. 2024), IDDPM NNichol Dhariwal, 2021 and EDM (Karras et al., 2022) architectures using the MNIST (Deng 2012).
CIFAR-10 (Krizhevsky al. 2009), and FFHQ (Karras al. 2019) datasets.
Our main contributions are follows. (i) We show that simple drop-in LoRA conditioning on the attention layers improves the image generation quality, as measured by lower FID scores, while incurring minimal (~10%) added memory and compute costs. (ii) We identify the problem of whether to and how lo condition attention layers in diffusion models and provide the positive answer that attention layers should be conditioned and LoRA an effoctive approach that outperforms the prior approaches of no conditioning or conditioning with adaLN (Peebles Xie 2023).
Our results advocate for incorporating LoRA conditioning into the larger state-of-the-art U-Net-based diffusion models and the newer experimental architectures.
2 Prior work and preliminaries
2.1
Diffusion models
Diffusion models (Sehl-Dickstein et al. 2015
Song & Ermon 2019 Ho et al. 2020 Song et al. 2021b
generate images by iteratively removing noise from a noisy image. This denoising process is defined by the reverse process of the forward diffusion process: given data To ~ go. progressively inject noise to wo by
g/ (x-1) = N (V1 -
for 1,..., T and 3, 1. If Be is sufficiently small, we can approximate the reverse process as a( | x N (Me(x)) 3,1)
3

---

Published in Transactions on Machine Learning Research (09/2024)
where
1
pc(x))
VT- (xr + 3. log
A diffuuson model is trained to approximate the score function V log Pt (+4) with score network 80, which
is often modeled with a U-Net architecture (Ronneberger et al. 2015 Song & Ermon 2019). With se
V (zee,, the diffusion model approximates the reverse process as
Po( = N (Vvi"
(xt + Beso 1)) (1- -1 Br
q(xc-11 xt).
To sample from a trained diffusion model, one starts with Gaussian noise ~ N (0, )I) where a₁ = 1I (1, and progressively denoise the image by sampling from pe (1 with = T,T 1 ,,...,.. ,, soquentially to obtain a clean image Zg-
The above discrete-time description of diffusion models has a contimuous-time counterpart based on the theory of stochastic differential equation (SDE) for the forward-corrupton process and reversing it based on Anderson's reversc-time SDE (Anderson 1982 reversc-time ordinary differential equation (ODE) with
equivalent marginal probabilities (Song et al. 2021a Higher-order integrators have been used to reduce the
discretization errors in solving the differential equations (Karras et al. 2022
Architecture for diffusion models. The initial work of Song & Ermon 2019) first utilized the CNN-based U-Net architecture (Ronncberger et al.
2015) as the architecture for the score network. Several improvements
have been made by later works (Ho et al. 2020 Nichol & Dhariwal 2021
Dhariwal & Nichol 2021 Hoogeboomi
et al. 2023 incorporating multi-head self-attention (Vaswani et al.
2017). group normalization (Wu & 2018). Recently, several alternative
He (018), and adaptive layer normalizate lization (adaLN) (Perez et al.
architectures have been proposed. Jabri et al. (2023) proposed Recurrent Interface Network (RIN), which
decouples the core computation and the dimension of the data for more scalable image generation. Peebles &
Xie (2023) ; Bao et al. (2023) Gao et al. 20223) Hatamizadeh et al. (2023) investigated the effectiveness of transformer-base based architectures (Dosovitskiy et al. 2021) for diffuusin models. Yan et al. (2023) utilized
state space models (uu et al. 2022) in DiffuSSM to present an attention-free diffusion model architecture. In this work, we propose a conditioning method for attention layers and test on several CNN-based U-Net
architectures. Note that our proposed method is applicable to all diffusion models utilizing attention layers.
2.2 Low-rank adaptation
Using trainable adapters for specific tasks has been an effective approach for fine-tuning models in the realm of natural language processing (NLP) (Houlsby al. 2019 Pfeiffer et al. 2020). Low-rank adpatation (LoRA, Hu et al. 2022 is carameter--effcicnt finc-tuning method that updates a low-rank adapter: to fine-tune a pre-trained dense weight matrix W E Rdon xdm, LoRA parameterizes the fine-tuning update AW with a low-rank factorization
W = W BA,
where B E dod Xr A E Rrxdm and r < minfdin dout).
LoRA and diffusion. Although initially proposed for fine-tuning LLMs, LoRA is generally applicable to wide range of other deep-learning modalities. Recent works used LoRA with diffusion models for various
lasks including image generation [Ry 2023 Gu et al.) 2023 Go et al., 2023 image editing (Sh et al., 2023), continual learning (Smith et al. 2023), and distillation (Golnari 2023 Wang et al. 2023b). While all these works demonstrate the flexibility and efficacy of the LoRA architecture used for fine-tuning diffusion models, the best of our knowledge, our work is the first attempt to use LoRA as part of the core U-Net for diffusion models for full training, not fine-tuning.
4

---

Published in Transactions on Machine Learning Research (09/2024)
2.3 Conditioning the score network
For diffusion models to work properly, it is crucial that the score network sp is conditioned appropriate side information. In the base formulation, the score function V.Pi(z), which the score network 89 learns, depends on the time so this i-dependence must be incorporated into the model via time conditioning. When class-labeled training data is available, class-conditional sampling requires class conditioning of the score network (Ho & Salimans 2021). To take advantage of data augmentation and thereby avoid overfitting, EDM (Karras et al. 2022 utilizes angmentation conditioning (Jun et al. 2020). where the model conditioned on the data augmentation information such as the degree of image rotation or blurring. Similarly, SDXL (Podell et al. 2024) uses micro-condiiiooiing where the network is conditioned on image resolution or cropping information. Finally, poxt-to-image diffusion models (Saharia et al., 2022 Ramesh et al. 2022 Rombach et al, 2022: Podell et al. 2024) use text conditioning, which conditions the score network with caption embeddings
80 that the model generates images aligned with the text description.
Conditioning attention layers. Prior diffusion models using CNN-based U-Net architectures condition only convolutional layers in the residual blocks by applying scale-and-shift or adaLN (see (left) of Figure 2 In particular, attention blocks are not directly conditioned in such models. This includes the state-of-the-ars diffusion models such as Imagen (Saharia al. 2022), DALL-F 2 (Ramesh et al., 2022), Stable Diffusion [Rombach et al., 2022), and SDXL (Podell et al. 2024 To clarify, Latent Diffusion Model (Rombach) et al. 2022) based models use cross-allention method for class and text conditioning, but they still utilize scale-and-shift for time conditioning,
There is a line of rescarch proposing transformer-based architecturcs (without convolutions) for diffusion models, and these work do propose methods for conditioning atlention layers. For instance, DiT (Peebles & Xie (023) conditioned attention layers using adaLN and DiffiT (Hatamizadeh al. 2023) introduced time-dependent. multi-head self-attention (TMSA), which can be viewed scalc-and-shift conditioning applied to attention layers. Although such dransformer-based archilectures have shown to be elfective, whether conditioning the attention layers with adaLN or scale-and-shift. is optimal was not investigated. Tn Section 5.5 of this work, we compare our proposed LoRA conditioning on attention layers with the prior adaLN conditioning on attention layers, and show that LoRA is the more effective mechanism for conditioning atlention layers.
Diffusion models as multi-task learners. Multi-task learning (Caruanaa 1997p framework where a single model is trained multiple related tasks simultaneously, leveraging shared representations between the tasks. If one views the denoising tasks for different timesteps (or SNR) of diffusion models related but different tasks, the training of diffusion models can be interpreted as an instance of the multi-task learning.
Following the use of trainable lightweight adapters for Mixture-of-Expert (MoE) (Jacobs et al. 1991 Ma
et al.
2018). several works have utilized LoRA as the expert adapter for the multi-task learning CCacia
et al.
2023 Wang et al. 2023a 2024 Zadouri et al.) 2024). Similarly, MORRIS (Audibert et al. 2023 and
LoRAHub Huang al. 2023 proposed using the weighted sum of multiple LoRA adapters to effectively
fackle general tasks. In this work, took inspiration from theses works by using composition of LoRA adapters to condition diffusion models.
3 Discrete-time time LoRA conditioning
Diffusion models such as DDPM Ho et al. 2020 and IDDPM (Nichol & Dhariwal 2021 have predetermined number of discrete timesteps { - 1,2, used for both training and sampling. We refer to this setting as the discrele-lime setting.
We first propose a method to condition the attention layers with LoRA in the discrete-time sctting. In particular, we implement LoRA conditioning on IDDPM by conditioning the score network with (discrete) time and (discrete) class information.
5

---

Published in Transactions on Machine Learning Research (09/2024)
3.1 TimeLoRA
TimeLoRA conditions the score network for the discrete time steps 1= 1.,...,, T. In prior architeclures, time information is typically injected into only the residual blocks containing convolutional layers. TimeLoRA instead conditions the attention blocks. See (right) of Figure
Non-compositional LoRA.
Non-compositional LoRA instantiates independent rank-r LoRA weights B1, B Br.
A1, , Ar,
The dense layer at time becomes
W = W W AW(()) = W B,A,
for 1 - ,,..., T. To clarify, the trainable parameters for each linear layer are W, A A2,..., Ar, and B1, .2...,.. Br. In particular,
trained concurrently with A1. Ar. and B1, Be...., Br.
However, this approach has two drawbacks. First, since T typically large (up to 4000), instantiating T independent LoRAs can occupy significant memory. Second, since each LoRA (Ar, B.) is trained independently, it disregards the fact that LoRAs of nearby time steps should likely be correlated/ /smilar... It would be preferable for the architecture to incorporate the inductive bias that the behavior at nearby timesteps are similar.
Compositional LoRA. Compositional LoRA composes m LoRA bases, AI...., Am and BI...., Bm where m & T. Each LoRA basis (A,, B,) corresponds time & for 1 1 < tm S T. The dense layer time becomes
W. = W AW W Σ
where Wt = ((w)) ...... ....., (w))m) is the time-dependent trainable weights composing the LoRA bases. To clarify, the trainable paramcters for cach lincar layer are W, A1, Ai...., Am, B1, BI...., Bm, and Wi.
Since the score network a continuous function of 4, we expect we v wy if / Rs l'. Therefore, to exploit the task similarity between nearby timesteps, we initialize (wi), with a linear interpolation scheme: for t; S < 1z+1:
li+1 - i i=j
1 bj
(wi);
tj
i=j+1
ty+1 -t tj
otherwise.
In short, at initialization, AW(L) uses linear combination of the two closest LoRA bases. During training, Wt learn to utilize more than two LoRA bases, i.e., Wi learn to have more than two non-zeros through training. Specifically, (WT) E m is represented as an m x T trainable table implemented as nn. Embedding in Pytorch.
3.2 ClassLoRA
Consider a conditional diffusion model with C' classes. ClassLoRA conditions the attention layers in the score network with the class label. Again, this contrasts with the typical approach of injecting class information only into the residual blocks containing convolutional layers. See (right) of Figure
Since is small for CIFAR-10 (C 10) and the correlations between different classes are likely not strong, only use the non-compositional ClassLoRA:
W. W + AW(() = W + B
for I.....C. C. In other words, each LoRA (A', (') handles a single class c. When is large, such as the case of Ik, one may consider using compositional version of ClassLoRA (See Appendix B

---

Published in Transactions on Machine Learning Research (09/2024)
4 Continuous-SNF LoRA conditioning
Motivated by (Kingma al. 2021). some recent models such as EDM (Karras al. 2022) consider parameterizing the score function as function of noise or signal-to-noise ratio (SNR) level instead of time. In particular, EDM (Karras al. 2022 considers the probability flow ODE
X; = ====((()(()00())((())) dl,
where (z is the score network conditioned on the SNR level a. We refer to this setting as the continuous- SNR setting.
The main distinction between Sections and is in the discrete vs. continuous parameterization, since continuous-timg and Con parameterizations of score functions are equivalent. We choose consider contimuous-SNR (instead of continuous-time) parameterizations for the sake of consistency with the
EDM model [Karaa et al. 2022
Two additional issues arise in the present setup compared to the setting of Section 3 First, by considering a continuum of SNR levels, there is no intuitive way to assign a single basis LoRA to a specific noise level. Second, to accommodate additional conditioning clements such as augmentations or even captions, allocating independent LoRA for each conditioning element could lead to memory inelficiency.
4.1 Unified compositional LoRA (UC-LORA)
Consider the general setting where the diffusion model is conditioned with N attributes cond), condy, which can a mixture of contimous and discrete information. In our EDM experiments, we condition the score network with N attributes: SNR level (time), class, and augmentation information.
Unified compositional LoRA (UC-LoRA) composes m LoRA bases Am. and B Bn simulta- neously condition the information of cond1, ... cond into the attention layer. The compositional weight, (Wi...., Wm) of the UC-LoRA is obtained by passing condj, condy through MLP.
Prior diffusion models typically process condj...., condy with an MLP to obtain condition embedding v, which is then shared by all residual blocks for conditioning. For the j-th residual block, w is further processed by an MLP to get scale and shift parameters 7 and By:
SharedML.(condd condyy
(7g, BB) (L)
The (1,, is then used for the scale-and-shift conditioning of the j-th residual block in the prior architectures. In our UC-LoRA, we similarly use the shared embedding v and an individual MLP for the j-th attention block to obtain the composition weight wy(v):
U = SharedMLP(cond, ...nnd
(v( (v),
Then, the j-th dense layer of the atlention block becomes
W (condd..... condx) W + AW ACoodd (condd.......o
= W + (,)
To clarify, the trainable parameters for the j-th dense layer are W, A Ag...., Am. B Bz,..., Bu and the weights MLP,. Shared across the entire architecture, the weights in SharedMLP are also trainable parameters.

---

Published in Transactions on Machine Learning Research (09/2024)
5 Experiments
In this section, present experimental findings. Section 5.1 describes the experimental setup. Section 5.2
first presents a loy, proof-of-concepl experiment to validate the proposed LoRA conditioning. Section 5.3
evaluates the effectiveness of LoRA conditioning on attention layers with a quantitative comparison between diffusion models with (baseline) conventional scale-and-shift conditioning on convolutional layers; (only LoRA) LoRA conditioning on attention layers without conditioning convolutional layers; and (with LoRA) conditioning both convolutional layers and attention layers with scale-and-shie and LoRA conditioning, respectively. Section 5.4 investigates the effect of tuning the LoRA rank and the number of LoRA bases. Section 5.5 compares our proposed LoRA conditioning with the adaLN conditioning attention layers. Section 5.6 explores the robustness of ClassLoRA conditioning compared to conventional scale-and-shift conditioning in extrapolating conditioning information.
5.1 Experimental Setup
Diffusion models. We implement LoRA conditioning on three different diffuusinn models: nano diffusion
[Llargg al.) 2024). IDDPM (Nichol & Dhariwal 2021), and EDM-vp (Karras et al.) 2022). With nano diffusion, conduct proof-of-concep experiment. With IDDPM, we test TimeLoRA and ClassLoRA for the discrete-time setting, and with EDM, we test UC-LoRA for the contimuous-SNT sctting.
Datasets. For nano diffusion, use MNIST. For IDDPM, class-conditional sampling, and ImageNet64, a downsampled version of the ImageNetlk, for unconditional sampling. For EDM-vp, we also use CIFAR-10 for both unconditional and class-conditional sampling and FFHQ64 for unconditional sampling.
CIFAR-10 for both unconditional and
Configurations. We follow the training and architecture configurations proposed by the baseline works and only tune the LoRA adapters. For IDDPM, we train the model for 500K iterations for CIFAR-10 with batch size of 128 and learning rate of 10-4, and 1.5M iterations for ImageNet64 with batch size of 128 and learning rate of 10-1 For EDM, train the model with batch size of 512 and learning rate of 1 x 10-3 for CIFAR-10, 10, and with batch size of 256 and learning rate of 2 10 for FFHQ64. For sampling, in IDDPM, we use 4000 and 4001 timesteps for the baseline and LoRA conditioning respectively, and in EDM, we use the proposed Heun's method and sample images with 18 timesteps (35 NFE) for CIFAR-10 and 40 timesteps (79 NFE) for FFHQ64. Here, NFE is the number of forward evaluation of the score network and it differs from the number of timesteps by a factor of 2 because Heun's method is a 2-stage Runge Kutta method. Appendix provides further details of the experiment configurations.
Note that the baseline works heavily optimized the hyperparameters such learning rate, dropout probability, and augmentations. Although we do not modify any configurations of the baseline and simply add LoRA conditioning in drop-in fashion, we expect further improvements from further optimizing the configuration for the entire architecture and training procedure.
LoRA. We use the standard LoRA initialization as in the original LoRA paper Hu et al. 2022 for the LoRA matrices (A, B) with rank A is initialized as Aij ~ (N and B as the zero matrix. Following Ryu 2023). we set the rank of each basis LoRA to 4. For TimeLoRA and ClassLoRA, we use 11 and 10 LoRA bases, and for UC-LoRA we use 18 and LoRA bases for CIFAR-10 and FFHQ.
Due to our constrained computational budget, we were not able to conduct a full investigation on the optimal LoRA rank the number LoRA bases. However, we experiment with the effect of rank and number of LoRA bases to limited extent and report the result in Section 5.4
5.2 Proof-of-concept experiments
We conduct toy experiments with nano diffusion for both discrete-time and Contimuous-SNR scttings. Nano diffusion is small diffusion model with CNN- U-Net architecture with no skip connections with about 500, 000 trainable parameters. We train nano diffuuson unconditional MNIST generation with
8

---

Published in Transactions on Machine Learning Research (09/2024)
3 different conditioning methods: conventional scale-and-shift, TimeLoRA, and UC-LoRA. As shown in Figure 3 conditioning with TimeLoRA or UC-LoRA yields competitive result compared to the conventional scale-and-shift conditioning.
2
0
4
2
2
6
4
7
4
8
3
+
6
T
7
9
1
6
7
9
8
5 1/
)
6
3
R
$
7
6
2
Figure 3: MNIST samples generated by nano diffusion trained with (1st row) conventional scale-and-shift conditioning; (2nd row) TimeLoRA with linear interpolation initialization; (3rd row) UC-LoRA; and (4th row) TimeLoRA with random initialization.
Initialization of wt(t) for TimeLoRA. As shown in Figure 3 the choice of initialization of (t) for TimeLoRA impacts performance. With randomly initialized wi(t), nano diffusion did not converge after 100 epochs, whereas with (t) initialized with the linear interpolation scheme, it did converge. Moreover, Figure 4 shows that even in UC-LoRA, w(() shows higher similarity between nearby timesteps than between distant timesteps after training. This consistent with our expectation that (t( R (t() if x tt.
1000
1.0
800
0.5
600
0.0
400
-0.5
200
-1.0
250
500
750
1000
250
500
750
1000
f
f1
Figure 4: Cosine similarity between w(t)) and (t2) for UC-LoRA applied to nano diffusion (left) at initialization and (right) after training. At initialization, the cosine similarity between w(t1) and w(22) has discernible pattern. After training, however, the cosine similarity between w(1)) and w(l2) for x l2 is close to 1, implying their high similarity.
5.3 Main quantitative results
Simply adding LoRA conditioning yields improvements. To evaluate the effectiveness of the drop-in addition of LoRA conditioning to the attention layers, we implement TimeLoRA and ClassLoRA to IDDPM and UC-LoRA to EDM, both with the conventional scalc-and-shift conditioning on the convolutional layers unchanged. We train IDDM with CIFAR-10, ImageNe164 and EDM with CIFAR-10, FFHQ64. As reported in Table [ the addition of LoRA conditioning to the attention layers consistently improves the image generation quality as measured by FID scores (Heusel et al., 2017) across diffcrent diffusion models and datasets with only (~10%) addition of the parameter counts. Note these improvements are achieved without tuning any hyperparameters of the base model components.

---

Published in Transactions on Machine Learning Research (09/2024)
Initializing the base model with pre-trained weights. We further test UC-LoRA pre-trained EDM base models for unconditional CIFAR-10 and FFHQ64 generations. As reported in Table 1 using pre-trained weights showed additional gain on FID score with fewer number of interations (~ 50%). To clarify, although we initialize the base model with pre-trained weights, we fully train both basc model and LoRA modules rather than finetuning.
LoRA can even replace scalo-and-shift. We further evaluate the effectiveness of LoRA conditioning by replacing the scale-and-shift, conditioning for the convolutional layers in residual blocks with LoRA conditioning for the attention blocks. The results of Table | suggest that solely using LoRA conditioning attention layers achieves competitive FID scores while being more efficient in memory compared to the baseline score network trained with scale-and-shif conditioning on convolutional layers. For IDDPM, using LoRA in place of the conventional scale-and-shift conditioning consistently produces better results. Significant improvement observed especially for class-conditional generation of CIFAR-10. For EDM, replacing the scale-and-shif conditioning did not yield an improvement, but nevertheless performed comparably, We note that in all cases, LoRA conditioning is more parameter-efficiet (~10%) than the conventional scale-and-shift conditioning.
5.4 Effect of LoRA rank and number of LoRA bases
We investigate the effect of tuning the LoRA rank and the number of LoRA bases on the EDM model for unconditional CIFAR-10 generation and report the results in Table [ Our findings indicate that using LoRA bases consistently improves the quality of image generations. On the other hand, increasing LoRA rank docs not guarantee better performance. These findings suggest an avenue of further optimizing and improving our main quantilative results of Section 5.3 and Table which we have not yet been able to pursue due to our constrained computational budget.
# basis
rank
FID
# Params
1.99
57185519
Varying basis
18
4
1.96
57745499
36
4
1.95
58865459
18
2
1.93
57192539
Varying rank
18
4
1.96
57745499
18
8
1.96
58851419
Table 2: Effect of the number of LoRA bases and the LoRA rank on unconditional CIFAR-10 sampling of EDM with LoRA
5.5
Comparison with adaLN
We compare the effectiveness our proposed LoRA conditioning with adaLN conditioning applied to attention layers. Specifically, we conduct an experiment on EDM with scale-and-she conditioning on convolutional layers removed and with (i) adaLN conditioning attention layers or (ii) LoRA conditioning attention layers. We compare the sample quality of unconditional and class-conditional CIFAR-10 generation and report the results in Table B We find that LoRAA conditioning significantly outperforms adaLN conditioning for both unconditional and conditional CIFAR-10 generation. This indicates that our proposed LoRA conditioning is the more effective mechanism for conditioning atteitoon layers in the U-Net architectures for diffusion models.
Type
uncond.
cond.
adaLN conditioning 2.16
2.0
LoRA conditioning
1.99
1.82
Table 3: Comparison of adaLN conditioning and LoRA conditioning attention layers EDM (without conditioning convolutional layers). We consider both unconditional and conditional CIFAR-10 generation.
10

---

Published in Transactions on Machine Learning Research (09/2024)
5.6 Extrapolating conditioning information
We conduct experiment comparing two class-conditional EDM models each conditioned by scale-and-shift and ClassLoRA, for the CIFAR-10 dataset. During training, both models receive size-10 one-hot vectors (ca), = 6ij representing the class information.
First, we input. the linear interpolation ore; +eg+00 +(1++(1000 (1-a)e (0 S a S 1) of two class inpuls e; and ej (corresponding 'airplane' and 'horse', respectively) to observe the continuous transition between classes, As shown in the top of Figure 5 both the scale-and-shift EDM and ClassLoRA EDM models effectively interpolate semantic information across different classes. However, when scaled imput Se, is received, with & ranging from -1 to 1, scale-and-shift EDM generates unrecognizable images when B < 0, while ClassLoRA EDM generates plausible images throughout the whole range, as shown in the bottom of Figure 5 This toy experiment shows that LoRA--asee conditioning may be more robust to extrapolating conditioning information beyond the range encountered during training. Appendix D provides further details.
Figure 5: Results of (Top) interpolation of class labels in class-condition EDM with (row1) ClassLoRA; (row2) scale-and-shift; (bottom) extrapolation of class labels in class-conditional EDM with (rowl) Class- LoRA; (row2) scale-and-shift
6 Conclusion
In this work, we show that simply adding Low-Rank Adaptation (LoRA) conditioning the attention layers in the U-Net architectures improves the performance of the diffusion models. Our work shows that we should condition the attention layers in diffusion models and provides a prescription for effectively doing so. Some prior works have conditioned attention layers in diffusion models with adaLN or scale-and-shift operations, but we find that LoRA conditioning is much more effective as discussed in Section 5.5
Implementing LoRA conditioning on different and larger diffiusion model architectures a natural and interesting direction of future work. Since almost all stale-of-the-arl (SOTA) or near-SOTA diffusion models utilize attention layers, LoRA conditioning is broadly and immediately applicable all such architectures. In particular, incorporating LoRA conditioning into largo-scale diffusion models such Imagen (Saharia
et al.
2022 DALL-E [Ramesh et al. 2022). Stable Diffusion (Rombach et al. 2022 e. and SDXL (Podell
et al.
2024 or transformer-basee diffusion models such as U-ViT (Bao et al. 2023 DiT (Peebles & Xie
2023), and DiffT (Hatamizadeh et al.) (023)) are interesting directions. Finally, using LoRA for the text conditioning of text-to-image diffusion models is another direction with much potential impact.
7
Acknowledgements
This research was supported by a grant from KRAFTON AI. AN was supported in part by the Ministry of Science and ICT (MSIT), South Korea, under the Information Technology Research Center (ITRC) Support Program 2LTTP-2024RSS222 2022-001669... EKR was supported by the National Research Foundation of Korca(NRF) grant funded by the Korean government (No.RS-2024-00-2) Dongwhan Rho for their careful reviews and valuable feedback. We also thank the anonymous reviewers for their insightful comments.
We thank Schyun Kwon and
11

---

Published in Transactions on Machine Learning Research (09/2024)
References
Brian D.O, Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 1982.
Alexandre Audibert, Massih R Amini, Konstantin Usevich, and Marianne Clausel. Low-rank updates of pre-trained weights for multi-task learning. ACL, 2023.
Fan Bao, Shen Nic, Kaiwen Xuc, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A VIT backbone for diffusion models. CVPR, 2023.
Lucas Caccia, Edoardo Ponti, Zhan Su, Matheus Pereira, Nicolas Le Roux, and Alessandro Sordoni. Multi-head adapter routing for cross-task generalization. NewrIPS, 2023.
Rich Caruana. Multitask learning. Machine Learning, 1997.
Shoufa Chen, Chongjian GE, Zhan Tong, Jiangliu Wang, Yibing Song, Jue Wang, and Ping Luo. Adaptformer: Adapting vision transformers for scalable visual recognition. NewrIPS, 2022.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 2012.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. NeurIPS, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image worth 16x16
words: Transformers for image recognition at scale. ICLR, 2021.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusson transformer is strong image synthesizer. ICCV, 2023.
Hyojun Go, Yunsung Lee, Jin-Jonnn Young Kim, Seunghyun Lee, Myeongho Jeong, Hyun Seung Lee, and Seungtaek Choi. Towards practical plug-and-play diffusion models. CVPR, 2023.
Pareesa Amench Golnari. LoRA-Fnhanced distillation on guided diffusion models. ar preprint
arXiv:2312.06892 2023.
Yuan Gong, Hongyin Luo, Alexander H Liu, Leonid Karlinsky, and James Glass. Listen, think, and understand. ICLR, 2024.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. ICL.R 2022.
Yuchao Gu, Xintao Wang, Jay Zhangjic Wu, Yujun Shi, Yunpeng Chen, Zihan Fan, Wuyon Xiao, Rui Zhao, Shuning Chang, Weijia Wu, Yixiao Ge, Ying Shan, and Mike Zheng Shou. Mix-of-show: Decentralized low-rank adaptation for multi-concept customization of diffusion models. NewrIPS, 2023.
Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. DiffIT: Diffusion vision transformers for image generation. arXiv preprint 2xxi::3122.02399 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to local nash equilibrium. NeurIPS, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. NewrIPS Workshop on Deep Generative Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NewrIPS, 2020.
Emiel Hoogeboom, Jonathan Hoek, and Tim Salimans. Simple diffusion: End-to-end diffusion for high resolution images. ICML, 2023.
12

---

Published in Transactions on Machine Learning Research (09/2024)
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficieu transfer learning for NLP. TCML, 2019.
Edward J. 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. ICLR, 2022.
Chengsong Huang, Qian Liu, Bill Yuchen Lin, Tianyu Pang, Chao Du, and Min Lin. LoraHub: Efficient. cross-task generalization via dynamic LoRA composition. arXie preprint arXiv:2307.1366. 2023.
Allan Jabri, David J Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. ICML, 2023.
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Newral Computation, 1991.
Heewoo Jun, Rewon Child, Mark Chen, John Schulman, Aditya Ramesh, Alec Radford, and Tlya Sutskever. Distribution augmentation for generative modeling, ICML, 2020.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. CVPR, 2019.
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. NewrIPS, 2022.
Dongjun Kim, Yeongmin Kim, Se Jung Kwon, Wanmo Kang, and I1-Chul Moon. Refining generative process with discriminator guidance in basoe diffusion models. ICML, 2023.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. NewrIPS, 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Mingi Kwon, Jaeseok Jeong, and Youngiung Uh. Diffusion models already have a semantic latent space. ICLR, 2023.
Mare Lelarge, Andrei Bursuc, and Jill-Jênn Vie. Dataflowr. 2024.
Yan-Bo Lin, Yi-Lin Sung, Jie Lei, Mohit Bansal, and Gedas Bertasius. Vision transformers are parameter- efficient andio-visual Iearners. CVPR, 2023.
Jiaqi Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H. Chi. Modeling task relationships in multi-task learning with multi-gate mixture-ofexxpers. ACM SIGKDD, 2018.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. ICML, 2021.
Junting Pan, Ziyi Lin, Xiatian Zhu, Jing Shao, and Hongsheng Li. St-adapter: Parameter-efficient image-to- video transfer learning. NewrIPS, 2022.
William Peebles and Saining Xie. Scalable diffusion models with transformers. ICCV, 2023.
Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville, FiLM: Visual reasoning with a general conditioning layer. AAAI, 2018
Jonas Pfeiffer, Andreas Rücklé, Clifton Poth, Aishwarya Kamath, Ivan Vulić, Sebastian Ruder, Kyunghyun Cho, and Iryna Gurevych. AdapterHub: A framework for adapting transformers. EMNLP, 2020.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent dilfusion models for high-resolution image synthesis. ICLR, 2024.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hicrarchical text-conditional image generation with clip latents. Xiv preprint ar Xiv:2204.061- 2022.
13

---

Published in Transactions on Machine Learning Research (09/2024)
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. CVPR, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. MICCAI, 2015.
Simo Ryu.
Low-rank adaptation for fast text-to-image
diffusion
fine-tuning.
https:// githtutthhhggg https://github.com cot.auuccccmmmoooooffsiioooo/// /ora 2023.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J Fleet, and Mohamnad Norouzi. Photorealistic Lext-lo-image diffusion models with deep language understanding. NearIPS, 2022.
Yujun Shi, Chuhui Xue, Jiachun Pan, Wenqing Zhang, Vincent YF Tan, and Song Bai. DragDiffusion: Harnessing diffusion models for interactive point- point-based image editing, arXiv preprint arXic:2306. 14435, 2023.
James Seale Smith, Yen-Chang Hsu, Lingyu Zhang, Ting Hua, Zsolt Kira, Yilin Shen, and Hongxia Jin. Continual diffusion: Continual customization of text-lo-image diffusson with c-lora. arXiv preprinl arXiv:2304.06027, 2023.
Jascha Sick-Dick,tein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. ICML, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. ICLR, 2021a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. NeurIPS, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modcling through stochastic differential equations. ICLR, 2021b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NewrIPS, 2017.
Haowen Wang, Tao Sun, Cong Fan, and Jinjic Gu. Customizable combination of parametor-cfficiet modules for multi-task learning. ICLR, 2024.
Yiming Wang, Yu Lin, Xiaodong Zeng, and Guannan Zhang. Multilora: Democratizing lora for better multi-task learning, arXiv preprint arzXi::211.11100 2023a.
Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. ProlificDreamer: High-fidelity and diverse text-to-3D generation with variational score distillation. NewrIPS, 2023b.
Yuxin Wu and Kaiming He. Group normalization. ECCV, 2018.
Jing Nathan Yan, Jiatao Gu, and Alexander M Rush. Diffusson models without attention. orXio preprinl arXio22311.18257, 2023.
Ted Zadouri, Ahmet Üstün, Arash Ahmadian, Beyza Ermiş, Acyr Locatelli, and Sara Hooker. Pushing mixture of experts the limit: Extremely parameter efficient moe for instruction tuning. ICLR, 2024.
14

---

Published in Transactions on Machine Learning Research (09/2024)
A Experimental details
Here, we provide the detailed experiment settings. Note that apart from the conditioning method, we mostly follow base codes and configurations provided by Dataflowr nano diffusion, IDDPM and EDM respectively.
IDDMM repository and EDM repository for
A.1 Nano Diffusion
We mostly followed the base code provided by Dataflowr with exceptions. First, the simusoidal embedding implemented in the original code not correctly implemented. Although it did not have visible impact on TimeLoRA and UC-LoRA, it significantly deteriorated the sample quality of the scale-and-shift conditioning (see Figure Second, during training process, the input image is normalized with mean = 0.5 and std = 0.5. However, is not considered in the visualization process. Lastly, we extended the number of training epochs from 50 to 100 for better convergence.
6
9
3
3
5
9
9
7
4
g
5.
J
e
9
Figure 6: MNIST image generated with incorrect sinusoidal embedding and scale-and-shift conditioning, (top) after 100 epochs of training, (bottom) after 400 epochs training.
A.2 IDDPM
Here we provide the training setting used IDDPM experiments based on the IDDPM repository
A.2.1 CIFAR-10
For CIFAR-10 training, we construct U-Net with model channel of 128 channels, and 3 residual blocks per each U-Net blocks. We use Adam optimizer with learning rate of x 10-4, momentum of (81, (2) (0.9, 0.999), weight decay, and dropout rate of 0.3. We train the model with Luybrid proposed in the original paper for 500k iterations with batch size of 128. The noise is scheduled with the cosine scheduler and the timestep is sampled with uniform sampler at training. For sampling, we use checkpoints saved every 50k iterations with exponential moving average of rate 0.9999, and sample image for 4000 steps and 4001 steps for the baseline and LoRA conditioning, respectively.
Training flags used for unconditional CIFAR-10 training.
MODET_ -TAGS="--Imas se --num_channe la 128 --num_reee res_ o loks ---earr_siggne sigma True --dropous 0.3*
OLTFUSION_FLAGS"" FLAGS-" diffusion_stepe 4000 joise_schedule cosine"
PRAIN_PLAGS- --1 le-d --batch_eize 123*
Training flags used for class-conditional CIFAR-10 training.
XODEL "ossss---mmmmm -im--mmmaaeeeeee size --num_channele lesnu_rre__beeem --Learn_sigga True --dropous --class_cond DIFFUSION, -iifuucioon 4000 scchedle cosine"
16-4 --batch_stz2 21 128*
A.2.2 ImageNet64
For ImageNet64 training, we use the same U-Net construction used in CIFAR-10 experiment. The model channel set as 128 channels, and cach U-Net block contains residual blocks. We use Adam optimizer with learning rate of 10-4, momentum (8,, B2) (0.9, 0.999), no weight decay and no dropout. We train the model with Taybeid proposed in the original paper for 1.5M iterations with batch size of 128. The noise
15

---

Published in Transactions on Machine Learning Research (09/2024)
is scheduled with the cosine scheduler and the timestep is sampled with uniform sampler at training. For sampling, use checkpoints saved every 500k iterations with exponential moving average of rate 0.9999, and sample images for 4000 steps and 4001 steps for the baseline and LoRA conditioning, respectively.
Training flags used for ImageNct64 training.
FEDEL_FLAAS-- -aiAs====iimmgo...... size 64 --num_chane channela 128 wewuunreee_lloeee -=leern_oigga True" DIPFOSTON, tffusiinnn 400d --nol.se_schedull coaine" TRAIN_FLAGS=" batch_size 128*
A.3
EDM
Here we provide the training setting used in EDM experiments based on the EDM repository
A.3.1 CIFAR-10
For CIFAR-10 training, we use the DDPM++ (Song et al. 2021bb with channel per resolution as 2, 2, 2. We use Adam optmizer with learning rate of 1 x 10-h, momentum of (81,32) = (0.9, 0.999), no weight decay and dropout rate of 0.13. We train the model with EDM preconditioning and EDM loss proposed in the paper for 200M training images (counting repetition) with batch size of 512 and augmentation probability of 0.13. For sampling, we used checkpoints saved every 50 iterations with exponential moving average of 0.99929, and sample image using Heun's method for 18 steps (35 NFE).
Training flags used for unconditional CIFAR-10 training.
--cond=0
-aarchaddpmpp
Training flags used for class-conditional CIFAR-10 training.
cond=1
azch-ddpapp
A.3.2 FFHQ64
For FFHQ training, we use the DDPM++ (Song al. 2021b) with channel per resolution 1, 2, 2, 2. We use Adam optmizer with learning rate of 2 10 momentum of (B1,82) (0.9, 0.999), weight decay and dropout rate of 0.05. We train the model with EDM preconditioning and EDM loss proposed in the paper for 200M training images (counting repetition) with batch size of 256 and augmentation probability of 0.15. For sampling, we used chockpoints saved every 50 iterations with exponential moving average about 0.99965, and sample image using Heun's method for 40 steps (79 NFE).
Training flags used for FFHQ training.
--oond=D -aarchaddpmpp --5 --preg=1, 2,,,, --1r=2e-4 --dropout=c.0s ->augment=8.15
A.4 LoRA basis for TimeLoRA
As IDDPM uniformly samples the timestep during training, we follow the similar procedure for selecting the
timestop f; assigned for the LoRA basis (A,, B,) in TimeLoRA. To be specific, we set f = 1, fm = T, and equally distribute &; in between:
T-1 m
4 1 (i
Note here, for simplicity we assumed that m divides T and choose T 4001 instead of T 4000 used in the baseline work.
16

---

Published in Transactions on Machine Learning Research (09/2024)
A.5 MLP for composition weights of UC-LoRA
For the composition weights of UC-LoRA, we used 3-layer MLP with group normalization and SiLU aclivation. Specifically, each LoRA module contains MLP consisting of two linear layers with by group normalization and SILU activation followed by a output lincar layer. MLP takes the shared condition embedding e E R as the imput and outputs the composition weight w(v) E R", where demb is the embedding dimension (512 for FDM) and m is the number of ToRA bases. This implemented in Pytorch as
self comp_oomppeee weights on Sequentiall
Linear (in_ febeertreseeeeeeddie GroupNorm tnum_channels-Nl, ops=le-6), torch. (1)11101),
vut_Eeatures-128),
Linear (in_ teatures-k1, out teatureeeee GroupNorm channels N2, sps-le 8), torch.ne na. (n.si1u0), (), Linear featuress,,,, uul_featuree--mnn_aasse
In our experiment, we set (NI, = (50, 50) for nano diffusion and (NI, = (128, 64) for EDM. Note the choice of depth, and the width of the MLP is somewhat arbitrary and can be further optimzied.
B Compositional ClassLoRA
To verify the effectiveness of a compositional version of ClassToRA, we conducted a proof-of-concepl experiment with EDM on CIFAR-10. We finetuned a pretrained unconditional EDM for class-conditinnal CIFAR-10 sampling using a single LoRA basis and randomly initialized compositional weights. Note in the non-compositional ClassLoRA, we used 10 LoRA bases, one for each class. As result, with the addition of only 61560 parameter counts, EDM and improve the FID score from 1.97 to 1.81.
able to convert (fine-tune) an unconditional EDM to conditional
C Cosine similarity between wis in nano diffusion
We present cosine similarity between wy and Wsoo for all LoRA bases in UC-LoRA for nano diffusion in Figure As observed in Section 5.2 and Figure I cosine similarity is consistently high for i close to 500, proving that there is a task similarity between nearby timesteps for all layers of the network. However, the patterns varied depending on the depth of the layer within the network. This could be an interesting point for future research to further understand the learning dynamics the diffusion models.
17

---

Published in Transactions on Machine Learning Research (09/2024)
1.00
CosSmllat) ((000))
0.75
CosSmiul)) w.15000))
0.50
0.25
0.00
-0.25
at intt
-0.5
init
-0.50
after training
training
200
600
1000
200
400
600
800
1000
Time step
Time step
Down block
Up block 3
1.0
CosSinturt), w/5001
0.8
CosSimtu(()) w(5000))
0.8
0.6
0.6
0.4
0.2
0.2
0.0
init
at kiiit
-0.2
training
training
200
400
800
1000
200
400
600
800
1000
Time step
Time step
Down block 2
Up block 2
1.00
CosSinto (5000)
CosSimtutff..ufB (f)))))))))
ConSimfcot)..c000000 c1130))
0.5
0.50
0.26
0.0
0.00
-0.5
after training
init
0.6
at Init after training
-0.25
at ink after training
-0.50
200
400
600
800
1000
0
200
400
600
800
1000
200
400
600
800
1000
Time step
Time step C
Time step t
Down block 3
Middle block
Up block 1
Figure 7: Cosine similarity between wi and w500 for UC-LoRA applied to nano diffusion.
D Extrapolating conditioning information with Class LoRA
We present a more detailed analysis along with comprehensive results of the experiments introduced Section 5.6 The class-conditional EDM model used for comparison was trained with the default training configurations explained in Appendix A.3 For ClassLoRA-con ClassLoRA-conditionee EDM, we used the unconditional EDM as the base model and applied ClassLoRA introduced in Section 3.2 for the discrete-time setting. We did not use UC-LoRA for this ablation study, with the intent of focusing the effect of LoRA conditioning class information.
We provide interpolated and extrapolated images as introduced in 5.6 for various classes in Figure and Figure g corroborating the consistent difference between the two class conditioning schemes across classes. We also experiment with strengthening the class input, where we use scaled class information input 80, with I0 shows that LoRA conditioning shows more robustness in this range as well.
3 >
Fig.
Considering the formulation of ClassLoRA, interpolating scaling the class input is equivalent to interpolating or scaling the class LoRA adapters, resulting in formulation similar to Compositional LoRA or UC-LoRA:
W. = W + (W = W + Σ
w. B A'
18

---

Published in Transactions on Machine Learning Research (09/2024)
where Wa corresponds to the interpolation or scaling of the class inputs. From this perspective, the weights could be interpreted as a natural 'latent. vector' capturing the semantic information of the image, in a similar sense that was highlighted in Kwon et al. (2023). While our current focus was exclusively on class information, we hypothesize that this method could be extended to train style or even text conditioning, especially considering the effectiveness of LoRA for fine-tuning.
Class-conditional EDM with ClassLoRA
Class-conditional EDM with scale ande-fnd-shift
Figure 8: Results of interpolation of class labels for various classes
19

---

Published in Transactions on Machine Learning Research (09/2024)
(a) Class-conditional EDM with ClassLoRA
(b) Class-conditional EDM with scale-and-shift.
Figure 9: Results of extrapolation of class labels for various classes
20

---

Published in Transactions on Machine Learning Research (09/2024)
(a) Class-conditional EDM with ClassLoRA
(b) Class-conditional EDM with scale-and-shift.
Figure 10: Results of class labels scaled from 1 to 5 for various classes
21

---

Published in Transactions on Machine Learning Research (09/2024)
E Image generation samples
E.1
IDDPM
E.1.1
Unconditional CIFAR-10
Figure 11: Unconditional CIFAR-10 samples generated by IDDPM with LoRA
E.1.2 CIFAR-10
(a) Bird
(b) Car
Figure 12: Class conditional CIFAR-10 samples generated by IDDPM with LoRA
22

---

Published in Transactions on Machine Learning Research (09/2024)
(c) Cat
(d) Deer
(e) Dog
(I) Frog
Figure 12: Class conditional CLEAR-10 samples generated by IDDPM with LoRA
23

---

Published in Transactions on Machine Learning Research (09/2024)
Horse
(h) Plane
i Ship
(i) Truck
Figure 12: Class conditional CIFAR-10 samples generated by IDDPM with LoRA
24

---

Published in Transactions on Machine Learning Research (09/2024)
E.1.3 ImageNet64
Figure 13: ImageNet64 samples generated by IDDPM with LoRA
25

---

Published in Transactions on Machine Learning Research (09/2024)
E.2
EDM
E.2.1
Unconditional CIFAR-10
Figure 14: Unconditional CIFAR-10 samples generated by EDM with LoRA
E.2.2
Class-conditional CIFAR-10
(a) Bird
(b) Car
Figure 15: Class conditional CIFAR-10 samples generated by EDM (vp) with LoRA
26

---

Published in Transactions on Machine Learning Research (09/2024)
(c) Cat
(d) Deer
(e) Dog
(f) Frog
Figure 15: Class conditional CIFAR-10 samples generated by EDM (vp) with LoRA
27

---

Published in Transactions on Machine Learning Research (09/2024)
Horse
(h) Plane
ii Ship
(i) Truck
Figure 15: Class conditional CIEAR-10 samples generated by EDM (vp) with LoRA
28

---

Published in Transactions on Machine Learning Research (09/2024)
E.2.3 FFHQ
Figure 16: FFHQ samples generated by EDM (vp) with LoRA
29


Reviews:
[{"review_id": "BvsepIDh6L", "reviewer": "Reviewer_f6Uc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The aim of the paper is straightforward and clear. Including timestep information in the attention layers of diffusion-based UNets seems to be a missing component of the current paradigm and it is useful to explore ways to do so.\n* The proposed LoRA method is efficient and aligns nicely with existing diffusion techniques. The handling of LoRA module design makes intuitive sense and seems natural.\n* The experiments are conducted in a clear and reproducible way. There is solid evidence that consistent (although modest) refinement can be achieved by including the proposed modules.\n  \nWeaknesses:\n* The introduction of two different LoRAs techniques for two different cases seems somewhat unnecessary. Discrete timesteps can be treated as continuous inputs, and classes can be represented as vectors in one-hot or binary encoding, etc. which can also be treated as continuous inputs. I would be interested in seeing evidence whether UC-LoRA can match Time and Class LoRA in the cases where Time and Class LoRA were used. If so, it might be more straightforward to adopt a single approach. This is not a major concern, but it is worth exploring.\n* I am not sure how interpolation between classes would be performed for Class LoRA in the case of a large number of classes. The interpolation method for timesteps seems to depend on an ordering of the timesteps, which doesn't hold for classes. I would be interested in hearing how handling a large number of classes could be done.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work examines the problem of including scalar/vector conditional information such as timestep and class label in the attention layers of UNet-based diffusion models. The majority of current methods include this type of conditional information only in the residual layer of the UNet, which seems suboptimal. The paper proposes to use LoRA to adapt the input QKV and output projection layers of UNet attention modules, with the motivation that this provides a lightweight and constrained adjustment to the operation basic attention layer. The major contribution of the work is designing two types of LoRA layers: one for discrete ordered timesteps or classes, and the other for continuous timesteps and continuous conditions. Experiments are conducted showing that including the proposed modules provides improvement in generative quality metrics.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have no major requested changes. It would be interesting to see whether UC-LoRA can serve as a single solution for all the cases considered, but this is not essential.", "broader_impact_concerns": "Broader impacts were not discussed, but that does not influence my view of the paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YumI2PBccp", "reviewer": "Reviewer_Rtei", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The fact that this is a drop-in improvement makes is a strong point, since otherwise, with architectural improvement, it can be difficult to tell if there is a meaningful improvement if there is significant hyper-parameter tuning involved.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper suggests that LoRA blocks should be integrated within a U-net block as a part of a conditioning process. The authors show that architectural change leads to a small improved performance w.r.t the FID metric, at the cost of some additional parameters.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think the paper is overall well written and the figures are nice. I was not particularly knowledgeable about this topic, but the paper did a good enough job of helping me understand the topic. I would say there could be some more information in the introduction about the theoretical motivation behind the paper.", "broader_impact_concerns": "None to note.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2vVE5P2ydm", "reviewer": "Reviewer_XaAm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: Using LoRA in this problem is novel and leads to promising result\n\nWeakness: Lack of theory, and the motivation is not too strong. For the present manuscript, the motivation is just (paraphrased) \"we should condition the attention layers because we also conditioned other.\" This is not a strong motivation for conditioning the attention layers, nor a strong motivation to condition the attention in the specific way the authors proposed", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes to apply drop-in lora (low rank layer) to the attention layers of the attention layers of a U-net improves the performance of the image generation tasks", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I believe having a more convincing and not-so-arbitrary motivation would be much better", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a new technique of including conditioning in the parameterization of diffusion models using LoRA for UNets. The paper has a clear objective and well-defined contributions for which enough empirical results are presented. 

The main strengths of the paper are the following:
- The paper has a clear goal of injecting timestep information in attention layers of UNet-based diffusion models to improve the model ability to adapt to different denoising timesteps and improve sampling results.
- The proposed LoRA method is efficient. 
- The experiments are conducted in a clear and reproducible way. There is solid evidence that consistent (although modest) refinement can be achieved by including the proposed modules.

There are no major weaknesses.

All reviewers vote to accept the paper. I agree with them and recommend accepting the paper.
Decision: Accept (poster)