Fine-tuning Diffusion Models with Limited Data
Abstract
Diffusion models have recently shown remarkable progress, demonstrating state-ofthe-art image generation qualities.
Like the other high-fidelity generative models, diffusion models require a large amount of data and computing time for stable training, which hinders the application of diffusion models for limited data settings.
To overcome this issue, one can employ a pre-trained diffusion model built on a large-scale dataset and fine-tune it on a target dataset.
Unfortunately, as we show empirically, this easily results in overfitting.
In this paper, we propose an efficient fine-tuning algorithm for diffusion models that can efficiently and robustly train on limited data settings.
We first show that fine-tuning only the small subset of the pre-trained parameters can efficiently learn the target dataset with much less overfitting.
Then we further introduce a lightweight adapter module that can be attached to the pre-trained model with minimal overhead and show that fine-tuning with our adapter module significantly improves the image generation quality.
We demonstrate the effectiveness of our method on various real-world image datasets.
Introduction
Diffusion models [36, 14] have become a mainstream approach in image synthesis, quickly taking over the crown of Generative Adversarial Networks ( GAN s)
[11] in the literature.
Compared to
Variational Autoencoder ( VAE ) variants [21, 42, 5] or autoregressive models [43, 9] which are trained to maximize likelihoods instead of directly optimizing image quality, diffusion models can be trained to generate more realistic images with less number of parameters [14, 27].
Similar to the other generative models, diffusion models require tons of data and training times to produce diverse images of high quality.
Even a simple diffusion model would require several GPU days to train on CIFARor ImageNet benchmark.
This computational cost acts as a bottleneck, limiting the scalability of diffusion models to applications where data are scarce.
To resolve this issue common in high-fidelity generative models, several approaches have been proposed.
We offer a detailed review of recent works in Appendix B.
However, for diffusion models, except for fine-tuning methods for conditional generation [35, 10], there has been little work on fine-tuning pre-trained unconditional diffusion models for target domain with limited data.
In this paper, we propose a novel fine-tuning algorithm for unconditional diffusion models under limited data settings.
This is a challenging task because it is prone to overfitting due to small data and any regularization technique depending on conditional inputs is not available unlike conditional settings [32].
A na¨ve fine-tuning method of updating all weights easily results in overfitting, so as done in existing transfer learning studies [45, 41, 25], a thorough study on which weights to freeze and tune is required.
Through empirical analysis, we first identify that the attention modules common in the diffusion models are crucial for quick adaptation to the target domain, and fine-tuning only the attention modules could significantly alleviate the overfitting.
This is somewhat surprising considering that the attention modules take only a small fraction (about 10%) of parameters in the entire network for the diffusion model.
Inspired by the adapter approach for language models [16], we further design a novel adapter module that can be attached to the pre-trained model and fine-tuned with negligible additional parameters.
Unlike the adapter module commonly applied for language models, due to the nature of the diffusion models, our adapter module is time-aware, and we find this time awareness is crucial for successful fine-tuning.
Especially, we find that the adapter module plays an important role under a large domain gap between source pre-training and target fine-tuning data.
Combining the idea of fine-tuning only the attention modules in the pre-trained model and the fine-tuning with the adapter module, we present Adapter-Augmented Attention Fine-tuning ( A 3 FT ) which efficiently learns on target tasks with limited data, converging faster with much less overfitting than existing fine-tuning.
We empirically demonstrate the effectiveness of our method on real-world image synthesis tasks.
Methods
Naïvely fine-tuning entire parameters of a pre-trained model with limited target data would be a straightforward approach, but it is well-known that this is prone to overfitting [46, 26].
As we empirically validate this for diffusion models in Fig.
1, the naïve fine-tuning causes overfitting in the early stage of training.
In Appendix D.1, we further verify that the quality and diversity of the images generated from the overfitted models are seriously damaged, and the models tend to memorize few images in target training data.
2.2 Attention block fine-tuning
Most of the modern diffusion models employ U-Net [31] as a backbone architecture, which consists of a stack of residual and attention blocks intertwined with skip connections.
At a high level, a U-Net is divided into an encoder and a decoder, where the encoder gradually downsamples images into multi-scale features, and the decoder upsamples them back to the original sizes.
Instead of fine-tuning all the parameters, we fine-tune attention blocks while freezing residual blocks.
As shown in Fig.
1, the attention block fine-tuning works surprisingly well, significantly outperforming naïve fine-tuning.
Especially, attention block fine-tuning is not only more stabilized, but also achieves better image quality in terms of cFID score [29].
Attention block tuning is more memory-efficient since it only takes 10.3% of the total parameters.
2.3 Time-aware adapter
While the attention block fine-tuning works well, as demonstrated by the adapter-based methods for language models [16, 17], it is often beneficial to attach additional modules to the pre-trained model for better adaptation to target data.
Inspired by this, we develop a novel adapter module tailored for diffusion models.
Importantly, our adapter is time-aware , since the networks being used for diffusion models should produce different outputs with respect to different time steps.
As depicted in Fig.
2, the time-aware adapter is placed between the attention block and its successive linear layer so as to modulate the intermediate feature vector given the target data.
Compared to the conventional adapters, our adapter additionally takes time embedding vector t embed ∈ R d t as an input along with feature tensor z in ∈ R c
z × h
z
× w z , hence the name time-aware.
A time-aware adapter is composed of three submodules; a time-fusion module, a cross-attention module, and a time-scaling module.
Time-fusion module Time-fusion network combines the time-embedding vector t embed and an input feature vector z in into another feature vector through feed-forward layers, which is then layernormalized [1].
Refer to Appendix C.4 for details.
The resulting feature vector z query ∈
R c z × h
z
× w z is fed into the cross-attention module as a query.
That is,
[Equation]
where LN( · ) denotes layer normalization and FF( · ) denotes feed-forward layers.
Then, crossattention module takes z query as a query, and z in as both key and value vectors with the same dimension.
We then apply channel-wise cross attention followed by layer normalization to produce an output tensor z out ∈ [Equation]
× [Equation]
z
[Equation]
Time-fusion modules use approximately 1.7M parameters in total ( ≈ 1.4% of the pre-trained model).
Time-scaling module Motivated by Hong et al.
[15], the time-scaling module computes the channelwise mixture factor α t ∈ [0 , 1] c z that regulates the information generated by the time-aware adapter.
Here, the mixture factor α t is also time-aware so that the effect of the adapter output can vary along time steps.
The final output of the time-aware adapter is computed as
[Equation]
where ⊙ denotes the elementwise multiplication.
Time-scaling modules use approximately 0.5M parameters ( ≈ 0.4% of the pre-trained model).
Adapter-augmented attention fine-tuning
Combining the attention block fine-tuning and time-aware adapter, we present Adapter-Augmented
Attention Fine-tuning ( A 3 FT ) which fine-tunes attention blocks along with the time-aware adapter attached to the pre-trained model.
Compared with the conventional adapter-based methods [16], we adapt the weight-sharing for the time-aware adapters.
Based on the weight-sharing of our modules, we observe that it not only alleviates the overfitting in our case but also accelerates the convergence.
We will demonstrate through experiments in Section 3, the time-aware adapter can seamlessly be trained with the attention blocks, leading to a significant performance gain across datasets.
Experiments
Our method is tested on both low and high resolution datasets: 64 × 64 resolution (CelebA [24], Pokemon [34]) and 256 × 256 resolution (MetFaces [18]).
See Appendix C for further experimental details.
Table 1 summarizes our main results.
We validate that naïve fine-tuning easily overfits to small-size training datasets yielding high cFID and KID.
On the other hand, attention-block finetuning largely outperforms naïve fine-tuning on both datasets.
Further, with the help of time-aware adapter, A 3 FT achieves the lowest cFID and KID score.
In Fig.
3, we notice the distinct difference in sample quality.
While naïvely fine-tuned model produces relatively plain and similar images, A 3 FT yield both diverse and high-quality images.
We provide the sampled images on all datasets in Appendix D.2.
[Table]
Attention-block fine-tuning 7.
{+A 3 FT 6 .0 .38 .
+}
Conclusion
In this paper, we empirically showed that naïvely fine-tuning diffusion models incurs overfitting in the early stage of training.
As a solution, we discovered that fine-tuning only attention blocks can be a simple yet solid fine-tuning technique.
To enhance the fine-tuning performance, we introduced the time-aware adapter that integrates a time-embedding vector so as to fit unconditional diffusion models.
We further proposed a novel fine-tuning method, Adapter-Augmented Attention Fine-tuning, which fine-tunes attention blocks in conjunction with the time-aware adapter.
Experiments validated that our approach enables faster and more stable fine-tuning with limited data as well as retaining high-quality samples, especially, for a large domain gap between the pre-training and the fine-tuning datasets.
A Background
We provide a brief summary of diffusion models.
We follow the standard formulation and notations [36, 38, 14].
A diffusion model is a class of latent variable models with a Markov chain of latent variables describing the generative process.
Given a data x 0 ∼ p data ( x ) :
= q ( x ) , a forward process injects a series of Markovian Gaussian noises with a sequence of controlled variances { β t } Tt =1 over T steps, until the perturbed data x T is nearly isotropic Gaussian distributed:
[Equation]
Reversing the forward process, a backward diffusion process will bring an isotropic Gaussian noise x T ∼ N ( 0 , I ) to a sample from the data distribution x 0 ∼ q ( x 0 ) .
Since the backward process q ( x t | x t − 1 ) is intractable
[Equation]
we introduce an approximate backward process defined with the model p θ ( x t | x t − 1 ) parameterized by θ :
[Equation]
Then we can optimize the parameter θ by maximizing the variational lower-bound, or minimizing the negative of the variational lower bound.
[Equation]
By Bayes’ rule and Markov property
[Equation]
one can easily derive that
[Equation]
and the variational objective (6) boils down to
[Equation]
where ≡ denotes the equality up to a constant.
Employing reparameterization for Gaussians [21], Ho et al.
[14] proposes to further simplify the objective as
[Equation]
where we parameterize the noise ϵ θ instead of the mean µ θ and the scaling factors ∥
Σ θ ( x t , t ) ∥ 2 are discarded to have uniform loss weights for all time steps.
While Ho et al.
[14] demonstrated that using fixed variance schedules for the model distributions empirically working well, that is, setting Σ θ ( x t , t )
=
σ 2 t I , Nichol and Dhariwal [27] later showed that one can improve upon fixed variance with the variances parameterized by weighting vector v , Σ θ ( x t , t ) =
exp( v log β t
+ (1 − v ) log ˜ β t ) .
Since the simplified objective (10) dropping the scaling factors does not contain the variance terms, Nichol and Dhariwal [27] proposed to optimize a hybrid objective defined as
[Equation]
(11) for some λ > 0 , together with cosine annealing variance scheduling which is specifically tailored for the low-resolution images (e.g., 64 × 64 or 32 × 32 images).
We follow the training scheme and model architecture described in Nichol and Dhariwal [27] throughout our experiments.
B Related works
Generative models for image synthesis Generative models for image synthesis are largely categorized into likelihood-based approaches [21, 42, 5, 43, 9] and GAN s [11, 19, 3].
Diffusion models have recently gained popularity due to their ability to generate high-quality images [37, 14, 7].
Most modern generative models use deep neural networks requiring a large amount of training data, so they usually suffer from overfitting when trained with limited data.
Image synthesis with limited data Multiple works have been proposed to improve the performance of image synthesis when the number of data is limited.
For GAN -based methods, Karras et al. [18], Zhao et al.
[48] propose data augmentation techniques to prevent discriminator overfitting.
Liu et al.
[23] propose a lightweight network architecture suitable for small datasets.
Meanwhile, leveraging pre-trained models on larger data and fine-tuning them on smaller data sets is becoming as popular as in other fields.
Mo et al.
[25] show that freezing discriminators in GAN fine-tuning gives stable results.
Ojha et al.
[28] constrain the adapted model to share latent space with the source model.
Zhao et al.
[47] shows that the low-level filters of pre-trained GAN s are helpful in learning limited training data even when they are perceptually distant from the target data.
Sauer et al.
[34], Kumari et al.
[22] used pre-trained models as feature extractors for a discriminator and show the effectiveness of faster and more stable convergence.
However, fine-tuning with limited data for unconditional diffusion models is still less explored yet.
Fine-tuning for diffusion models Recently, a high-performance pre-trained diffusion model has been published [30, 33].
As of now, few pioneering works investigate fine-tuning diffusion models;
Giannone et al.
[10] proposing few-shot learning diffusion models that can quickly learn conditional generative model for a set of images, and Ruiz et al.
[32] proposing a prompt-like approach for textto-image diffusion models.
Wang et al.
[44] propose two-stage finetuning scheme for image-to-image translation.
However, these works focus on conditional generation and fine-tune the pre-trained model without network modification.
To the best of our knowledge, our method is the first one to study module-level effects for fine-tuning an unconditional diffusion model with limited data.
Adapter-based method for transfer learning In a natural language processing task, Houlsby et al. [16] proposed a memory-efficient method for transformers that only updates adapter modules while freezing other weights during fine-tuning.
The adapter-based method requires only a small number of additional parameters, so it is efficient when adapting to a variety of tasks.
Karimi Mahabadi et al. [17] developed a more compact and faster adapter module based on matrix factorization and weight sharing.
Chen et al.
[4], Sung et al.
[41] adopt the adapter method for the vision task and showed effectiveness on reducing memory and training time.
C Experimental Details C.1 Configurations
Our algorithm is implemented in Python 3.9.5 and Pytorch 1.10.2.
All experiments are performed on four NVIDIA Tesla V100s.
C.2 Setting
We present detailed experimental settings for experiment results reproduction.
For the low-resolution datasets, we use the ImageNet 64 × 64 pre-trained model ( ≈ 120M parameters) in Nichol and Dhariwal [27].
Our model consists of a total of 15 attention blocks and 30 residual blocks.
Input images go under four downsampling steps, each with [2 C, 2 C, 2 C, 4 C ] channels where C = 128 , and then upsampled via four symmetric steps.
The pre-trained model employs multi-head attention blocks with 4 heads.
For the high-resolution datasets, we use the FFHQ 256 × 256 pre-trained model ( ≈ 90M parameters) presented in [6], which is composed of a total of 4 attention blocks andresidual blocks.
We employ Adam [20] for all experiments, with a learning rate of 2 · 10 − 4 and epsilon of 1 · 10 − 8 .
We set the batch size to 128 and 5 for fine-tuning and sampling, respectively.
Opposed to Nichol and Dhariwal [27], we do not use neither dropout [40] nor exponential moving average over model parameters.
To fine-tune diffusion models, we use the hybrid loss in (11) with cosine beta scheduling such that ( β 1 , β 2 ) =
(1 · 10 − 4 , 2 · 10 − 2 ) .
We set diffusion timesteps as T = 1000 .
We fine-tune the pre-trained model with 1000 CelebA images, 833 Pokemon images and 1336 MetFaces images.
Evaluation metrics are computed between the whole training dataset and 10K/5K/1.4K samples from CelebA/Pokemon/MetFaces dataset.
C.3 Metrics
Since quantifying image generation performance remains a challenge, metrics must be carefully chosen.
We use two different evaluation metrics: (i) clean Fréchet Inception Distance (cFID) [29] and (ii) Kernel Inception Distance (KID) [2].
While FID [12] is a popular choice for capturing both image fidelity and diversity, Parmar et al.
[29] points out that FID is inconsistent with image resizing and quantization.
To resolve such an issue, we use cFID by default.
Due to our relatively small dataset size, we also implement KID which is proven to be unbiased and asymptotically normal.
Metrics are computed between the whole training dataset and generated images.
C.4 Time-aware adapter
In this section, we review a detailed structure of the time-aware adapter implemented in low-resolution experiment.
In conventional adapter-based methods [16, 17], independent adapters are employed across all blocks.
In Houlsby et al.
[16] and Karimi Mahabadi et al.
[17], every single transformer layer needs two separate adapters so that the number of adapters needed grows as to the number of layers.
We alleviate such an issue using a weight-sharing technique.
We let the same adapter submodule fit into several attention blocks.
We build two time-fusion modules which integrate into attention blocks with 384 and 512 channels.
Time-fusion module first match the dimension of two inputs, t e mbed and z i n .
The time-embedding vector t e mbed goes through two linear layers to match the dimension of z i n , and we add such two vectors before passing on to the final linear layer.
Similarly, we build only two cross-attention modules which are attached to attention blocks with 384 and 512 channels.
We implement vanilla single-head attention operation.
Importantly, we construct separate time-scaling modules for all attention blocks analogous to conventional NLP adapters.
This is because we want to regulate the information from time-aware adapter in a block-wise manner.
Time-scaling module is merely an
MLP with two fully connected layers with SiLU activation [8].
D Additional results D.1 Sample diversity